minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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)