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 9, 10:00 UTC
support returning tokenids in responses api (#33212)
Signed-off-by: Christian Munley <cmunley@nvidia.com>
cmunley1 committed
Jan 29, 2026 at 00:52 UTC
3bba2edb0ffb9c15daa73c56289c2e12448249d8
3 files changed
+50
-9
tests/entrypoints/openai/test_return_tokens_as_ids.py
+39
@@ -121,3 +121,42 @@ async def test_chat_return_tokens_as_token_ids_completion(server_fixture):
121
for logprob_content in response.choices[0].logprobs.content:
122
token_ids.append(int(logprob_content.token.removeprefix("token_id:")))
123
assert tokenizer.decode(token_ids, skip_special_tokens=True) == text
124
+
125
+
126
+def test_responses_api_logprobs_with_return_tokens_as_token_ids():
127
+ """Test that return_tokens_as_token_ids works in Responses API logprobs."""
128
+ from unittest.mock import MagicMock
129
+
130
+ from vllm.entrypoints.openai.engine.serving import OpenAIServing
131
+ from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
132
+ from vllm.logprobs import Logprob as SampleLogprob
133
+
134
+ serving = MagicMock(spec=OpenAIServingResponses)
135
+ serving.return_tokens_as_token_ids = True
136
+ serving._get_decoded_token = OpenAIServing._get_decoded_token
137
+
138
+ tokenizer = MagicMock()
139
+ tokenizer.decode = lambda token_id: "decoded"
140
+
141
+ token_ids = [100, 200, 300]
142
+ sample_logprobs = [
143
+ {100: SampleLogprob(logprob=-0.5, decoded_token="hello")},
144
+ {200: SampleLogprob(logprob=-1.2, decoded_token="world")},
145
+ {300: SampleLogprob(logprob=-0.8, decoded_token="!")},
146
+ ]
147
+
148
+ result = OpenAIServingResponses._create_response_logprobs(
149
+ serving,
150
+ token_ids=token_ids,
151
+ logprobs=sample_logprobs,
152
+ tokenizer=tokenizer,
153
+ top_logprobs=1,
154
+ )
155
+
156
+ assert len(result) == 3
157
+ assert result[0].token == "token_id:100"
158
+ assert result[1].token == "token_id:200"
159
+ assert result[2].token == "token_id:300"
160
+ assert result[0].logprob == -0.5
161
+ assert result[1].logprob == -1.2
162
+ assert result[2].logprob == -0.8
vllm/entrypoints/openai/engine/serving.py
+1
-1
@@ -1528,7 +1528,7 @@ class OpenAIServing:
1528
"Unable to get tokenizer because `skip_tokenizer_init=True`"
1529
)
1530
1531
- return tokenizer.decode(token_id)
1531
+ return tokenizer.decode([token_id])
1532
1533
def _is_model_supported(self, model_name: str | None) -> bool:
1534
if not model_name:
vllm/entrypoints/openai/responses/serving.py
+10
-8
@@ -836,10 +836,11 @@ class OpenAIServingResponses(OpenAIServing):
836
for i, (token_id, _logprob) in enumerate(logprobs.items()):
837
if i >= top_logprobs:
838
break
839
- text = (
840
- _logprob.decoded_token
841
- if _logprob.decoded_token is not None
842
- else tokenizer.decode([token_id])
839
+ text = self._get_decoded_token(
840
+ logprob=_logprob,
841
+ token_id=token_id,
842
+ tokenizer=tokenizer,
843
+ return_as_token_id=self.return_tokens_as_token_ids,
844
)
845
out.append(
846
LogprobTopLogprob(
@@ -865,10 +866,11 @@ class OpenAIServingResponses(OpenAIServing):
866
for i, token_id in enumerate(token_ids):
867
logprob = logprobs[i]
868
token_logprob = logprob[token_id]
868
- text = (
869
- token_logprob.decoded_token
870
- if token_logprob.decoded_token is not None
871
- else tokenizer.decode([token_id])
869
+ text = self._get_decoded_token(
870
+ logprob=token_logprob,
871
+ token_id=token_id,
872
+ tokenizer=tokenizer,
873
+ return_as_token_id=self.return_tokens_as_token_ids,
874
)
875
out.append(
876
Logprob(