obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-20 12:40:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Frontend] Expose logprob_token_ids on Python OpenAI endpoints (#43463)

Browse Source 
Signed-off-by: Lang Zhao <[email protected]>
Co-authored-by: Claude <[email protected]>
This commit is contained in:
langzhao-netizen
2026-07-13 14:40:21 -07:00
committed by GitHub
co-authored by Claude
parent 7738ef35b8
commit 8b8af2caf7
9 changed files with 401 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
+34
View File
@@ -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,
}
tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py
+208
View File
@@ -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:<int>` 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
tests/entrypoints/openai/completion/test_completion.py
+55
View File
@@ -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",
vllm/entrypoints/openai/chat_completion/batch_serving.py
+4 -1
View File
@@ -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,
)
vllm/entrypoints/openai/chat_completion/protocol.py
+42 -1
View File
@@ -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")
vllm/entrypoints/openai/chat_completion/serving.py
+14 -3
View File
@@ -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),
),
)
)
vllm/entrypoints/openai/completion/protocol.py
+38 -1
View File
@@ -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(
vllm/entrypoints/openai/completion/serving.py
+4 -1
View File
@@ -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
}
)
vllm/sampling_params.py
+2
View File
@@ -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,