Add return_token_ids support to completions and chat completions APIs (#30917)

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

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)

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)
 

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})
 
 

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):
+    def test_convert_to_internal_request_rejects_stream_token_ids(self):
-        req = ChatCompletionRequest(
+        for field in ("return_prompt_token_ids", "return_token_ids"):
-            model="x",
+            req = ChatCompletionRequest(
-            messages=[{"role": "user", "content": "Hi?"}],
+                model="x",
-            stream=True,
+                messages=[{"role": "user", "content": "Hi?"}],
-            return_prompt_token_ids=True,
+                stream=True,
-        )
+                **{field: True},
-
+            )
-        with self.assertRaisesRegex(
+            with self.subTest(field=field), self.assertRaisesRegex(ValueError, field):
-            ValueError, "return_prompt_token_ids is not supported with streaming"
+                self.chat._convert_to_internal_request(req, self.fastapi_request)
-        ):
-            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):

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):
+    def test_non_streaming_response(self):
-        """Test that logprobs=False doesn't cause KeyError in non-streaming response."""
         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."""