@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2e56975657
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 17:15 UTC

Generative Scoring (#34539)

Signed-off-by: Vedant Jhaveri <vjhaveri@linkedin.com> Co-authored-by: Vedant Jhaveri <vjhaveri@linkedin.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Vedant V Jhaveri committed Mar 31, 2026 at 16:02 UTC 2e569756570bbe1297694ec1dcc93844198ed685
13 files changed +1265 -3
docs/serving/openai_compatible_server.md
+70 -2
@@ -73,8 +73,11 @@ In addition, we have the following custom APIs:
73 - [Cohere Embed API](../models/pooling_models/embed.md#cohere-embed-api) (`/v2/embed`)
74 - Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
75 - Works with any [embedding model](../models/pooling_models/embed.md#supported-models), including multimodal models.
76 -- [Score API](../models/pooling_models/scoring.md#score-api) (`/score`)
77 - - Applicable to [score models](../models/pooling_models/scoring.md).
76 +- [Score API](../models/pooling_models/scoring.md#score-api) (`/score`, `/v1/score`)
77 + - Applicable to [score models](../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction).
78 +- [Generative Scoring API](#generative-scoring-api) (`/generative_scoring`)
79 + - Applicable to [CausalLM models](../models/generative_models.md) (task `"generate"`).
80 + - Computes next-token probabilities for specified `label_token_ids`.
81 - [Rerank API](../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
82 - Implements [Jina AI's v1 rerank API](https://jina.ai/reranker/)
83 - Also compatible with [Cohere's v1 & v2 rerank APIs](https://docs.cohere.com/v2/reference/rerank)
@@ -481,6 +484,71 @@ This approach is more robust than index-based access (`messages[0]`, `messages[1
484
485 Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
486
487 +### Generative Scoring API
488 +
489 +The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
490 +
491 +This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
492 +
493 +**Requirements:**
494 +
495 +- The `label_token_ids` parameter is **required** and must contain **at least 1 token ID**.
496 +- When 2 label tokens are provided, the score equals `P(label_token_ids[0]) / (P(label_token_ids[0]) + P(label_token_ids[1]))` (softmax over the two labels).
497 +- When more labels are provided, the score is the softmax-normalized probability of the first label token across all label tokens.
498 +
499 +#### Example
500 +
501 +```bash
502 +curl -X POST http://localhost:8000/generative_scoring \
503 + -H "Content-Type: application/json" \
504 + -d '{
505 + "model": "Qwen/Qwen3-0.6B",
506 + "query": "Is this city the capital of France?",
507 + "items": ["Paris", "London", "Berlin"],
508 + "label_token_ids": [9454, 2753]
509 + }'
510 +```
511 +
512 +Here, each item is appended to the query to form prompts like `"Is this city the capital of France? Paris"`, `"... London"`, etc. The model then predicts the next token, and the score reflects the probability of "Yes" (token 9454) vs "No" (token 2753).
513 +
514 +??? console "Response"
515 +
516 + ```json
517 + {
518 + "id": "generative-scoring-abc123",
519 + "object": "list",
520 + "created": 1234567890,
521 + "model": "Qwen/Qwen3-0.6B",
522 + "data": [
523 + {"index": 0, "object": "score", "score": 0.95},
524 + {"index": 1, "object": "score", "score": 0.12},
525 + {"index": 2, "object": "score", "score": 0.08}
526 + ],
527 + "usage": {"prompt_tokens": 45, "total_tokens": 48, "completion_tokens": 3}
528 + }
529 + ```
530 +
531 +#### How it works
532 +
533 +1. **Prompt Construction**: For each item, builds `prompt = query + item` (or `item + query` if `item_first=true`)
534 +2. **Forward Pass**: Runs the model on each prompt to get next-token logits
535 +3. **Probability Extraction**: Extracts logprobs for the specified `label_token_ids`
536 +4. **Softmax Normalization**: Applies softmax over only the label tokens (when `apply_softmax=true`)
537 +5. **Score**: Returns the normalized probability of the first label token
538 +
539 +#### Finding Token IDs
540 +
541 +To find the token IDs for your labels, use the tokenizer:
542 +
543 +```python
544 +from transformers import AutoTokenizer
545 +
546 +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
547 +yes_id = tokenizer.encode("Yes", add_special_tokens=False)[0]
548 +no_id = tokenizer.encode("No", add_special_tokens=False)[0]
549 +print(f"Yes: {yes_id}, No: {no_id}")
550 +```
551 +
552 ## Ray Serve LLM
553
554 Ray Serve LLM enables scalable, production-grade serving of the vLLM engine. It integrates tightly with vLLM and extends it with features such as auto-scaling, load balancing, and back-pressure.
tests/entrypoints/openai/generative_scoring/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/entrypoints/openai/generative_scoring/test_generative_scoring.py new
+325
@@ -0,0 +1,325 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the Generative Scoring API.
4 +
5 +Tests cover:
6 +1. Protocol models (request/response construction)
7 +2. Probability computation (softmax normalization)
8 +3. Input validation
9 +4. Score formula: P(token[0]) / (P(token[0]) + P(token[1]))
10 +5. Prompt building and item ordering
11 +"""
12 +
13 +import math
14 +from dataclasses import dataclass, field
15 +from typing import Any
16 +from unittest.mock import MagicMock
17 +
18 +import pytest
19 +
20 +from vllm.config.multimodal import MultiModalConfig
21 +from vllm.entrypoints.openai.engine.protocol import ErrorResponse
22 +from vllm.entrypoints.openai.generative_scoring.serving import (
23 + GenerativeScoringItemResult,
24 + GenerativeScoringRequest,
25 + GenerativeScoringResponse,
26 + OpenAIServingGenerativeScoring,
27 +)
28 +from vllm.entrypoints.openai.models.protocol import BaseModelPath
29 +from vllm.entrypoints.openai.models.serving import OpenAIServingModels
30 +from vllm.logprobs import Logprob
31 +from vllm.outputs import CompletionOutput, RequestOutput
32 +from vllm.tokenizers import get_tokenizer
33 +from vllm.v1.engine.async_llm import AsyncLLM
34 +
35 +MODEL_NAME = "Qwen/Qwen3-0.6B"
36 +BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
37 +
38 +
39 +@dataclass
40 +class MockHFConfig:
41 + model_type: str = "any"
42 +
43 +
44 +@dataclass
45 +class MockModelConfig:
46 + task = "generate"
47 + runner_type = "generate"
48 + tokenizer = MODEL_NAME
49 + trust_remote_code = False
50 + tokenizer_mode = "auto"
51 + max_model_len = 100
52 + tokenizer_revision = None
53 + multimodal_config = MultiModalConfig()
54 + hf_config = MockHFConfig()
55 + logits_processor_pattern = None
56 + logits_processors: list[str] | None = None
57 + diff_sampling_param: dict | None = None
58 + allowed_local_media_path: str = ""
59 + allowed_media_domains: list[str] | None = None
60 + encoder_config = None
61 + generation_config: str = "auto"
62 + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
63 + skip_tokenizer_init = False
64 + vocab_size = 151936
65 +
66 + def get_diff_sampling_param(self):
67 + return self.diff_sampling_param or {}
68 +
69 + def get_vocab_size(self):
70 + return self.vocab_size
71 +
72 +
73 +def _create_mock_engine():
74 + """Create a mock AsyncLLM engine."""
75 + mock_engine = MagicMock(spec=AsyncLLM)
76 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
77 + mock_engine.errored = False
78 + mock_engine.model_config = MockModelConfig()
79 + mock_engine.input_processor = MagicMock()
80 + mock_engine.io_processor = MagicMock()
81 +
82 + # renderer is accessed by OpenAIServing.__init__ and serving.py
83 + mock_renderer = MagicMock()
84 + mock_renderer.tokenizer = get_tokenizer(MODEL_NAME)
85 + mock_engine.renderer = mock_renderer
86 +
87 + return mock_engine
88 +
89 +
90 +def _create_serving(mock_engine) -> OpenAIServingGenerativeScoring:
91 + """Create an OpenAIServingGenerativeScoring instance with mocks."""
92 + models = OpenAIServingModels(
93 + engine_client=mock_engine,
94 + base_model_paths=BASE_MODEL_PATHS,
95 + )
96 + return OpenAIServingGenerativeScoring(mock_engine, models, request_logger=None)
97 +
98 +
99 +def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput:
100 + """Create a mock RequestOutput with specified logprobs."""
101 + logprobs_with_objs = {
102 + tid: Logprob(logprob=lp, rank=i + 1)
103 + for i, (tid, lp) in enumerate(logprobs_dict.items())
104 + }
105 + completion_output = CompletionOutput(
106 + index=0,
107 + text="",
108 + token_ids=[100],
109 + cumulative_logprob=-1.0,
110 + logprobs=[logprobs_with_objs],
111 + finish_reason="length",
112 + )
113 + return RequestOutput(
114 + request_id="test-request",
115 + prompt="test prompt",
116 + prompt_token_ids=[1, 2, 3],
117 + prompt_logprobs=None,
118 + outputs=[completion_output],
119 + finished=True,
120 + )
121 +
122 +
123 +class TestProtocolModels:
124 + """Tests for GenerativeScoringRequest and GenerativeScoringResponse."""
125 +
126 + def test_request_and_response_all_fields(self):
127 + """Test request construction with all field types and response structure."""
128 + # Test request with string inputs
129 + req_str = GenerativeScoringRequest(
130 + query="Is this the capital?",
131 + items=["Paris", "London"],
132 + label_token_ids=[9454, 2753],
133 + )
134 + assert req_str.query == "Is this the capital?"
135 + assert req_str.items == ["Paris", "London"]
136 + assert req_str.label_token_ids == [9454, 2753]
137 + assert req_str.apply_softmax is True # default
138 + assert req_str.item_first is False # default
139 + assert req_str.add_special_tokens is True # default
140 +
141 + # Test request with pre-tokenized inputs and custom options
142 + req_tok = GenerativeScoringRequest(
143 + query=[100, 200, 300],
144 + items=[[400, 500], [600, 700]],
145 + label_token_ids=[1234, 5678],
146 + apply_softmax=False,
147 + item_first=True,
148 + add_special_tokens=False,
149 + )
150 + assert req_tok.query == [100, 200, 300]
151 + assert req_tok.items == [[400, 500], [600, 700]]
152 + assert req_tok.apply_softmax is False
153 + assert req_tok.item_first is True
154 + assert req_tok.add_special_tokens is False
155 +
156 + # Test response structure
157 + response = GenerativeScoringResponse(
158 + model="test-model",
159 + data=[
160 + GenerativeScoringItemResult(index=0, score=0.7),
161 + GenerativeScoringItemResult(index=1, score=0.4),
162 + ],
163 + usage={"prompt_tokens": 10, "total_tokens": 12, "completion_tokens": 2},
164 + )
165 + assert response.object == "list"
166 + assert response.model == "test-model"
167 + assert len(response.data) == 2
168 + assert response.data[0].score == 0.7
169 + assert response.data[0].object == "score"
170 + assert response.data[1].score == 0.4
171 + assert response.usage.prompt_tokens == 10
172 +
173 +
174 +class TestProbabilityComputation:
175 + """Tests for _compute_probabilities with both softmax modes."""
176 +
177 + @pytest.mark.parametrize(
178 + "label_logprobs,apply_softmax,should_sum_to_one",
179 + [
180 + ({100: -1.0, 200: -2.0}, True, True),
181 + ({100: -100.0, 200: -100.5}, True, True), # numerical stability
182 + ({100: -1.0, 200: -2.0}, False, False),
183 + ],
184 + ids=["softmax_basic", "softmax_extreme_values", "true_probs"],
185 + )
186 + def test_compute_probabilities(
187 + self, label_logprobs, apply_softmax, should_sum_to_one
188 + ):
189 + """Test probability computation for softmax and true probability modes."""
190 + serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
191 + probs = serving._compute_probabilities(
192 + label_logprobs, apply_softmax=apply_softmax
193 + )
194 +
195 + # Verify sum behavior
196 + total = sum(probs.values())
197 + if should_sum_to_one:
198 + assert abs(total - 1.0) < 1e-6
199 + else:
200 + assert total < 1.0
201 +
202 + # Verify math
203 + if apply_softmax:
204 + max_lp = max(label_logprobs.values())
205 + exp_vals = {k: math.exp(v - max_lp) for k, v in label_logprobs.items()}
206 + sum_exp = sum(exp_vals.values())
207 + for tid, lp in label_logprobs.items():
208 + assert abs(probs[tid] - exp_vals[tid] / sum_exp) < 1e-9
209 + else:
210 + for tid, lp in label_logprobs.items():
211 + assert abs(probs[tid] - math.exp(lp)) < 1e-9
212 +
213 + def test_score_formula(self):
214 + """Test the score formula: P(token[0]) / (P(token[0]) + P(token[1]))."""
215 + serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
216 +
217 + # With logprobs -0.5 and -2.0, softmax gives higher prob to first token
218 + logprobs = {9454: -0.5, 2753: -2.0}
219 + probs = serving._compute_probabilities(logprobs, apply_softmax=True)
220 +
221 + # Score = P(9454) / (P(9454) + P(2753)) = P(9454) since they sum to 1
222 + score = probs[9454]
223 +
224 + # Manual calculation
225 + exp_0 = math.exp(-0.5)
226 + exp_1 = math.exp(-2.0)
227 + expected_score = exp_0 / (exp_0 + exp_1)
228 +
229 + assert abs(score - expected_score) < 1e-9
230 + assert score > 0.5 # First token has higher logprob, so higher probability
231 +
232 +
233 +class TestValidation:
234 + """Tests for input validation errors."""
235 +
236 + @pytest.mark.asyncio
237 + @pytest.mark.parametrize(
238 + "request_kwargs,expected_error",
239 + [
240 + (
241 + {"query": "q", "items": ["i"], "label_token_ids": [999999, 999998]},
242 + "out of vocabulary",
243 + ),
244 + (
245 + {"query": "q", "items": [], "label_token_ids": [100, 200]},
246 + "at least one item",
247 + ),
248 + ],
249 + ids=["invalid_token_id", "empty_items"],
250 + )
251 + async def test_validation_errors(self, request_kwargs, expected_error):
252 + """Test that invalid inputs return appropriate errors."""
253 + mock_engine = _create_mock_engine()
254 + serving = _create_serving(mock_engine)
255 + request = GenerativeScoringRequest(model=MODEL_NAME, **request_kwargs)
256 + result = await serving.create_generative_scoring(request, None)
257 +
258 + assert isinstance(result, ErrorResponse)
259 + assert expected_error in result.error.message.lower()
260 +
261 +
262 +class TestPromptBuilding:
263 + """Tests for prompt construction and item ordering."""
264 +
265 + @pytest.mark.asyncio
266 + @pytest.mark.parametrize(
267 + "item_first,expected",
268 + [
269 + (False, [[100, 101, 200, 201], [100, 101, 300, 301]]), # query + item
270 + (True, [[200, 201, 100, 101], [300, 301, 100, 101]]), # item + query
271 + ],
272 + ids=["query_first", "item_first"],
273 + )
274 + async def test_item_ordering(self, item_first, expected):
275 + """Test that item_first flag controls prompt concatenation order."""
276 + mock_engine = _create_mock_engine()
277 + serving = _create_serving(mock_engine)
278 +
279 + request = GenerativeScoringRequest(
280 + query=[100, 101],
281 + items=[[200, 201], [300, 301]],
282 + label_token_ids=[500, 501],
283 + item_first=item_first,
284 + )
285 + engine_inputs, _ = await serving._build_prompts(
286 + request, MagicMock(), max_model_len=4096
287 + )
288 +
289 + for i, exp in enumerate(expected):
290 + assert engine_inputs[i]["prompt_token_ids"] == exp
291 +
292 +
293 +class TestGeneration:
294 + """Tests for the full generation flow with mocked engine."""
295 +
296 + @pytest.mark.asyncio
297 + async def test_successful_generation(self):
298 + """Test successful score generation returns valid response."""
299 + mock_engine = _create_mock_engine()
300 + serving = _create_serving(mock_engine)
301 +
302 + mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0}
303 + mock_output = _create_mock_request_output(mock_logprobs)
304 +
305 + async def mock_generate(*args, **kwargs):
306 + yield mock_output
307 +
308 + mock_engine.generate = mock_generate
309 +
310 + request = GenerativeScoringRequest(
311 + model=MODEL_NAME,
312 + query="Is Paris the capital?",
313 + items=["Yes", "No"],
314 + label_token_ids=[1234, 5678],
315 + )
316 + result = await serving.create_generative_scoring(request, None)
317 +
318 + assert isinstance(result, GenerativeScoringResponse)
319 + assert len(result.data) == 2
320 + for item_result in result.data:
321 + assert 0.0 <= item_result.score <= 1.0
322 +
323 +
324 +if __name__ == "__main__":
325 + pytest.main([__file__, "-v"])
tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py new
+157
@@ -0,0 +1,157 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end tests for the Generative Scoring API.
4 +
5 +Tests verify the full HTTP request/response flow using RemoteOpenAIServer.
6 +"""
7 +
8 +import pytest
9 +import requests
10 +
11 +from ....utils import RemoteOpenAIServer
12 +
13 +MODEL_NAME = "Qwen/Qwen3-0.6B"
14 +
15 +
16 +@pytest.fixture(scope="module")
17 +def server():
18 + args = [
19 + "--dtype",
20 + "bfloat16",
21 + "--max-model-len",
22 + "512",
23 + "--enforce-eager",
24 + "--max-num-seqs",
25 + "32",
26 + ]
27 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
28 + yield remote_server
29 +
30 +
31 +class TestGenerativeScoringAPI:
32 + """End-to-end tests for the Generative Scoring API."""
33 +
34 + @pytest.mark.asyncio
35 + async def test_basic_score_and_response_structure(self, server: RemoteOpenAIServer):
36 + """Test basic generative scoring request and verify response structure."""
37 + response = requests.post(
38 + server.url_for("generative_scoring"),
39 + json={
40 + "model": MODEL_NAME,
41 + "query": "Is Paris the capital of France? Answer Yes or No: ",
42 + "items": ["Paris is beautiful.", "London is rainy."],
43 + "label_token_ids": [9454, 2753],
44 + },
45 + )
46 + assert response.status_code == 200, f"Response: {response.text}"
47 + data = response.json()
48 +
49 + # Verify response structure
50 + assert data["id"].startswith("generative-scoring-")
51 + assert data["object"] == "list"
52 + assert "model" in data
53 + assert "usage" in data
54 + assert len(data["data"]) == 2
55 +
56 + # Verify each result
57 + for i, result in enumerate(data["data"]):
58 + assert result["index"] == i
59 + assert result["object"] == "score"
60 + assert 0.0 <= result["score"] <= 1.0
61 +
62 + # Verify usage tracking
63 + usage = data["usage"]
64 + assert usage["prompt_tokens"] > 0
65 + assert usage["completion_tokens"] > 0
66 + assert (
67 + usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
68 + )
69 +
70 + @pytest.mark.asyncio
71 + async def test_multiple_items(self, server: RemoteOpenAIServer):
72 + """Test generative scoring request with multiple items."""
73 + response = requests.post(
74 + server.url_for("generative_scoring"),
75 + json={
76 + "model": MODEL_NAME,
77 + "query": "Is this city a capital? ",
78 + "items": ["Paris", "London", "Berlin", "New York", "Tokyo"],
79 + "label_token_ids": [9454, 2753],
80 + },
81 + )
82 + assert response.status_code == 200
83 + data = response.json()
84 + assert len(data["data"]) == 5
85 +
86 + @pytest.mark.asyncio
87 + async def test_validation_missing_label_token_ids(self, server: RemoteOpenAIServer):
88 + """Test that missing label_token_ids returns a validation error."""
89 + response = requests.post(
90 + server.url_for("generative_scoring"),
91 + json={
92 + "model": MODEL_NAME,
93 + "query": "Test query",
94 + "items": ["item1", "item2"],
95 + },
96 + )
97 + # Missing required field returns 400 (manual JSON parsing)
98 + assert response.status_code == 400
99 +
100 + @pytest.mark.asyncio
101 + async def test_validation_empty_items(self, server: RemoteOpenAIServer):
102 + """Test that empty items returns an error."""
103 + response = requests.post(
104 + server.url_for("generative_scoring"),
105 + json={
106 + "model": MODEL_NAME,
107 + "query": "Test query",
108 + "items": [],
109 + "label_token_ids": [100, 200],
110 + },
111 + )
112 + assert response.status_code == 400
113 +
114 + @pytest.mark.asyncio
115 + @pytest.mark.parametrize(
116 + "label_token_ids,expected_status",
117 + [
118 + ([9999999999, 9999999998], 400), # Out of vocab range
119 + ],
120 + ids=["invalid_token_ids"],
121 + )
122 + async def test_validation_errors(
123 + self, server: RemoteOpenAIServer, label_token_ids, expected_status
124 + ):
125 + """Test validation errors for various invalid inputs."""
126 + response = requests.post(
127 + server.url_for("generative_scoring"),
128 + json={
129 + "model": MODEL_NAME,
130 + "query": "Test query",
131 + "items": ["item1"],
132 + "label_token_ids": label_token_ids,
133 + },
134 + )
135 + assert response.status_code == expected_status
136 +
137 + @pytest.mark.asyncio
138 + async def test_score_consistency(self, server: RemoteOpenAIServer):
139 + """Test that scores are deterministic across identical requests."""
140 + request_body = {
141 + "model": MODEL_NAME,
142 + "query": "Is this consistent? ",
143 + "items": ["Yes it is."],
144 + "label_token_ids": [100, 200],
145 + }
146 +
147 + r1 = requests.post(server.url_for("generative_scoring"), json=request_body)
148 + r2 = requests.post(server.url_for("generative_scoring"), json=request_body)
149 +
150 + assert r1.status_code == 200 and r2.status_code == 200
151 + r1_score = r1.json()["data"][0]["score"]
152 + r2_score = r2.json()["data"][0]["score"]
153 + assert abs(r1_score - r2_score) < 1e-6
154 +
155 +
156 +if __name__ == "__main__":
157 + pytest.main([__file__, "-v"])
vllm/entrypoints/openai/api_server.py
+14
@@ -246,6 +246,13 @@ def build_app(
246
247 register_pooling_api_routers(app, supported_tasks, model_config)
248
249 + if "generate" in supported_tasks:
250 + from vllm.entrypoints.openai.generative_scoring.api_router import (
251 + register_generative_scoring_api_router,
252 + )
253 +
254 + register_generative_scoring_api_router(app)
255 +
256 app.root_path = args.root_path
257 app.add_middleware(
258 CORSMiddleware,
@@ -413,6 +420,13 @@ async def init_app_state(
420
421 init_pooling_state(engine_client, state, args, request_logger, supported_tasks)
422
423 + if "generate" in supported_tasks:
424 + from vllm.entrypoints.openai.generative_scoring.api_router import (
425 + init_generative_scoring_state,
426 + )
427 +
428 + await init_generative_scoring_state(engine_client, state, args, request_logger)
429 +
430 state.enable_server_load_tracking = args.enable_server_load_tracking
431 state.server_load_metrics = 0
432
vllm/entrypoints/openai/generative_scoring/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/entrypoints/openai/generative_scoring/api_router.py new
+87
@@ -0,0 +1,87 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from http import HTTPStatus
4 +from typing import TYPE_CHECKING
5 +
6 +from fastapi import APIRouter, Depends, FastAPI, Request
7 +from fastapi.responses import JSONResponse
8 +
9 +from vllm.entrypoints.openai.engine.protocol import ErrorResponse
10 +from vllm.entrypoints.openai.generative_scoring.serving import (
11 + GenerativeScoringResponse,
12 + OpenAIServingGenerativeScoring,
13 +)
14 +from vllm.entrypoints.openai.utils import validate_json_request
15 +from vllm.entrypoints.utils import load_aware_call, with_cancellation
16 +from vllm.logger import init_logger
17 +
18 +if TYPE_CHECKING:
19 + from argparse import Namespace
20 +
21 + from starlette.datastructures import State
22 +
23 + from vllm.engine.protocol import EngineClient
24 + from vllm.entrypoints.logger import RequestLogger
25 +
26 +router = APIRouter()
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +def generative_scoring(request: Request) -> OpenAIServingGenerativeScoring | None:
32 + return request.app.state.serving_generative_scoring
33 +
34 +
35 +@router.post(
36 + "/generative_scoring",
37 + dependencies=[Depends(validate_json_request)],
38 + responses={
39 + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
40 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
41 + },
42 +)
43 +@with_cancellation
44 +@load_aware_call
45 +async def create_generative_scoring(raw_request: Request):
46 + handler = generative_scoring(raw_request)
47 + if handler is None:
48 + raise NotImplementedError(
49 + "The model does not support the Generative Scoring API"
50 + )
51 +
52 + raw_body = await raw_request.json()
53 +
54 + from vllm.entrypoints.openai.generative_scoring.serving import (
55 + GenerativeScoringRequest,
56 + )
57 +
58 + gen_request = GenerativeScoringRequest(**raw_body)
59 + result = await handler.create_generative_scoring(gen_request, raw_request)
60 +
61 + if isinstance(result, ErrorResponse):
62 + return JSONResponse(content=result.model_dump(), status_code=result.error.code)
63 + elif isinstance(result, GenerativeScoringResponse):
64 + return JSONResponse(content=result.model_dump())
65 +
66 + raise ValueError(f"Unexpected response type: {type(result)}")
67 +
68 +
69 +def register_generative_scoring_api_router(app: FastAPI):
70 + app.include_router(router)
71 +
72 +
73 +async def init_generative_scoring_state(
74 + engine_client: "EngineClient",
75 + state: "State",
76 + args: "Namespace",
77 + request_logger: "RequestLogger | None",
78 +):
79 + from vllm.entrypoints.openai.generative_scoring.serving import (
80 + OpenAIServingGenerativeScoring,
81 + )
82 +
83 + state.serving_generative_scoring = OpenAIServingGenerativeScoring(
84 + engine_client,
85 + state.openai_serving_models,
86 + request_logger=request_logger,
87 + )
vllm/entrypoints/openai/generative_scoring/serving.py new
+491
@@ -0,0 +1,491 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Generative Scoring implementation for generative models.
4 +
5 +This module implements generative scoring functionality that computes the
6 +probability of specified token IDs appearing as the next token after a
7 +given query+item prompt. This works on any generative model that produces
8 +logits (task="generate").
9 +"""
10 +
11 +import asyncio
12 +import math
13 +import time
14 +from collections.abc import AsyncGenerator, Mapping
15 +from typing import Literal
16 +
17 +from fastapi import Request
18 +from pydantic import Field
19 +
20 +from vllm.engine.protocol import EngineClient
21 +from vllm.entrypoints.logger import RequestLogger
22 +from vllm.entrypoints.openai.engine.protocol import (
23 + ErrorResponse,
24 + OpenAIBaseModel,
25 + UsageInfo,
26 +)
27 +from vllm.entrypoints.openai.engine.serving import OpenAIServing
28 +from vllm.entrypoints.openai.models.serving import OpenAIServingModels
29 +from vllm.inputs import EngineInput, tokens_input
30 +from vllm.logger import init_logger
31 +from vllm.outputs import RequestOutput
32 +from vllm.sampling_params import SamplingParams
33 +from vllm.tokenizers import TokenizerLike
34 +from vllm.tracing import (
35 + contains_trace_headers,
36 + extract_trace_headers,
37 + log_tracing_disabled_warning,
38 +)
39 +from vllm.utils import random_uuid
40 +from vllm.utils.async_utils import merge_async_iterators
41 +
42 +logger = init_logger(__name__)
43 +
44 +
45 +# ============================================================================
46 +# Protocol definitions
47 +# ============================================================================
48 +
49 +
50 +class GenerativeScoringRequest(OpenAIBaseModel):
51 + """Request for computing generative scoring.
52 +
53 + Attributes:
54 + model: The model to use for scoring. Optional, follows existing patterns.
55 + query: The query text or pre-tokenized query token IDs.
56 + items: The item text(s) or pre-tokenized item token IDs.
57 + label_token_ids: List of token IDs to compute probabilities for.
58 + apply_softmax: Whether to normalize probabilities using softmax over only
59 + the label_token_ids (True) or return true model probabilities over
60 + the full vocab for those ids (False).
61 + item_first: If True, prepend items to query. Otherwise append items to query.
62 + add_special_tokens: Whether to add special tokens when tokenizing.
63 + """
64 +
65 + model: str | None = None
66 + query: str | list[int] = Field(
67 + ...,
68 + description="The query text or pre-tokenized query token IDs.",
69 + )
70 + items: list[str] | list[list[int]] = Field(
71 + ...,
72 + description="List of item texts or pre-tokenized item token IDs.",
73 + )
74 + label_token_ids: list[int] = Field(
75 + ...,
76 + description="List of token IDs to compute probabilities for.",
77 + )
78 + apply_softmax: bool = Field(
79 + default=True,
80 + description=(
81 + "If True, normalize probabilities using softmax over only the "
82 + "label_token_ids. If False, return the true model probabilities "
83 + "over the full vocab for those ids."
84 + ),
85 + )
86 + item_first: bool = Field(
87 + default=False,
88 + description="If True, prepend items to query. Otherwise append items to query.",
89 + )
90 + add_special_tokens: bool = Field(
91 + default=True,
92 + description="Whether to add special tokens when tokenizing.",
93 + )
94 + priority: int = Field(
95 + default=0,
96 + description=(
97 + "The priority of the request (lower means earlier handling; default: 0)."
98 + ),
99 + )
100 + request_id: str = Field(
101 + default_factory=random_uuid,
102 + description="The request_id related to this request.",
103 + )
104 +
105 +
106 +class GenerativeScoringItemResult(OpenAIBaseModel):
107 + """Result for a single item in the generative scoring response.
108 +
109 + Attributes:
110 + index: The index of this item in the input items list.
111 + object: Type of object, always "score".
112 + score: The probability score for the first label token.
113 + """
114 +
115 + index: int
116 + object: Literal["score"] = "score"
117 + score: float
118 +
119 +
120 +class GenerativeScoringResponse(OpenAIBaseModel):
121 + """Response from the generative scoring computation.
122 +
123 + Attributes:
124 + id: Unique identifier for this response.
125 + object: Type of object, always "list".
126 + created: Unix timestamp of when the response was created.
127 + model: The model used for scoring.
128 + data: List of scoring results, one per input item.
129 + usage: Token usage information.
130 + """
131 +
132 + id: str = Field(default="")
133 + object: Literal["list"] = "list"
134 + created: int = Field(default_factory=lambda: int(time.time()))
135 + model: str
136 + data: list[GenerativeScoringItemResult]
137 + usage: UsageInfo
138 +
139 +
140 +# ============================================================================
141 +# Serving class
142 +# ============================================================================
143 +
144 +
145 +class OpenAIServingGenerativeScoring(OpenAIServing):
146 + """Serving class for generative scoring computation.
147 +
148 + This class handles computing the probability of specified token IDs
149 + appearing as the next token after concatenating query and item prompts.
150 +
151 + The key operation is:
152 + 1. For each item, build a prompt: query + item (or item + query if item_first)
153 + 2. Run a forward pass to get the next token distribution
154 + 3. Extract probabilities for the specified label_token_ids
155 + 4. Normalize either over the full vocab (apply_softmax=False) or
156 + over just the label_token_ids (apply_softmax=True)
157 + """
158 +
159 + def __init__(
160 + self,
161 + engine_client: EngineClient,
162 + models: OpenAIServingModels,
163 + *,
164 + request_logger: RequestLogger | None,
165 + ) -> None:
166 + super().__init__(
167 + engine_client=engine_client,
168 + models=models,
169 + request_logger=request_logger,
170 + )
171 +
172 + async def create_generative_scoring(
173 + self,
174 + request: GenerativeScoringRequest,
175 + raw_request: Request | None = None,
176 + ) -> GenerativeScoringResponse | ErrorResponse:
177 + """Create generative scoring for the given request.
178 +
179 + Args:
180 + request: The GenerativeScoringRequest containing query, items, and
181 + label_token_ids.
182 + raw_request: The raw FastAPI request object.
183 +
184 + Returns:
185 + GenerativeScoringResponse with probabilities for each item, or
186 + ErrorResponse if an error occurred.
187 + """
188 + # Check model
189 + error_check_ret = await self._check_model(request) # type: ignore[arg-type]
190 + if error_check_ret is not None:
191 + return error_check_ret
192 +
193 + # Check if engine is alive
194 + if self.engine_client.errored:
195 + raise self.engine_client.dead_error
196 +
197 + # Get tokenizer
198 + tokenizer = self.renderer.tokenizer
199 + if tokenizer is None:
200 + return self.create_error_response(
201 + "Tokenizer not available. Cannot process generative scoring request."
202 + )
203 +
204 + # Validate label_token_ids
205 + vocab_size = self.model_config.get_vocab_size()
206 + for token_id in request.label_token_ids:
207 + if token_id < 0 or token_id >= vocab_size:
208 + return self.create_error_response(
209 + f"label_token_id {token_id} is out of vocabulary range "
210 + f"[0, {vocab_size}). Please provide valid token IDs."
211 + )
212 +
213 + if len(request.label_token_ids) == 0:
214 + return self.create_error_response(
215 + "label_token_ids must contain at least one token ID."
216 + )
217 +
218 + # Validate items
219 + if len(request.items) == 0:
220 + return self.create_error_response("items must contain at least one item.")
221 +
222 + # Note: Mixed item types (string and token list) are validated by
223 + # Pydantic at request parsing time, so we don't need to check here.
224 +
225 + try:
226 + lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type]
227 + except (ValueError, TypeError, RuntimeError) as e:
228 + logger.exception("Error preparing request components")
229 + return self.create_error_response(e)
230 +
231 + base_id = self._base_request_id(raw_request, default=request.request_id)
232 + request_id = f"generative-scoring-{base_id}"
233 + created_time = int(time.time())
234 +
235 + # Build prompts for each item
236 + try:
237 + engine_inputs, prompt_token_counts = await self._build_prompts(
238 + request, tokenizer, self.model_config.max_model_len
239 + )
240 + except (ValueError, TypeError) as e:
241 + logger.exception("Error building prompts")
242 + return self.create_error_response(e)
243 +
244 + # Create sampling params for scoring
245 + # We use max_tokens=1 with logprob_token_ids to efficiently get
246 + # logprobs for only the specified label tokens (not full vocab)
247 + # Note: temperature/top_k/top_p don't affect logprobs - they only
248 + # affect the sampling distribution. Logprobs are computed from raw
249 + # logits via log_softmax before any sampling transformations.
250 + sampling_params = SamplingParams(
251 + max_tokens=1,
252 + logprobs=len(request.label_token_ids),
253 + logprob_token_ids=request.label_token_ids,
254 + n=1,
255 + )
256 +
257 + # Get trace headers
258 + trace_headers = (
259 + None
260 + if raw_request is None
261 + else await self._get_trace_headers(raw_request.headers)
262 + )
263 +
264 + # Schedule requests for all inputs
265 + generators: list[AsyncGenerator[RequestOutput, None]] = []
266 + for i, engine_input in enumerate(engine_inputs):
267 + request_id_item = f"{request_id}-{i}"
268 +
269 + self._log_inputs(
270 + request_id_item,
271 + engine_input,
272 + params=sampling_params,
273 + lora_request=lora_request,
274 + )
275 +
276 + generator = self.engine_client.generate(
277 + engine_input,
278 + sampling_params,
279 + request_id_item,
280 + lora_request=lora_request,
281 + trace_headers=trace_headers,
282 + priority=request.priority,
283 + )
284 + generators.append(generator)
285 +
286 + # Collect results
287 + result_generator = merge_async_iterators(*generators)
288 + results: list[RequestOutput | None] = [None] * len(engine_inputs)
289 +
290 + try:
291 + async for i, res in result_generator:
292 + results[i] = res
293 + except asyncio.CancelledError:
294 + return self.create_error_response("Client disconnected")
295 + except Exception as e:
296 + logger.exception("Error during generation")
297 + return self.create_error_response(e)
298 +
299 + # Process results to extract label token probabilities
300 + item_results: list[GenerativeScoringItemResult] = []
301 + total_prompt_tokens = 0
302 + total_completion_tokens = 0
303 +
304 + for i, result in enumerate(results):
305 + if result is None:
306 + return self.create_error_response(
307 + f"Failed to generate result for item {i}"
308 + )
309 +
310 + # Check for errors
311 + if result.outputs and result.outputs[0].finish_reason == "error":
312 + return self.create_error_response(f"Generation error for item {i}")
313 +
314 + # Get logprobs from the generated token
315 + if not result.outputs or len(result.outputs) == 0:
316 + return self.create_error_response(f"No output generated for item {i}")
317 +
318 + output = result.outputs[0]
319 + if output.logprobs is None or len(output.logprobs) == 0:
320 + return self.create_error_response(
321 + f"No logprobs available for item {i}. "
322 + "This might indicate an issue with logprobs configuration."
323 + )
324 +
325 + # The logprobs dict maps token_id -> Logprob object
326 + # For logprobs=-1, this contains all vocab tokens
327 + logprobs_dict = output.logprobs[0]
328 +
329 + # Extract logprobs for label tokens
330 + label_logprobs: dict[int, float] = {}
331 + missing_tokens = []
332 + for token_id in request.label_token_ids:
333 + if token_id in logprobs_dict:
334 + label_logprobs[token_id] = logprobs_dict[token_id].logprob
335 + else:
336 + missing_tokens.append(token_id)
337 +
338 + if missing_tokens:
339 + return self.create_error_response(
340 + f"Token IDs {missing_tokens} not found in logprobs for item {i}. "
341 + "This might indicate the tokens are outside the model's vocabulary."
342 + )
343 +
344 + # Compute probabilities based on apply_softmax setting
345 + token_probs = self._compute_probabilities(
346 + label_logprobs,
347 + apply_softmax=request.apply_softmax,
348 + )
349 +
350 + # Use the first label token's probability as the score
351 + first_label_id = request.label_token_ids[0]
352 + score = token_probs[first_label_id]
353 +
354 + item_results.append(
355 + GenerativeScoringItemResult(
356 + index=i,
357 + score=score,
358 + )
359 + )
360 +
361 + # Update token counts
362 + total_prompt_tokens += prompt_token_counts[i]
363 + total_completion_tokens += len(output.token_ids)
364 +
365 + # Build response
366 + model_name = self.models.model_name(lora_request)
367 + response = GenerativeScoringResponse(
368 + id=request_id,
369 + created=created_time,
370 + model=model_name,
371 + data=item_results,
372 + usage=UsageInfo(
373 + prompt_tokens=total_prompt_tokens,
374 + total_tokens=total_prompt_tokens + total_completion_tokens,
375 + completion_tokens=total_completion_tokens,
376 + ),
377 + )
378 +
379 + return response
380 +
381 + async def _build_prompts(
382 + self,
383 + request: GenerativeScoringRequest,
384 + tokenizer: TokenizerLike,
385 + max_model_len: int,
386 + ) -> tuple[list[EngineInput], list[int]]:
387 + """Build prompts by concatenating query and items.
388 +
389 + Uses the Renderer's tokenizer to tokenize text inputs, then
390 + creates EngineInput via tokens_input() for engine consumption.
391 +
392 + Args:
393 + request: The request containing query, items, and settings.
394 + tokenizer: The tokenizer to use.
395 + max_model_len: Maximum model context length for truncation.
396 +
397 + Returns:
398 + Tuple of (list of EngineInput, list of prompt token counts).
399 + """
400 + # Tokenize query if it's a string
401 + if isinstance(request.query, str):
402 + query_token_ids = tokenizer.encode(
403 + request.query,
404 + add_special_tokens=request.add_special_tokens,
405 + )
406 + else:
407 + query_token_ids = request.query
408 +
409 + engine_inputs: list[EngineInput] = []
410 + prompt_token_counts: list[int] = []
411 +
412 + for item in request.items:
413 + # Tokenize item if it's a string
414 + if isinstance(item, str):
415 + # Don't add special tokens for items to avoid duplicate BOS/EOS
416 + item_token_ids = tokenizer.encode(
417 + item,
418 + add_special_tokens=False,
419 + )
420 + else:
421 + item_token_ids = item
422 +
423 + # Concatenate based on item_first setting
424 + if request.item_first:
425 + prompt_token_ids = item_token_ids + query_token_ids
426 + else:
427 + prompt_token_ids = query_token_ids + item_token_ids
428 +
429 + # Truncate to max_model_len - 1 to leave room for 1 output token
430 + max_prompt_len = max_model_len - 1
431 + if len(prompt_token_ids) > max_prompt_len:
432 + prompt_token_ids = prompt_token_ids[:max_prompt_len]
433 +
434 + engine_inputs.append(tokens_input(prompt_token_ids))
435 + prompt_token_counts.append(len(prompt_token_ids))
436 +
437 + return engine_inputs, prompt_token_counts
438 +
439 + def _compute_probabilities(
440 + self,
441 + label_logprobs: dict[int, float],
442 + apply_softmax: bool,
443 + ) -> dict[int, float]:
444 + """Compute probabilities from logprobs.
445 +
446 + Args:
447 + label_logprobs: Dictionary mapping token_id to logprob.
448 + apply_softmax: If True, normalize over only the label tokens.
449 + If False, return true model probabilities (exp(logprob)).
450 +
451 + Returns:
452 + Dictionary mapping token_id to probability.
453 + """
454 + if apply_softmax:
455 + # Normalize over only the label tokens (subset softmax)
456 + # softmax(gathered_logits) over the subset
457 + logprobs_list = list(label_logprobs.values())
458 + max_logprob = max(logprobs_list)
459 +
460 + # Compute exp(logprob - max) for numerical stability
461 + exp_values = {
462 + token_id: math.exp(logprob - max_logprob)
463 + for token_id, logprob in label_logprobs.items()
464 + }
465 + sum_exp = sum(exp_values.values())
466 +
467 + return {
468 + token_id: exp_val / sum_exp for token_id, exp_val in exp_values.items()
469 + }
470 + else:
471 + # Return true model probabilities
472 + # Since logprobs are already log(softmax(logits)),
473 + # we just need to exp() them
474 + return {
475 + token_id: math.exp(logprob)
476 + for token_id, logprob in label_logprobs.items()
477 + }
478 +
479 + async def _get_trace_headers(
480 + self,
481 + headers: Mapping[str, str],
482 + ) -> Mapping[str, str] | None:
483 + """Extract trace headers from request headers."""
484 + if not contains_trace_headers(headers):
485 + return None
486 +
487 + if not await self.engine_client.is_tracing_enabled():
488 + log_tracing_disabled_warning()
489 + return None
490 +
491 + return extract_trace_headers(headers)
vllm/sampling_params.py
+6
@@ -232,6 +232,12 @@ class SamplingParams(
232 prompt_logprobs: int | None = None
233 """Number of log probabilities to return per prompt token.
234 When set to -1, return all `vocab_size` log probabilities."""
235 + logprob_token_ids: list[int] | None = None
236 + """Specific token IDs to return logprobs for. More efficient than
237 + logprobs=-1 when you only need logprobs for a small set of tokens.
238 + When set, logprobs for exactly these token IDs will be returned,
239 + in addition to the sampled token. This is useful for scoring tasks
240 + where you want to compare probabilities of specific label tokens."""
241 flat_logprobs: bool = False
242 """Whether to return logprobs in flatten format (i.e. FlatLogprob)
243 for better performance.
vllm/v1/sample/metadata.py
+5
@@ -40,5 +40,10 @@ class SamplingMetadata:
40 # Loaded logits processors
41 logitsprocs: LogitsProcessors
42
43 + # Specific token IDs to compute logprobs for (more efficient than full vocab)
44 + # When set, logprobs are computed only for these token IDs using gather
45 + # req_index -> list of token IDs to get logprobs for
46 + logprob_token_ids: dict[int, list[int]] | None = None
47 +
48 # Speculative token ids
49 spec_token_ids: list[list[int]] | None = None
vllm/v1/sample/sampler.py
+86 -1
@@ -13,6 +13,7 @@ from vllm.v1.sample.ops.bad_words import apply_bad_words
13 from vllm.v1.sample.ops.logprobs import batched_count_greater_than
14 from vllm.v1.sample.ops.penalties import apply_all_penalties
15 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
16 +from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs
17
18 _SAMPLING_EPS = 1e-5
19
@@ -102,8 +103,16 @@ class Sampler(nn.Module):
103 # return int32 (while PyTorch argmax and topk return int64).
104 sampled = sampled.long()
105
106 + # Handle logprob_token_ids if specified (more efficient than full vocab)
107 + # This is used by generative_scoring API to get logprobs for specific tokens
108 + logprob_token_ids_tensors = None
109 + if sampling_metadata.logprob_token_ids:
110 + logprob_token_ids_tensors = self.gather_specific_token_logprobs(
111 + logits, sampling_metadata.logprob_token_ids, sampled
112 + )
113 +
114 if num_logprobs is None:
106 - logprobs_tensors = None
115 + logprobs_tensors = logprob_token_ids_tensors
116 elif num_logprobs == -1:
117 # Return the full unsorted and unranked logprobs.
118 logprobs_tensors = LogprobsTensors(
@@ -115,6 +124,11 @@ class Sampler(nn.Module):
124 raw_logprobs, num_logprobs, token_ids=sampled
125 )
126
127 + # If we have both num_logprobs and logprob_token_ids, prefer
128 + # logprob_token_ids as it's more specific
129 + if logprob_token_ids_tensors is not None and num_logprobs is not None:
130 + logprobs_tensors = logprob_token_ids_tensors
131 +
132 # Use int32 to reduce the tensor size.
133 sampled = sampled.to(torch.int32)
134
@@ -128,6 +142,77 @@ class Sampler(nn.Module):
142 )
143 return sampler_output
144
145 + def gather_specific_token_logprobs(
146 + self,
147 + logits: torch.Tensor,
148 + logprob_token_ids: dict[int, list[int]],
149 + sampled: torch.Tensor,
150 + ) -> LogprobsTensors | None:
151 + """Compute logprobs for specific token IDs using Triton kernel.
152 +
153 + This method handles heterogeneous token ID lists across requests by
154 + padding shorter lists to max length and using a fused Triton kernel
155 + for efficient log_softmax + gather computation.
156 +
157 + Benchmarks show the Triton kernel approach is ~1.4x faster than sparse
158 + gather for batch sizes > 1 due to the fused kernel reducing memory
159 + bandwidth requirements.
160 +
161 + Args:
162 + logits: [batch_size, vocab_size] tensor of logits
163 + logprob_token_ids: dict mapping req_index -> list of token IDs
164 + sampled: [batch_size] tensor of sampled token IDs
165 +
166 + Returns:
167 + LogprobsTensors with logprobs for the specified tokens, or None
168 + if no requests have logprob_token_ids.
169 + """
170 + if not logprob_token_ids:
171 + return None
172 +
173 + batch_size = logits.shape[0]
174 + device = logits.device
175 +
176 + # Find max number of tokens across all requests
177 + max_num_tokens = max(len(tids) for tids in logprob_token_ids.values())
178 +
179 + # Create padded token_ids tensor: [batch_size, max_num_tokens + 1]
180 + # +1 for sampled token in first position
181 + token_ids_tensor = torch.zeros(
182 + batch_size, max_num_tokens + 1, dtype=torch.int64, device=device
183 + )
184 + token_ids_tensor[:, 0] = sampled # First column is sampled token
185 +
186 + # Create mask for valid positions (True = valid, False = padded)
187 + valid_mask = torch.zeros(
188 + batch_size, max_num_tokens + 1, dtype=torch.bool, device=device
189 + )
190 + valid_mask[:, 0] = True # Sampled token is always valid
191 +
192 + # Fill in token IDs for each request
193 + for req_idx, token_ids in logprob_token_ids.items():
194 + num_tokens = len(token_ids)
195 + token_ids_tensor[req_idx, 1 : num_tokens + 1] = torch.tensor(
196 + token_ids, dtype=torch.int64, device=device
197 + )
198 + valid_mask[req_idx, 1 : num_tokens + 1] = True
199 +
200 + # Compute logprobs using the fused Triton kernel (log_softmax + gather)
201 + logprobs = compute_token_logprobs(logits, token_ids_tensor)
202 +
203 + # Mask invalid (padded) positions with -inf
204 + logprobs = logprobs.masked_fill(~valid_mask, float("-inf"))
205 +
206 + # Compute ranks for the sampled token
207 + sampled_logits = logits.gather(-1, sampled.unsqueeze(-1))
208 + token_ranks = (logits > sampled_logits).sum(dim=-1)
209 +
210 + return LogprobsTensors(
211 + logprob_token_ids=token_ids_tensor.to(torch.int32),
212 + logprobs=logprobs,
213 + selected_token_ranks=token_ranks,
214 + )
215 +
216 @staticmethod
217 def apply_temperature(
218 logits: torch.Tensor,
vllm/v1/worker/gpu_input_batch.py
+19
@@ -235,6 +235,10 @@ class InputBatch:
235
236 self.num_logprobs: dict[str, int] = {}
237
238 + # req_id -> list of specific token IDs to compute logprobs for
239 + # More efficient than num_logprobs=-1 when only a few tokens are needed
240 + self.logprob_token_ids: dict[str, list[int]] = {}
241 +
242 # To accumulate prompt logprobs tensor chunks across prefill steps.
243 self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
244
@@ -395,6 +399,10 @@ class InputBatch:
399 else sampling_params.logprobs
400 )
401
402 + # Store specific token IDs to compute logprobs for (more efficient)
403 + if sampling_params.logprob_token_ids is not None:
404 + self.logprob_token_ids[req_id] = sampling_params.logprob_token_ids
405 +
406 if sampling_params.allowed_token_ids:
407 self.has_allowed_token_ids.add(req_id)
408 if self.allowed_token_ids_mask_cpu_tensor is None:
@@ -522,6 +530,7 @@ class InputBatch:
530 self.repetition_penalties_reqs.discard(req_id)
531 self.generators.pop(req_index, None)
532 self.num_logprobs.pop(req_id, None)
533 + self.logprob_token_ids.pop(req_id, None)
534 self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
535 if self.prev_req_id_to_index is not None:
536 self.prev_req_id_to_index.pop(req_id, None)
@@ -865,6 +874,15 @@ class InputBatch:
874 )
875 allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs]
876
877 + # Build per-request logprob_token_ids mapping: req_index -> token_ids
878 + logprob_token_ids_by_index: dict[int, list[int]] | None = None
879 + if self.logprob_token_ids:
880 + logprob_token_ids_by_index = {}
881 + for req_id, token_ids in self.logprob_token_ids.items():
882 + if req_id in self.req_id_to_index:
883 + req_index = self.req_id_to_index[req_id]
884 + logprob_token_ids_by_index[req_index] = token_ids
885 +
886 return SamplingMetadata(
887 temperature=temperature,
888 all_greedy=self.all_greedy,
@@ -873,6 +891,7 @@ class InputBatch:
891 top_k=None if self.no_top_k else self.top_k[:num_reqs],
892 generators=self.generators,
893 max_num_logprobs=self.max_num_logprobs,
894 + logprob_token_ids=logprob_token_ids_by_index,
895 prompt_token_ids=prompt_token_ids,
896 frequency_penalties=self.frequency_penalties[:num_reqs],
897 presence_penalties=self.presence_penalties[:num_reqs],
vllm/v1/worker/gpu_model_runner.py
+1
@@ -5600,6 +5600,7 @@ class GPUModelRunner(
5600 top_k=dummy_tensors(logits.size(1) - 1),
5601 generators={},
5602 max_num_logprobs=None,
5603 + logprob_token_ids=None,
5604 no_penalties=True,
5605 prompt_token_ids=None,
5606 frequency_penalties=dummy_tensors(0.1),