minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jimmy Shong
2026-07-23 14:41:52 -07:00
committed by GitHub
parent ebe3ab29e4
commit 410ab4fde5
6 changed files with 159 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/openai/protocol.py
+17
View File
@@ -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
+13 -2
View File
@@ -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
+31 -1
View File
@@ -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
+3
View File
@@ -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
+15 -13
View File
@@ -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):
test/registered/unit/entrypoints/openai/test_serving_completions.py
+80 -3
View File
@@ -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."""