@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 24f4d1a224
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 14, 05:30 UTC

Add return_token_ids parameter to OpenAI API endpoints (#22587)

Signed-off-by: Yuge Zhang <scottyugochang@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Simon Mo <simon.mo@hey.com>

Yuge Zhang committed Aug 20, 2025 at 00:48 UTC 24f4d1a224ef1da6cb3878628c39ad09de480c98
5 files changed +480 -30
tests/entrypoints/openai/test_openai_schema.py
+38 -25
@@ -74,31 +74,44 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
74 -d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
75 http://localhost:8000/v1/chat/completions
76 """ # noqa: E501
77 - if (hasattr(case, "body") and isinstance(case.body, dict)
78 - and "messages" in case.body
79 - and isinstance(case.body["messages"], list)
80 - and len(case.body["messages"]) > 0):
81 -
82 - for message in case.body["messages"]:
83 - if not isinstance(message, dict):
84 - continue
85 -
86 - # Check for invalid file type in tokenize endpoint
87 - if op.method.lower() == "post" and op.path == "/tokenize":
88 - content = message.get("content", [])
89 - if (isinstance(content, list) and len(content) > 0 and any(
90 - item.get("type") == "file" for item in content)):
91 - return False
92 -
93 - # Check for invalid tool_calls with non-function types
94 - tool_calls = message.get("tool_calls", [])
95 - if isinstance(tool_calls, list):
96 - for tool_call in tool_calls:
97 - if isinstance(tool_call, dict):
98 - if tool_call.get("type") != "function":
99 - return False
100 - if "custom" in tool_call:
101 - return False
77 + if hasattr(case, "body") and isinstance(case.body, dict):
78 + if ("messages" in case.body
79 + and isinstance(case.body["messages"], list)
80 + and len(case.body["messages"]) > 0):
81 +
82 + for message in case.body["messages"]:
83 + if not isinstance(message, dict):
84 + continue
85 +
86 + # Check for invalid file type in tokenize endpoint
87 + if op.method.lower() == "post" and op.path == "/tokenize":
88 + content = message.get("content", [])
89 + if (isinstance(content, list) and len(content) > 0
90 + and any(
91 + item.get("type") == "file"
92 + for item in content)):
93 + return False
94 +
95 + # Check for invalid tool_calls with non-function types
96 + tool_calls = message.get("tool_calls", [])
97 + if isinstance(tool_calls, list):
98 + for tool_call in tool_calls:
99 + if isinstance(tool_call, dict):
100 + if tool_call.get("type") != "function":
101 + return False
102 + if "custom" in tool_call:
103 + return False
104 +
105 + # Sometimes guided_grammar is generated to be empty
106 + # Causing a server error in EBNF grammar parsing
107 + # https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
108 + guided_grammar = case.body.get("guided_grammar")
109 +
110 + if guided_grammar == '':
111 + # Allow None (will be handled as no grammar)
112 + # But skip empty strings
113 + return False
114 +
115 return True
116
117 return strategy.filter(no_invalid_types)
tests/entrypoints/openai/test_return_token_ids.py new
+374
@@ -0,0 +1,374 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.transformers_utils.tokenizer import get_tokenizer
7 +
8 +from ...utils import RemoteOpenAIServer
9 +
10 +MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
11 +
12 +
13 +@pytest.fixture(scope="module")
14 +def server():
15 + args = [
16 + "--max-model-len",
17 + "2048",
18 + "--max-num-seqs",
19 + "128",
20 + "--enable-auto-tool-choice",
21 + "--tool-call-parser",
22 + "hermes",
23 + "--enforce-eager",
24 + ]
25 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
26 + yield remote_server
27 +
28 +
29 +@pytest.mark.asyncio
30 +async def test_basic_completion_with_emoji(server):
31 + """Test basic completion with emoji to verify token_ids field."""
32 + async with server.get_async_client() as client:
33 + # Test with return_token_ids enabled
34 + completion = await client.completions.create(
35 + model=MODEL_NAME,
36 + prompt="Complete this sentence with emojis: I love coding 🚀",
37 + max_tokens=10,
38 + temperature=0,
39 + logprobs=1,
40 + extra_body={"return_token_ids": True},
41 + )
42 +
43 + # Check the raw response to see the structure
44 + completion_dict = completion.model_dump()
45 +
46 + # Verify prompt_token_ids field is present in the completion response
47 + assert "prompt_token_ids" in completion_dict["choices"][0]
48 + assert isinstance(completion.choices[0].prompt_token_ids, list)
49 +
50 + # Check against the expected prompt token IDs
51 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
52 + encoded_tokens = tokenizer.encode(
53 + "Complete this sentence with emojis: I love coding 🚀")
54 + # Check that encoded_tokens is a subsequence of prompt_token_ids
55 + assert any(completion.choices[0].prompt_token_ids[i:i +
56 + len(encoded_tokens)]
57 + == encoded_tokens for i in range(
58 + len(completion.choices[0].prompt_token_ids) -
59 + len(encoded_tokens) + 1))
60 +
61 + # Verify token_ids field is present in the choice
62 + assert completion.choices[0].token_ids is not None
63 + assert isinstance(completion.choices[0].token_ids, list)
64 + assert len(completion.choices[0].token_ids) > 0
65 +
66 + # Verify decoding works correctly
67 + decoded_text = tokenizer.decode(completion.choices[0].token_ids)
68 + # The decoded text should contain a <|im_end|> at the end
69 + assert decoded_text.startswith(completion.choices[0].text)
70 +
71 + # Test without return_token_ids (should be None)
72 + completion_without = await client.completions.create(
73 + model=MODEL_NAME,
74 + prompt="Complete this sentence with emojis: I love coding 🚀",
75 + max_tokens=10,
76 + temperature=0,
77 + logprobs=1,
78 + extra_body={"return_token_ids": False},
79 + )
80 +
81 + completion_without_dict = completion_without.model_dump()
82 + assert completion_without_dict["choices"][0].get("token_ids") is None
83 + assert completion_without_dict.get("prompt_token_ids") is None
84 +
85 +
86 +@pytest.mark.asyncio
87 +async def test_chat_completion_with_tool_use(server):
88 + """Test chat completion with tool use (get_weather function)."""
89 + tools = [{
90 + "type": "function",
91 + "function": {
92 + "name": "get_weather",
93 + "description": "Get the current weather in a given location",
94 + "parameters": {
95 + "type": "object",
96 + "properties": {
97 + "location": {
98 + "type":
99 + "string",
100 + "description":
101 + "The city and state, e.g. San Francisco, CA",
102 + },
103 + "unit": {
104 + "type": "string",
105 + "enum": ["celsius", "fahrenheit"],
106 + "description": "The unit of temperature",
107 + },
108 + },
109 + "required": ["location"],
110 + },
111 + },
112 + }]
113 +
114 + async with server.get_async_client() as client:
115 + # Test with return_token_ids enabled
116 + response = await client.chat.completions.create(
117 + model=MODEL_NAME,
118 + messages=[
119 + {
120 + "role": "system",
121 + "content": "You are a helpful assistant."
122 + },
123 + {
124 + "role": "user",
125 + "content": "What's the weather like in Paris?"
126 + },
127 + ],
128 + tools=tools,
129 + tool_choice="auto",
130 + max_tokens=100,
131 + temperature=0,
132 + logprobs=True,
133 + extra_body={"return_token_ids": True},
134 + )
135 +
136 + # Verify token_ids field is present in choices
137 + assert response.choices[0].token_ids is not None
138 + assert isinstance(response.choices[0].token_ids, list)
139 +
140 + # Verify prompt_token_ids field is present
141 + assert response.prompt_token_ids is not None
142 + assert isinstance(response.prompt_token_ids, list)
143 +
144 + # Verify the prompt texts and response texts
145 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
146 + prompt_text = tokenizer.decode(response.prompt_token_ids)
147 + assert prompt_text.startswith(
148 + "<|im_start|>system\nYou are a helpful assistant.")
149 + assert prompt_text.endswith(
150 + "What's the weather like in Paris?<|im_end|>\n"
151 + "<|im_start|>assistant\n")
152 +
153 + response_text = tokenizer.decode(response.choices[0].token_ids)
154 + assert response_text.startswith('<tool_call>\n{"name": "get_weather"')
155 + assert response_text.endswith("</tool_call><|im_end|>")
156 +
157 + # If tool call was made, verify the response structure
158 + if response.choices[0].message.tool_calls:
159 + assert len(response.choices[0].message.tool_calls) > 0
160 + tool_call = response.choices[0].message.tool_calls[0]
161 + assert tool_call.function.name == "get_weather"
162 +
163 + # Test without return_token_ids
164 + response_without = await client.chat.completions.create(
165 + model=MODEL_NAME,
166 + messages=[
167 + {
168 + "role": "system",
169 + "content": "You are a helpful assistant."
170 + },
171 + {
172 + "role": "user",
173 + "content": "What's the weather like in Paris?"
174 + },
175 + ],
176 + tools=tools,
177 + tool_choice="auto",
178 + max_tokens=100,
179 + temperature=0,
180 + logprobs=True,
181 + extra_body={"return_token_ids": False},
182 + )
183 +
184 + assert response_without.choices[0].token_ids is None
185 + assert response_without.prompt_token_ids is None
186 +
187 +
188 +@pytest.mark.asyncio
189 +async def test_comparison_with_prompt_logprobs_and_logprobs(server):
190 + """
191 + Test that token_ids align with prompt_logprobs and
192 + logprobs when return_tokens_as_token_ids is enabled.
193 + """
194 + async with server.get_async_client() as client:
195 + # Test with both return_token_ids and return_tokens_as_token_ids enabled
196 + completion = await client.completions.create(
197 + model=MODEL_NAME,
198 + prompt="Hello, world! How are you today?",
199 + max_tokens=20,
200 + temperature=0,
201 + echo=True,
202 + logprobs=1,
203 + extra_body={
204 + "return_token_ids": True,
205 + "return_tokens_as_token_ids": True,
206 + "prompt_logprobs": 1
207 + },
208 + )
209 +
210 + # Verify all fields are present
211 + assert completion.choices[0].token_ids is not None
212 + assert completion.choices[0].prompt_token_ids is not None
213 + assert completion.choices[0].prompt_logprobs is not None
214 + assert completion.choices[0].logprobs is not None
215 +
216 + # Extract token IDs from logprobs
217 + # (when return_tokens_as_token_ids is True)
218 + logprobs_token_ids = []
219 + for token_str in completion.choices[0].logprobs.tokens:
220 + # Token format is "token_id:12345" when
221 + # return_tokens_as_token_ids is True
222 + if token_str.startswith("token_id:"):
223 + token_id = int(token_str.removeprefix("token_id:"))
224 + logprobs_token_ids.append(token_id)
225 +
226 + # When echo=True, the logprobs include both prompt and response tokens
227 + # The token_ids field should match the the suffix of response portion
228 + # The prompt_token_ids should match the prompt portion
229 + assert len(completion.choices[0].token_ids) < len(logprobs_token_ids)
230 + response_token_ids_length = len(completion.choices[0].token_ids)
231 + assert logprobs_token_ids[-response_token_ids_length:] == \
232 + completion.choices[0].token_ids
233 +
234 + # Verify tokenizer consistency
235 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
236 +
237 + # Decode prompt tokens
238 + if completion.choices[0].prompt_token_ids:
239 + prompt_text = tokenizer.decode(
240 + completion.choices[0].prompt_token_ids)
241 + # The decoded prompt should match or close to original prompt
242 + assert "Hello, world" in prompt_text
243 +
244 + # Decode response tokens
245 + if completion.choices[0].token_ids:
246 + response_text = tokenizer.decode(completion.choices[0].token_ids)
247 + assert completion.choices[0].text.endswith(response_text)
248 +
249 + # Test streaming mode
250 + stream = await client.completions.create(
251 + model=MODEL_NAME,
252 + prompt="Tell me a short fact about Python:",
253 + max_tokens=30,
254 + temperature=0,
255 + stream=True,
256 + echo=False,
257 + logprobs=1,
258 + extra_body={
259 + "return_token_ids": True,
260 + "return_tokens_as_token_ids": True
261 + },
262 + )
263 +
264 + # Collect streamed tokens
265 + streamed_prompt_token_ids = []
266 + streamed_token_ids = []
267 + streamed_logprob_token_ids = []
268 + first_chunk = True
269 + async for chunk in stream:
270 + for token_str in chunk.choices[0].logprobs.tokens:
271 + # Token format is "token_id:12345" when
272 + # return_tokens_as_token_ids is True
273 + if token_str.startswith("token_id:"):
274 + token_id = int(token_str.removeprefix("token_id:"))
275 + streamed_logprob_token_ids.append(token_id)
276 + if first_chunk:
277 + streamed_prompt_token_ids = chunk.choices[0].prompt_token_ids
278 + first_chunk = False
279 + streamed_token_ids += chunk.choices[0].token_ids
280 +
281 + # Verify we collected some tokens and first chunk had prompt_token_ids
282 + assert len(streamed_prompt_token_ids) > 0
283 + assert streamed_token_ids == streamed_logprob_token_ids
284 +
285 +
286 +@pytest.mark.asyncio
287 +async def test_chat_completion_with_emoji_and_token_ids(server):
288 + """Test chat completion with emojis to verify token_ids handling."""
289 + chat_messages = [
290 + {
291 + "role": "system",
292 + "content": "You like to use emojis in your responses."
293 + },
294 + {
295 + "role": "user",
296 + "content": "Repeat after me: I love cats 🐱"
297 + },
298 + ]
299 + async with server.get_async_client() as client:
300 + response = await client.chat.completions.create(
301 + model=MODEL_NAME,
302 + messages=chat_messages,
303 + max_tokens=50,
304 + temperature=0,
305 + logprobs=True,
306 + extra_body={"return_token_ids": True},
307 + )
308 +
309 + # Verify token_ids are present
310 + response_dict = response.model_dump()
311 + assert response.choices[0].token_ids is not None
312 + assert "prompt_token_ids" in response_dict
313 +
314 + # Verify the response contains the expected fields
315 + assert response.choices[0].message.content is not None
316 +
317 + # Decode token_ids and verify consistency
318 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
319 +
320 + decoded_prompt = tokenizer.decode(response.prompt_token_ids)
321 + assert decoded_prompt.startswith(
322 + "<|im_start|>system\nYou like to use emojis in your responses.")
323 + assert decoded_prompt.endswith(
324 + "I love cats 🐱<|im_end|>\n<|im_start|>assistant\n")
325 +
326 + decoded_response = tokenizer.decode(response.choices[0].token_ids)
327 + # The content should match the response text
328 + # except the ending <|im_end|>
329 + assert decoded_response == response.choices[
330 + 0].message.content + "<|im_end|>"
331 +
332 + # Test with streaming
333 + stream = await client.chat.completions.create(
334 + model=MODEL_NAME,
335 + messages=chat_messages,
336 + max_tokens=50,
337 + temperature=0,
338 + stream=True,
339 + extra_body={"return_token_ids": True},
340 + )
341 +
342 + collected_content = ""
343 + collected_token_ids = []
344 + first_chunk = True
345 +
346 + async for chunk in stream:
347 + if first_chunk:
348 + assert chunk.prompt_token_ids is not None
349 + assert isinstance(chunk.prompt_token_ids, list)
350 + # Check the prompt_token_ids match the initial prompt
351 + decoded_prompt_stream = tokenizer.decode(
352 + chunk.prompt_token_ids)
353 + assert decoded_prompt_stream == decoded_prompt
354 + first_chunk = False
355 + else:
356 + chunk_dump = chunk.model_dump()
357 + assert "prompt_token_ids" not in chunk_dump, \
358 + "Subsequent chunks should not have prompt_token_ids"
359 +
360 + if chunk.choices:
361 + if chunk.choices[0].delta.content:
362 + collected_content += chunk.choices[0].delta.content
363 + # token_ids may not present in all chunks
364 + choice_dump = chunk.choices[0].model_dump()
365 + if "token_ids" in choice_dump:
366 + collected_token_ids.extend(chunk.choices[0].token_ids)
367 +
368 + # Verify we got response and token_ids
369 + assert len(collected_content) > 0
370 + assert len(collected_token_ids) > 0
371 +
372 + # Verify token_ids decode properly
373 + decoded_response = tokenizer.decode(collected_token_ids)
374 + assert decoded_response == collected_content + "<|im_end|>"
vllm/entrypoints/openai/protocol.py
+30
@@ -576,6 +576,14 @@ class ChatCompletionRequest(OpenAIBaseModel):
576 "If specified with 'logprobs', tokens are represented "
577 " as strings of the form 'token_id:{token_id}' so that tokens "
578 "that are not JSON-encodable can be identified."))
579 + return_token_ids: Optional[bool] = Field(
580 + default=None,
581 + description=(
582 + "If specified, the result will include token IDs alongside the "
583 + "generated text. In streaming mode, prompt_token_ids is included "
584 + "only in the first chunk, and token_ids contains the delta tokens "
585 + "for each chunk. This is useful for debugging or when you "
586 + "need to map generated text back to input tokens."))
587 cache_salt: Optional[str] = Field(
588 default=None,
589 description=(
@@ -1062,6 +1070,14 @@ class CompletionRequest(OpenAIBaseModel):
1070 "If specified with 'logprobs', tokens are represented "
1071 " as strings of the form 'token_id:{token_id}' so that tokens "
1072 "that are not JSON-encodable can be identified."))
1073 + return_token_ids: Optional[bool] = Field(
1074 + default=None,
1075 + description=(
1076 + "If specified, the result will include token IDs alongside the "
1077 + "generated text. In streaming mode, prompt_token_ids is included "
1078 + "only in the first chunk, and token_ids contains the delta tokens "
1079 + "for each chunk. This is useful for debugging or when you "
1080 + "need to map generated text back to input tokens."))
1081
1082 cache_salt: Optional[str] = Field(
1083 default=None,
@@ -1480,7 +1496,9 @@ class CompletionResponseChoice(OpenAIBaseModel):
1496 "to stop, None if the completion finished for some other reason "
1497 "including encountering the EOS token"),
1498 )
1499 + token_ids: Optional[list[int]] = None # For response
1500 prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None
1501 + prompt_token_ids: Optional[list[int]] = None # For prompt
1502
1503
1504 class CompletionResponse(OpenAIBaseModel):
@@ -1511,6 +1529,10 @@ class CompletionResponseStreamChoice(OpenAIBaseModel):
1529 "to stop, None if the completion finished for some other reason "
1530 "including encountering the EOS token"),
1531 )
1532 + # not part of the OpenAI spec but for tracing the tokens
1533 + # prompt tokens is put into choice to align with CompletionResponseChoice
1534 + prompt_token_ids: Optional[list[int]] = None
1535 + token_ids: Optional[list[int]] = None
1536
1537
1538 class CompletionStreamResponse(OpenAIBaseModel):
@@ -1680,6 +1702,9 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
1702 finish_reason: Optional[str] = "stop"
1703 # not part of the OpenAI spec but included in vLLM for legacy reasons
1704 stop_reason: Optional[Union[int, str]] = None
1705 + # not part of the OpenAI spec but is useful for tracing the tokens
1706 + # in agent scenarios
1707 + token_ids: Optional[list[int]] = None
1708
1709
1710 class ChatCompletionResponse(OpenAIBaseModel):
@@ -1695,6 +1720,7 @@ class ChatCompletionResponse(OpenAIBaseModel):
1720
1721 # vLLM-specific fields that are not in OpenAI spec
1722 prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None
1723 + prompt_token_ids: Optional[list[int]] = None
1724 kv_transfer_params: Optional[dict[str, Any]] = Field(
1725 default=None, description="KVTransfer parameters.")
1726
@@ -1712,6 +1738,8 @@ class ChatCompletionResponseStreamChoice(OpenAIBaseModel):
1738 logprobs: Optional[ChatCompletionLogProbs] = None
1739 finish_reason: Optional[str] = None
1740 stop_reason: Optional[Union[int, str]] = None
1741 + # not part of the OpenAI spec but for tracing the tokens
1742 + token_ids: Optional[list[int]] = None
1743
1744
1745 class ChatCompletionStreamResponse(OpenAIBaseModel):
@@ -1721,6 +1749,8 @@ class ChatCompletionStreamResponse(OpenAIBaseModel):
1749 model: str
1750 choices: list[ChatCompletionResponseStreamChoice]
1751 usage: Optional[UsageInfo] = Field(default=None)
1752 + # not part of the OpenAI spec but for tracing the tokens
1753 + prompt_token_ids: Optional[list[int]] = None
1754
1755
1756 class TranscriptionResponseStreamChoice(OpenAIBaseModel):
vllm/entrypoints/openai/serving_chat.py
+18 -4
@@ -568,12 +568,17 @@ class OpenAIServingChat(OpenAIServing):
568 ),
569 logprobs=None,
570 finish_reason=None)
571 +
572 + # return prompt_token_ids at the first chunk ever
573 chunk = ChatCompletionStreamResponse(
574 id=request_id,
575 object=chunk_object_type,
576 created=created_time,
577 choices=[choice_data],
576 - model=model_name)
578 + model=model_name,
579 + prompt_token_ids=(res.prompt_token_ids
580 + if request.return_token_ids else
581 + None))
582
583 # if continuous usage stats are requested, add it
584 if include_continuous_usage:
@@ -912,7 +917,9 @@ class OpenAIServingChat(OpenAIServing):
917 index=i,
918 delta=delta_message,
919 logprobs=logprobs,
915 - finish_reason=None)
920 + finish_reason=None,
921 + token_ids=(as_list(output.token_ids)
922 + if request.return_token_ids else None))
923
924 # if the model is finished generating
925 else:
@@ -973,7 +980,9 @@ class OpenAIServingChat(OpenAIServing):
980 logprobs=logprobs,
981 finish_reason=output.finish_reason
982 if not auto_tools_called else "tool_calls",
976 - stop_reason=output.stop_reason)
983 + stop_reason=output.stop_reason,
984 + token_ids=(as_list(output.token_ids)
985 + if request.return_token_ids else None))
986
987 finish_reason_sent[i] = True
988
@@ -1260,7 +1269,10 @@ class OpenAIServingChat(OpenAIServing):
1269 logprobs=logprobs,
1270 finish_reason="tool_calls" if auto_tools_called else
1271 output.finish_reason if output.finish_reason else "stop",
1263 - stop_reason=output.stop_reason)
1272 + stop_reason=output.stop_reason,
1273 + token_ids=(as_list(output.token_ids)
1274 + if request.return_token_ids else None),
1275 + )
1276
1277 choices.append(choice_data)
1278
@@ -1301,6 +1313,8 @@ class OpenAIServingChat(OpenAIServing):
1313 choices=choices,
1314 usage=usage,
1315 prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
1316 + prompt_token_ids=(final_res.prompt_token_ids
1317 + if request.return_token_ids else None),
1318 kv_transfer_params=final_res.kv_transfer_params,
1319 )
1320
vllm/entrypoints/openai/serving_completion.py
+20 -1
@@ -42,7 +42,7 @@ from vllm.outputs import RequestOutput
42 from vllm.sampling_params import BeamSearchParams, SamplingParams
43 from vllm.sequence import Logprob
44 from vllm.transformers_utils.tokenizer import AnyTokenizer
45 -from vllm.utils import merge_async_iterators
45 +from vllm.utils import as_list, merge_async_iterators
46
47 logger = init_logger(__name__)
48
@@ -365,6 +365,11 @@ class OpenAIServingCompletion(OpenAIServing):
365 for output in res.outputs:
366 i = output.index + prompt_idx * num_choices
367
368 + # Useful when request.return_token_ids is True
369 + # Returning prompt token IDs shares the same logic
370 + # with the echo implementation.
371 + prompt_token_ids_to_return: Optional[list[int]] = None
372 +
373 assert request.max_tokens is not None
374 if request.echo and not has_echoed[i]:
375 assert prompt_token_ids is not None
@@ -385,6 +390,7 @@ class OpenAIServingCompletion(OpenAIServing):
390 *(prompt_logprobs or []),
391 *(output.logprobs or []),
392 ]
393 + prompt_token_ids_to_return = prompt_token_ids
394 has_echoed[i] = True
395 else:
396 # return just the delta
@@ -392,6 +398,12 @@ class OpenAIServingCompletion(OpenAIServing):
398 delta_token_ids = output.token_ids
399 out_logprobs = output.logprobs
400
401 + # has_echoed[i] is reused here to indicate whether
402 + # we have already returned the prompt token IDs.
403 + if not has_echoed[i]:
404 + prompt_token_ids_to_return = prompt_token_ids
405 + has_echoed[i] = True
406 +
407 if (not delta_text and not delta_token_ids
408 and not previous_num_tokens[i]):
409 # Chunked prefill case, don't return empty chunks
@@ -428,6 +440,9 @@ class OpenAIServingCompletion(OpenAIServing):
440 logprobs=logprobs,
441 finish_reason=finish_reason,
442 stop_reason=stop_reason,
443 + prompt_token_ids=prompt_token_ids_to_return,
444 + token_ids=(as_list(output.token_ids) if
445 + request.return_token_ids else None),
446 )
447 ],
448 )
@@ -548,6 +563,10 @@ class OpenAIServingCompletion(OpenAIServing):
563 finish_reason=output.finish_reason,
564 stop_reason=output.stop_reason,
565 prompt_logprobs=final_res.prompt_logprobs,
566 + prompt_token_ids=(prompt_token_ids
567 + if request.return_token_ids else None),
568 + token_ids=(as_list(output.token_ids)
569 + if request.return_token_ids else None),
570 )
571 choices.append(choice_data)
572