[Bugfix][Frontend] Fix batch chat endpoint corrupting logprobs when return_token_ids is set (#47384)
Signed-off-by: David Feng <fenghourun@meta.com> Co-authored-by: Claude <noreply@anthropic.com>
David Feng committed
Jul 3, 2026 at 06:01 UTC
18f658bb3185779ee58999a328246d09886d568b
3 files changed
+80
-1
tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
+76
@@ -111,3 +111,79 @@ async def test_batched_chat_completions_with_json_schema(
111
parsed = json.loads(choice["message"]["content"])
112
assert "answer" in parsed
113
assert parsed["answer"] in ("yes", "no")
114
+
115
+
116
+@pytest.mark.asyncio
117
+@pytest.mark.parametrize(
118
+ "model_name",
119
+ [MODEL_NAME],
120
+)
121
+async def test_batched_chat_completions_logprobs_not_token_id_placeholders(
122
+ server: RemoteOpenAIServer, model_name: str
123
+) -> None:
124
+ # Regression test: requesting `return_token_ids` alongside logprobs must not
125
+ # corrupt the logprob `token` fields into "token_id:{id}" placeholders. That
126
+ # placeholder rendering is controlled by `return_tokens_as_token_ids`, which
127
+ # this request leaves unset.
128
+ conversations = [
129
+ [{"role": "user", "content": "Reply with exactly the word: alpha"}],
130
+ ]
131
+
132
+ async with httpx.AsyncClient() as http_client:
133
+ response = await http_client.post(
134
+ f"{server.url_for('v1/chat/completions/batch')}",
135
+ json={
136
+ "model": model_name,
137
+ "messages": conversations,
138
+ "logprobs": True,
139
+ "top_logprobs": 1,
140
+ "return_token_ids": True,
141
+ },
142
+ timeout=60,
143
+ )
144
+
145
+ assert response.status_code == 200, response.text
146
+ data = response.json()
147
+
148
+ content = data["choices"][0]["logprobs"]["content"]
149
+ assert content
150
+ for entry in content:
151
+ assert not entry["token"].startswith("token_id:")
152
+ for top in entry["top_logprobs"]:
153
+ assert not top["token"].startswith("token_id:")
154
+
155
+
156
+@pytest.mark.asyncio
157
+@pytest.mark.parametrize(
158
+ "model_name",
159
+ [MODEL_NAME],
160
+)
161
+async def test_batched_chat_completions_return_tokens_as_token_ids(
162
+ server: RemoteOpenAIServer, model_name: str
163
+) -> None:
164
+ # Complementary check: when `return_tokens_as_token_ids` is explicitly set,
165
+ # the logprob tokens *should* be rendered as "token_id:{id}" placeholders,
166
+ # proving the new field is actually wired through.
167
+ conversations = [
168
+ [{"role": "user", "content": "Reply with exactly the word: alpha"}],
169
+ ]
170
+
171
+ async with httpx.AsyncClient() as http_client:
172
+ response = await http_client.post(
173
+ f"{server.url_for('v1/chat/completions/batch')}",
174
+ json={
175
+ "model": model_name,
176
+ "messages": conversations,
177
+ "logprobs": True,
178
+ "top_logprobs": 1,
179
+ "return_tokens_as_token_ids": True,
180
+ },
181
+ timeout=60,
182
+ )
183
+
184
+ assert response.status_code == 200, response.text
185
+ data = response.json()
186
+
187
+ content = data["choices"][0]["logprobs"]["content"]
188
+ assert content
189
+ assert all(entry["token"].startswith("token_id:") for entry in content)
vllm/entrypoints/openai/chat_completion/batch_serving.py
+1
-1
@@ -258,7 +258,7 @@ class OpenAIServingChatBatch(OpenAIServingChat):
258
top_logprobs=output.logprobs,
259
num_output_top_logprobs=request.top_logprobs,
260
tokenizer=tokenizer,
261
- return_as_token_id=request.return_token_ids,
261
+ return_as_token_id=request.return_tokens_as_token_ids,
262
)
263
else:
264
logprobs = None
vllm/entrypoints/openai/chat_completion/protocol.py
+3
@@ -1023,6 +1023,9 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
1023
include_stop_str_in_output: bool = False
1024
guided_decoding_backend: str | None = None
1025
echo: bool = False
1026
+ # None falls back to the server-level --return-tokens-as-token-ids default,
1027
+ # matching ChatCompletionRequest.return_tokens_as_token_ids.
1028
+ return_tokens_as_token_ids: bool | None = None
1029
return_token_ids: bool = False
1030
1031
@model_validator(mode="before")