From 410ab4fde5791ed7007c9b21a88e75d691e41f9d Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:41:52 -0700 Subject: [PATCH] Add return_token_ids support to completions and chat completions APIs (#30917) --- .../sglang/srt/entrypoints/openai/protocol.py | 17 ++++ .../srt/entrypoints/openai/serving_chat.py | 15 +++- .../entrypoints/openai/serving_completions.py | 32 ++++++- .../unit/entrypoints/openai/test_protocol.py | 3 + .../entrypoints/openai/test_serving_chat.py | 28 ++++--- .../openai/test_serving_completions.py | 83 ++++++++++++++++++- 6 files changed, 159 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 1cd8e12fa..824fb6b22 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -342,6 +342,7 @@ class CompletionRequest(BaseModel): return_routed_experts: bool = False routed_experts_start_len: int = 0 return_cached_tokens_details: bool = False + return_token_ids: bool = False # Extra parameters for SRT backend only and will be ignored by OpenAI models. top_k: int = -1 @@ -426,12 +427,18 @@ class CompletionResponseChoice(BaseModel): finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None matched_stop: Union[None, int, str] = None hidden_states: Optional[object] = None + token_ids: Optional[List[int]] = None + prompt_token_ids: Optional[List[int]] = None @model_serializer(mode="wrap") def _serialize(self, handler): data = handler(self) if self.hidden_states is None: data.pop("hidden_states", None) + if self.token_ids is None: + data.pop("token_ids", None) + if self.prompt_token_ids is None: + data.pop("prompt_token_ids", None) return data @@ -460,12 +467,18 @@ class CompletionResponseStreamChoice(BaseModel): finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None matched_stop: Union[None, int, str] = None hidden_states: Optional[object] = None + token_ids: Optional[List[int]] = None + prompt_token_ids: Optional[List[int]] = None @model_serializer(mode="wrap") def _serialize(self, handler): data = handler(self) if self.hidden_states is None: data.pop("hidden_states", None) + if self.token_ids is None: + data.pop("token_ids", None) + if self.prompt_token_ids is None: + data.pop("prompt_token_ids", None) return data @@ -746,6 +759,7 @@ class ChatCompletionRequest(BaseModel): routed_experts_start_len: int = 0 return_cached_tokens_details: bool = False return_prompt_token_ids: bool = False + return_token_ids: bool = False return_meta_info: bool = False reasoning_effort: ReasoningEffortType = Field( default=None, @@ -1056,6 +1070,7 @@ class ChatCompletionResponseChoice(BaseModel): matched_stop: Union[None, int, str] = None hidden_states: Optional[object] = None prompt_token_ids: Optional[List[int]] = None + token_ids: Optional[List[int]] = None meta_info: Optional[Dict[str, Any]] = None @model_serializer(mode="wrap") @@ -1065,6 +1080,8 @@ class ChatCompletionResponseChoice(BaseModel): data.pop("hidden_states", None) if self.prompt_token_ids is None: data.pop("prompt_token_ids", None) + if self.token_ids is None: + data.pop("token_ids", None) if self.meta_info is None: data.pop("meta_info", None) return data diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 021d22283..dcdd60bad 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -682,6 +682,12 @@ class OpenAIServingChat(OpenAIServingBase): "return_prompt_token_ids is not supported with streaming. " "Please set stream=false when using return_prompt_token_ids=true." ) + if request.return_token_ids: + raise ValueError( + "return_token_ids is not supported with streaming on " + "/v1/chat/completions. Please set stream=false when using " + "return_token_ids=true." + ) if request.return_meta_info: raise ValueError( "return_meta_info is not supported with streaming. " @@ -771,7 +777,8 @@ class OpenAIServingChat(OpenAIServingBase): video_max_dynamic_patch=vid_max_dynamic_patch, max_dynamic_patch=getattr(request, "max_dynamic_patch", None), use_audio_in_video=getattr(request, "use_audio_in_video", False), - return_prompt_token_ids=request.return_prompt_token_ids, + return_prompt_token_ids=request.return_prompt_token_ids + or request.return_token_ids, ) return adapted_request, request @@ -1539,9 +1546,12 @@ class OpenAIServingChat(OpenAIServingBase): # Extract prompt_token_ids if requested choice_prompt_token_ids = ( ret_item.get("prompt_token_ids") - if request.return_prompt_token_ids + if request.return_prompt_token_ids or request.return_token_ids else None ) + choice_token_ids = ( + ret_item["output_ids"] if request.return_token_ids else None + ) choice_meta_info = ( ret_item["meta_info"] if request.return_meta_info else None @@ -1568,6 +1578,7 @@ class OpenAIServingChat(OpenAIServingBase): ), hidden_states=hidden_states, prompt_token_ids=choice_prompt_token_ids, + token_ids=choice_token_ids, meta_info=choice_meta_info, ) choices.append(choice_data) diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index 0b88c1a2e..8be9bf41a 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -124,6 +124,7 @@ class OpenAIServingCompletion(OpenAIServingBase): return_hidden_states=request.return_hidden_states, return_routed_experts=request.return_routed_experts, routed_experts_start_len=request.routed_experts_start_len, + return_prompt_token_ids=request.return_token_ids, rid=request.rid, session_id=request.session_id, extra_key=self._compute_extra_key(request), @@ -224,6 +225,7 @@ class OpenAIServingCompletion(OpenAIServingBase): # State tracking for streaming stream_offsets = {} n_prev_tokens = {} + n_prev_token_ids = {} # Usage tracking prompt_tokens = {} @@ -313,8 +315,26 @@ class OpenAIServingCompletion(OpenAIServingBase): ) n_prev_tokens[index] = total_output_logprobs + chunk_token_ids = None + chunk_prompt_token_ids = None + if request.return_token_ids: + output_ids = content["output_ids"] + if ( + not self.tokenizer_manager.server_args.incremental_streaming_output + ): + n_prev_token_id = n_prev_token_ids.get(index, 0) + chunk_token_ids = output_ids[n_prev_token_id:] + n_prev_token_ids[index] = len(output_ids) + else: + chunk_token_ids = output_ids + if is_first_chunk: + chunk_prompt_token_ids = content.get("prompt_token_ids") + # Generate delta - delta = text[offset:] + if self.tokenizer_manager.server_args.incremental_streaming_output: + delta = text + else: + delta = text[offset:] stream_offsets[index] = len(content["text"]) finish_reason = content["meta_info"].get("finish_reason", None) finish_reason_type = finish_reason["type"] if finish_reason else None @@ -347,6 +367,8 @@ class OpenAIServingCompletion(OpenAIServingBase): if finish_reason and "matched" in finish_reason else None ), + token_ids=chunk_token_ids, + prompt_token_ids=chunk_prompt_token_ids, ) chunk = CompletionStreamResponse( id=content["meta_info"]["id"], @@ -547,6 +569,14 @@ class OpenAIServingCompletion(OpenAIServingBase): else None ), hidden_states=hidden_states, + token_ids=( + ret_item["output_ids"] if request.return_token_ids else None + ), + prompt_token_ids=( + ret_item.get("prompt_token_ids") + if request.return_token_ids + else None + ), ) choices.append(choice_data) diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index eceeea83f..23a4237e2 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -492,6 +492,7 @@ class TestModelSerialization(unittest.TestCase): ) default_data = default_choice.model_dump() self.assertNotIn("prompt_token_ids", default_data) + self.assertNotIn("token_ids", default_data) self.assertNotIn("meta_info", default_data) choice = ChatCompletionResponseChoice( @@ -499,10 +500,12 @@ class TestModelSerialization(unittest.TestCase): message=ChatMessage(role="assistant", content="Hello"), finish_reason="stop", prompt_token_ids=[1, 2, 3], + token_ids=[4, 5], meta_info={"prompt_tokens": 3}, ) data = choice.model_dump() self.assertEqual(data["prompt_token_ids"], [1, 2, 3]) + self.assertEqual(data["token_ids"], [4, 5]) self.assertEqual(data["meta_info"], {"prompt_tokens": 3}) diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 693cdd4b8..6e2a525e6 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -150,18 +150,16 @@ class ServingChatTestCase(unittest.TestCase): self.assertEqual(adapted.session_id, "session-1") self.assertEqual(processed, self.basic_req) - def test_convert_to_internal_request_rejects_stream_return_prompt_token_ids(self): - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi?"}], - stream=True, - return_prompt_token_ids=True, - ) - - with self.assertRaisesRegex( - ValueError, "return_prompt_token_ids is not supported with streaming" - ): - self.chat._convert_to_internal_request(req, self.fastapi_request) + def test_convert_to_internal_request_rejects_stream_token_ids(self): + for field in ("return_prompt_token_ids", "return_token_ids"): + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "Hi?"}], + stream=True, + **{field: True}, + ) + with self.subTest(field=field), self.assertRaisesRegex(ValueError, field): + self.chat._convert_to_internal_request(req, self.fastapi_request) def test_convert_to_internal_request_rejects_stream_return_meta_info(self): req = ChatCompletionRequest( @@ -1656,18 +1654,20 @@ class ServingChatTestCase(unittest.TestCase): }, ) - def test_non_streaming_chat_response_returns_requested_prompt_ids_and_meta_info( + def test_non_streaming_chat_response_returns_requested_token_ids_and_meta_info( self, ): req = ChatCompletionRequest( model="x", messages=[{"role": "user", "content": "Hi?"}], return_prompt_token_ids=True, + return_token_ids=True, return_meta_info=True, ) ret = [ { "text": "Answer", + "output_ids": [21, 22], "prompt_token_ids": [11, 12, 13], "meta_info": { "id": "chatcmpl-token-ids", @@ -1684,9 +1684,11 @@ class ServingChatTestCase(unittest.TestCase): choice = response.choices[0] self.assertEqual(choice.prompt_token_ids, [11, 12, 13]) + self.assertEqual(choice.token_ids, [21, 22]) self.assertEqual(choice.meta_info, ret[0]["meta_info"]) dumped_choice = response.model_dump()["choices"][0] self.assertEqual(dumped_choice["prompt_token_ids"], [11, 12, 13]) + self.assertEqual(dumped_choice["token_ids"], [21, 22]) self.assertEqual(dumped_choice["meta_info"], ret[0]["meta_info"]) def test_streaming_cached_tokens_details_emits_sglext(self): diff --git a/test/registered/unit/entrypoints/openai/test_serving_completions.py b/test/registered/unit/entrypoints/openai/test_serving_completions.py index bf460a4ac..0a7097d8e 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_completions.py +++ b/test/registered/unit/entrypoints/openai/test_serving_completions.py @@ -167,15 +167,20 @@ class ServingCompletionTestCase(unittest.TestCase): # (but might have json_schema from the legacy json_schema field) self.assertIsNone(sampling_params.get("structural_tag")) - def test_logprobs_false_non_streaming(self): - """Test that logprobs=False doesn't cause KeyError in non-streaming response.""" + def test_non_streaming_response(self): req = CompletionRequest( - model="x", prompt="Hello", max_tokens=10, logprobs=False + model="x", + prompt="Hello", + max_tokens=10, + logprobs=False, + return_token_ids=True, ) mock_ret = [ { "text": " world", + "output_ids": [3, 4], + "prompt_token_ids": [1, 2], "meta_info": { "id": "test-id", "prompt_tokens": 1, @@ -191,6 +196,8 @@ class ServingCompletionTestCase(unittest.TestCase): self.assertEqual(len(response.choices), 1) self.assertEqual(response.choices[0].text, " world") self.assertEqual(len(response.choices[0].logprobs.top_logprobs), 0) + self.assertEqual(response.choices[0].token_ids, [3, 4]) + self.assertEqual(response.choices[0].prompt_token_ids, [1, 2]) def test_streaming_abort_yields_error(self): """Test that an abort finish reason during streaming correctly yields an error and stops.""" @@ -258,6 +265,76 @@ class ServingCompletionTestCase(unittest.TestCase): self.assertGreaterEqual(len(chunks), 2) self.assertIn("error", chunks[0]) + def test_streaming_token_ids_deltas_cover_output_exactly(self): + req = CompletionRequest( + model="x", + prompt="Hi", + max_tokens=10, + stream=True, + return_token_ids=True, + ) + adapted_request, _ = self.sc._convert_to_internal_request(req) + self.sc.tokenizer_manager.server_args.stream_response_default_include_usage = ( + False + ) + + for incremental in (False, True): + with self.subTest(incremental_streaming_output=incremental): + self.sc.tokenizer_manager.server_args.incremental_streaming_output = ( + incremental + ) + texts = ("a", "b", "c") if incremental else ("a", "ab", "abc") + output_ids = ( + ([5], [6], [7]) if incremental else ([5], [5, 6], [5, 6, 7]) + ) + chunks = [ + { + "text": text, + "output_ids": ids, + "prompt_token_ids": [1, 2], + "meta_info": { + "id": "cmpl-test", + "prompt_tokens": 2, + "completion_tokens": i + 1, + "finish_reason": {"type": "stop"} if i == 2 else None, + }, + "index": 0, + } + for i, (text, ids) in enumerate(zip(texts, output_ids)) + ] + + async def _mock_generate(*args, _chunks=chunks, **kwargs): + for chunk in _chunks: + yield chunk + + self.sc.tokenizer_manager.generate_request = _mock_generate + + async def run_stream(): + return [ + chunk + async for chunk in self.sc._generate_completion_stream( + adapted_request, req, self.fastapi_request + ) + ] + + loop = get_or_create_event_loop() + raw_chunks = loop.run_until_complete(run_stream()) + + choices = [] + for raw in raw_chunks: + if not raw.startswith("data: ") or raw.strip() == "data: [DONE]": + continue + data = json.loads(raw[len("data: ") :]) + choices.extend(data.get("choices", [])) + + token_ids = [tid for c in choices for tid in c.get("token_ids", [])] + text = "".join(c["text"] for c in choices) + self.assertEqual(text, "abc") + self.assertEqual(token_ids, [5, 6, 7]) + self.assertEqual(choices[0]["prompt_token_ids"], [1, 2]) + for choice in choices[1:]: + self.assertNotIn("prompt_token_ids", choice) + def test_non_streaming_cached_tokens_details_emits_sglext(self): """Test that non-streaming completion responses emit cached token details in sglext."""