From 8b8af2caf739a7537b9ca848b836733c6533bf20 Mon Sep 17 00:00:00 2001 From: langzhao-netizen Date: Mon, 13 Jul 2026 14:40:21 -0700 Subject: [PATCH] [Frontend] Expose logprob_token_ids on Python OpenAI endpoints (#43463) Signed-off-by: Lang Zhao Co-authored-by: Claude --- .../test_batched_chat_completions.py | 34 +++ .../chat_completion/test_logprob_token_ids.py | 208 ++++++++++++++++++ .../openai/completion/test_completion.py | 55 +++++ .../openai/chat_completion/batch_serving.py | 5 +- .../openai/chat_completion/protocol.py | 43 +++- .../openai/chat_completion/serving.py | 17 +- .../entrypoints/openai/completion/protocol.py | 39 +++- vllm/entrypoints/openai/completion/serving.py | 5 +- vllm/sampling_params.py | 2 + 9 files changed, 401 insertions(+), 7 deletions(-) create mode 100644 tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index 87a206b566a..b529f481889 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -187,3 +187,37 @@ async def test_batched_chat_completions_return_tokens_as_token_ids( content = data["choices"][0]["logprobs"]["content"] assert content assert all(entry["token"].startswith("token_id:") for entry in content) + + +@pytest.mark.asyncio +async def test_batched_chat_completions_logprob_token_ids( + server: RemoteOpenAIServer, +) -> None: + conversations = [[{"role": "user", "content": "Hello"}]] + + async with httpx.AsyncClient() as http_client: + response = await http_client.post( + f"{server.url_for('v1/chat/completions/batch')}", + json={ + "model": MODEL_NAME, + "messages": conversations, + "max_tokens": 1, + "temperature": 0, + "logprobs": True, + "top_logprobs": 5, + "logprob_token_ids": [100, 1000, 5000], + "return_tokens_as_token_ids": True, + }, + timeout=60, + ) + + assert response.status_code == 200, response.text + content = response.json()["choices"][0]["logprobs"]["content"] + assert content + sampled_token = content[0]["token"] + assert {entry["token"] for entry in content[0]["top_logprobs"]} == { + "token_id:100", + "token_id:1000", + "token_id:5000", + sampled_token, + } diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py new file mode 100644 index 00000000000..aa04d787ccc --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the `logprob_token_ids` field on the OpenAI-compat +chat-completion endpoint. + +`logprob_token_ids` lets a caller pin the set of vocab ids whose logprobs +should appear in the response, independent of where those ids would rank in +the natural top-k distribution. This is the primitive that multilabel +scoring postprocessors use to gather logprobs at a fixed small label +vocabulary (e.g. PII detection where each label corresponds to a known +digit-token vocab id). +""" + +import math + +import pytest +from pydantic import ValidationError + +from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest + +MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--max-model-len", + "1024", + "--max-num-seqs", + "8", + "--enforce-eager", + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +def _logprob_entries(resp) -> list: + """Return the top_logprobs list from the first generated position.""" + return resp.choices[0].logprobs.content[0].top_logprobs + + +def _token_id(token: str) -> int: + assert token.startswith("token_id:"), ( + "expected return_tokens_as_token_ids=True to yield the " + f"`token_id:` form, got {token!r}" + ) + return int(token.removeprefix("token_id:")) + + +def _top_logprob_token_ids(resp) -> list[int]: + return [_token_id(entry.token) for entry in _logprob_entries(resp)] + + +def test_chat_request_decouples_top_k_from_explicit_token_ids(): + request = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Hello"}], + logprobs=True, + top_logprobs=5, + logprob_token_ids=[5000], + ) + + sampling_params = request.to_sampling_params( + max_tokens=1, default_sampling_params={} + ) + + assert sampling_params.logprobs is None + assert sampling_params.logprob_token_ids == [5000] + assert sampling_params.num_logprobs == 1 + + +def test_completion_request_decouples_top_k_from_explicit_token_ids(): + request = CompletionRequest( + model=MODEL_NAME, + prompt="Hello", + logprobs=5, + logprob_token_ids=[5000], + ) + + sampling_params = request.to_sampling_params(max_tokens=1) + + assert sampling_params.logprobs is None + assert sampling_params.logprob_token_ids == [5000] + assert sampling_params.num_logprobs == 1 + + +def test_completion_rejects_explicit_token_ids_without_generated_tokens(): + with pytest.raises(ValidationError, match="no output tokens are generated"): + CompletionRequest( + model=MODEL_NAME, + prompt="Hello", + echo=True, + max_tokens=0, + logprobs=5, + logprob_token_ids=[5000], + ) + + +def test_requests_reject_explicit_token_ids_with_beam_search(): + with pytest.raises(ValidationError, match="not supported with beam search"): + ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Hello"}], + logprobs=True, + logprob_token_ids=[5000], + use_beam_search=True, + ) + + with pytest.raises(ValidationError, match="not supported with beam search"): + CompletionRequest( + model=MODEL_NAME, + prompt="Hello", + logprobs=5, + logprob_token_ids=[5000], + use_beam_search=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_ids", "top_logprobs", "sampled_id"), + [ + pytest.param([5000], 5, 42, id="original-issue-shape"), + pytest.param([100, 1000, 5000], 3, 42, id="sampled-outside-set"), + pytest.param([100, 1000, 5000], 3, 5000, id="sampled-inside-set"), + ], +) +async def test_logprob_token_ids_returns_requested_ids( + server, requested_ids, top_logprobs, sampled_id +): + async with server.get_async_client() as client: + resp = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Hello"}], + max_tokens=1, + temperature=0, + logprobs=True, + top_logprobs=top_logprobs, + extra_body={ + "logprob_token_ids": requested_ids, + "allowed_token_ids": [sampled_id], + "return_tokens_as_token_ids": True, + }, + ) + entries = _logprob_entries(resp) + returned_token_ids = _top_logprob_token_ids(resp) + + assert _token_id(resp.choices[0].logprobs.content[0].token) == sampled_id + assert set(returned_token_ids) == {*requested_ids, sampled_id} + assert len(returned_token_ids) == len(set(returned_token_ids)) + + for e in entries: + assert isinstance(e.logprob, float) + assert not math.isnan(e.logprob) + assert not math.isinf(e.logprob) + + +@pytest.mark.asyncio +async def test_logprob_token_ids_stream_with_no_top_k(server): + async with server.get_async_client() as client: + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Hello"}], + max_tokens=1, + temperature=0, + logprobs=True, + top_logprobs=None, + stream=True, + extra_body={ + "logprob_token_ids": [5000], + "allowed_token_ids": [42], + "return_tokens_as_token_ids": True, + }, + ) + + returned_token_ids: list[int] = [] + async for chunk in stream: + if not chunk.choices or chunk.choices[0].logprobs is None: + continue + for content in chunk.choices[0].logprobs.content: + returned_token_ids.extend( + _token_id(entry.token) for entry in content.top_logprobs + ) + + assert set(returned_token_ids) == {42, 5000} + + +@pytest.mark.asyncio +async def test_logprob_token_ids_default_behavior_unchanged(server): + """Without `logprob_token_ids`, the response carries the natural top-k + most-likely tokens. This guards against the new field accidentally + changing the default-path output.""" + async with server.get_async_client() as client: + resp = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Hello"}], + max_tokens=1, + temperature=0, + logprobs=True, + top_logprobs=5, + ) + entries = _logprob_entries(resp) + assert len(entries) == 5 + # Default mode emits log_softmax, so all values are <= 0. + for e in entries: + assert e.logprob <= 0.0 diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 7b628207ac1..2e7fd090694 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -175,6 +175,61 @@ async def test_some_logprobs(client: openai.AsyncOpenAI, model_name: str): assert 5 <= len(choice.logprobs.top_logprobs[0]) <= 6 +@pytest.mark.asyncio +async def test_logprob_token_ids(client: openai.AsyncOpenAI): + completion = await client.completions.create( + model=MODEL_NAME, + prompt="Hello", + max_tokens=1, + temperature=0.0, + logprobs=5, + extra_body={ + "logprob_token_ids": [5000], + "allowed_token_ids": [42], + "return_tokens_as_token_ids": True, + }, + ) + + choice = completion.choices[0] + assert choice.logprobs is not None + assert choice.logprobs.tokens == ["token_id:42"] + assert choice.logprobs.top_logprobs is not None + assert set(choice.logprobs.top_logprobs[0]) == { + "token_id:42", + "token_id:5000", + } + + +@pytest.mark.asyncio +async def test_logprob_token_ids_stream(client: openai.AsyncOpenAI): + stream = await client.completions.create( + model=MODEL_NAME, + prompt="Hello", + max_tokens=1, + temperature=0.0, + logprobs=5, + stream=True, + extra_body={ + "logprob_token_ids": [5000], + "allowed_token_ids": [42], + "return_tokens_as_token_ids": True, + }, + ) + + returned_top_logprobs: list[dict[str, float]] = [] + async for chunk in stream: + logprobs = chunk.choices[0].logprobs + if logprobs is not None and logprobs.top_logprobs is not None: + returned_top_logprobs.extend( + top_logprobs + for top_logprobs in logprobs.top_logprobs + if top_logprobs is not None + ) + + assert len(returned_top_logprobs) == 1 + assert set(returned_top_logprobs[0]) == {"token_id:42", "token_id:5000"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "model_name", diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index 25f7c3ec5f5..17fff37cc13 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -251,12 +251,15 @@ class OpenAIServingChatBatch(OpenAIServingChat): for output in final_res.outputs: self._raise_if_error(output.finish_reason, request_id) - if request.logprobs and request.top_logprobs is not None: + if request.logprobs and ( + request.top_logprobs is not None or request.logprob_token_ids + ): assert output.logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=output.token_ids, top_logprobs=output.logprobs, num_output_top_logprobs=request.top_logprobs, + logprob_token_ids=request.logprob_token_ids, tokenizer=tokenizer, return_as_token_id=request.return_tokens_as_token_ids, ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 3cbbe1fcac9..76a03dd7202 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -270,6 +270,18 @@ class ChatCompletionRequest(OpenAIBaseModel): ), ) prompt_logprobs: int | None = None + logprob_token_ids: list[int] | None = Field( + default=None, + description=( + "Specific vocab token IDs to return logprobs for at each generated " + "position, in addition to the sampled token. More efficient than " + "`top_logprobs=-1` when only a small fixed label set is needed " + "(e.g. multilabel scoring " + "where each label corresponds to a known vocab id). When set, " + "this explicit token selection takes precedence over the natural " + "top-k selected by `top_logprobs`. Requires `logprobs=True`." + ), + ) allowed_token_ids: list[int] | None = None bad_words: list[str] = Field(default_factory=list) # --8<-- [end:chat-completion-sampling-params] @@ -690,8 +702,13 @@ class ChatCompletionRequest(OpenAIBaseModel): seed=self.seed, stop=self.stop, stop_token_ids=stop_token_ids, - logprobs=self.top_logprobs if self.logprobs else None, + logprobs=( + self.top_logprobs + if self.logprobs and not self.logprob_token_ids + else None + ), prompt_logprobs=prompt_logprobs, + logprob_token_ids=self.logprob_token_ids or None, ignore_eos=self.ignore_eos, max_tokens=max_tokens, min_tokens=self.min_tokens, @@ -756,6 +773,18 @@ class ChatCompletionRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_logprobs(cls, data): + if data.get("logprob_token_ids") and data.get("use_beam_search"): + raise VLLMValidationError( + "`logprob_token_ids` is not supported with beam search.", + parameter="logprob_token_ids", + ) + + if data.get("logprob_token_ids") and not data.get("logprobs"): + raise VLLMValidationError( + "when using `logprob_token_ids`, `logprobs` must be set to true.", + parameter="logprob_token_ids", + ) + if (prompt_logprobs := data.get("prompt_logprobs")) is not None: if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1): raise VLLMValidationError( @@ -1009,6 +1038,14 @@ class BatchChatCompletionRequest(OpenAIBaseModel): logit_bias: dict[str, float] | None = None logprobs: bool | None = False top_logprobs: int | None = 0 + logprob_token_ids: list[int] | None = Field( + default=None, + description=( + "Specific vocab token IDs to return logprobs for at each generated " + "position, in addition to the sampled token. Requires " + "`logprobs=True`." + ), + ) max_tokens: int | None = None max_completion_tokens: int | None = None n: int | None = 1 @@ -1058,6 +1095,10 @@ class BatchChatCompletionRequest(OpenAIBaseModel): "Batch chat completions do not support beam search. " "Please set `use_beam_search` to False." ) + if data.get("logprob_token_ids") and not data.get("logprobs"): + raise ValueError( + "when using `logprob_token_ids`, `logprobs` must be set to true." + ) response_format = data.get("response_format") rf_type = ( response_format.get("type") diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 1e9f58cc2f5..5d821ba5841 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -568,13 +568,16 @@ class OpenAIServingChat(GenerateBaseServing): if finish_reason_sent[i]: continue - if request.logprobs and request.top_logprobs is not None: + if request.logprobs and ( + request.top_logprobs is not None or request.logprob_token_ids + ): assert output.logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=output.token_ids, top_logprobs=output.logprobs, tokenizer=tokenizer, num_output_top_logprobs=request.top_logprobs, + logprob_token_ids=request.logprob_token_ids, return_as_token_id=request.return_tokens_as_token_ids, ) else: @@ -862,12 +865,15 @@ class OpenAIServingChat(GenerateBaseServing): token_ids = output.token_ids out_logprobs = output.logprobs - if request.logprobs and request.top_logprobs is not None: + if request.logprobs and ( + request.top_logprobs is not None or request.logprob_token_ids + ): assert out_logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=token_ids, top_logprobs=out_logprobs, num_output_top_logprobs=request.top_logprobs, + logprob_token_ids=request.logprob_token_ids, tokenizer=tokenizer, return_as_token_id=request.return_tokens_as_token_ids, ) @@ -1101,6 +1107,7 @@ class OpenAIServingChat(GenerateBaseServing): top_logprobs: int | None, tokenizer: TokenizerLike | None, should_return_as_token_id: bool, + return_all: bool = False, ) -> list[ChatCompletionLogProb]: return [ ChatCompletionLogProb( @@ -1116,7 +1123,9 @@ class OpenAIServingChat(GenerateBaseServing): bytes=list(token.encode("utf-8", errors="replace")), ) for i, p in enumerate(logprobs.items()) - if (top_logprobs and i < top_logprobs or top_logprobs == -1) + if return_all + or top_logprobs == -1 + or (top_logprobs is not None and i < top_logprobs) ] def _create_chat_logprobs( @@ -1125,6 +1134,7 @@ class OpenAIServingChat(GenerateBaseServing): top_logprobs: GenericSequence[dict[int, Logprob] | None], tokenizer: TokenizerLike | None, num_output_top_logprobs: int | None = None, + logprob_token_ids: list[int] | None = None, return_as_token_id: bool | None = None, ) -> ChatCompletionLogProbs: """Create OpenAI-style logprobs.""" @@ -1177,6 +1187,7 @@ class OpenAIServingChat(GenerateBaseServing): num_output_top_logprobs, tokenizer, should_return_as_token_id, + return_all=bool(logprob_token_ids), ), ) ) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 11ff63e8aec..6f8171f97ba 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -95,6 +95,19 @@ class CompletionRequest(OpenAIBaseModel): ) allowed_token_ids: list[int] | None = None prompt_logprobs: int | None = None + logprob_token_ids: list[int] | None = Field( + default=None, + description=( + "Specific vocab token IDs to return logprobs for at each generated " + "position, in addition to the sampled token. More efficient than " + "requesting the full vocab when only a small fixed label set is " + "needed (e.g. multilabel " + "scoring where each label corresponds to a known vocab id). When " + "set, this explicit token selection takes precedence over the " + "natural top-k selected by `logprobs`. Requires `logprobs` to be " + "set." + ), + ) bad_words: list[str] = Field(default_factory=list) # --8<-- [end:completion-sampling-params] @@ -368,11 +381,12 @@ class CompletionRequest(OpenAIBaseModel): seed=self.seed, stop=self.stop, stop_token_ids=stop_token_ids, - logprobs=self.logprobs, + logprobs=None if self.logprob_token_ids else self.logprobs, ignore_eos=self.ignore_eos, max_tokens=max_tokens if not echo_without_generation else 1, min_tokens=self.min_tokens, prompt_logprobs=prompt_logprobs, + logprob_token_ids=self.logprob_token_ids or None, skip_special_tokens=self.skip_special_tokens, spaces_between_special_tokens=self.spaces_between_special_tokens, include_stop_str_in_output=self.include_stop_str_in_output, @@ -459,6 +473,29 @@ class CompletionRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_logprobs(cls, data): + if data.get("logprob_token_ids") and data.get("use_beam_search"): + raise VLLMValidationError( + "`logprob_token_ids` is not supported with beam search.", + parameter="logprob_token_ids", + ) + + if ( + data.get("logprob_token_ids") + and data.get("echo") + and data.get("max_tokens") == 0 + ): + raise VLLMValidationError( + "`logprob_token_ids` is not supported when `echo=True` and " + "`max_tokens=0` because no output tokens are generated.", + parameter="logprob_token_ids", + ) + + if data.get("logprob_token_ids") and data.get("logprobs") is None: + raise VLLMValidationError( + "when using `logprob_token_ids`, `logprobs` must be set.", + parameter="logprob_token_ids", + ) + if (prompt_logprobs := data.get("prompt_logprobs")) is not None: if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1): raise VLLMValidationError( diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 545ed965fee..8791042c692 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -386,6 +386,7 @@ class OpenAIServingCompletion(GenerateBaseServing): top_logprobs=out_logprobs, num_output_top_logprobs=request.logprobs, tokenizer=tokenizer, + logprob_token_ids=request.logprob_token_ids, initial_text_offset=previous_text_lens[i], return_as_token_id=request.return_tokens_as_token_ids, ) @@ -560,6 +561,7 @@ class OpenAIServingCompletion(GenerateBaseServing): top_logprobs=out_logprobs, tokenizer=tokenizer, num_output_top_logprobs=request.logprobs, + logprob_token_ids=request.logprob_token_ids, return_as_token_id=request.return_tokens_as_token_ids, ) else: @@ -653,6 +655,7 @@ class OpenAIServingCompletion(GenerateBaseServing): top_logprobs: GenericSequence[dict[int, Logprob] | None], num_output_top_logprobs: int, tokenizer: TokenizerLike | None, + logprob_token_ids: list[int] | None = None, initial_text_offset: int = 0, return_as_token_id: bool | None = None, ) -> CompletionLogProbs: @@ -717,7 +720,7 @@ class OpenAIServingCompletion(GenerateBaseServing): return_as_token_id=should_return_as_token_id, ): max(top_lp[1].logprob, -9999.0) for i, top_lp in enumerate(step_top_logprobs.items()) - if num_output_top_logprobs >= i + if logprob_token_ids or num_output_top_logprobs >= i } ) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index f0966902d36..08580f6e8f6 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -383,6 +383,7 @@ class SamplingParams( extra_args: dict[str, Any] | None = None, skip_clone: bool = False, repetition_detection: RepetitionDetectionParams | None = None, + logprob_token_ids: list[int] | None = None, ) -> "SamplingParams": if logit_bias is not None: # Fast path uses a dict comprehension; on failure we iterate once @@ -433,6 +434,7 @@ class SamplingParams( min_tokens=min_tokens, logprobs=logprobs, prompt_logprobs=prompt_logprobs, + logprob_token_ids=logprob_token_ids, detokenize=detokenize, skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens,