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

[feature] Add response-level input/output token ids to chat completions via SglExt (#34488)

Browse Source
This commit is contained in:
Amy Chang
2026-09-04 16:45:13 -07:00
committed by GitHub
parent 9c254df8cd
commit 8a98f11078
6 changed files with 683 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/openai/protocol.py
+18
View File
@@ -441,6 +441,22 @@ class SglExt(BaseModel):
spec_tokens_details: Optional[Union[SpecTokensDetails, List[SpecTokensDetails]]] = (
None
)
input_ids: Optional[List[int]] = None
output_ids: Optional[List[List[int]]] = None
def split_ids(self) -> Tuple[Optional[SglExt], Optional[SglExt]]:
"""Split set fields into (non_ids, ids); a side with no set fields is None."""
non_ids: Dict[str, Any] = {}
ids: Dict[str, Any] = {}
for name in type(self).model_fields:
value = getattr(self, name)
if value is None:
continue
(ids if name in ("input_ids", "output_ids") else non_ids)[name] = value
return (
type(self)(**non_ids) if non_ids else None,
type(self)(**ids) if ids else None,
)
@model_serializer(mode="wrap")
def _serialize(self, handler):
@@ -869,6 +885,8 @@ class ChatCompletionRequest(BaseModel):
return_prompt_token_ids: bool = False
return_token_ids: bool = False
return_meta_info: bool = False
return_input_ids_in_sglext: bool = False
return_output_ids_in_sglext: bool = False
return_sampling_mask: bool = False
reasoning_effort: ReasoningEffortType = Field(
default=None,
python/sglang/srt/entrypoints/openai/serving_chat.py
+123 -18
View File
@@ -706,6 +706,20 @@ class OpenAIServingChat(OpenAIServingBase):
"""Post-process reasoning and tool_calls before building response."""
return reasoning_text, tool_calls
def _should_return_input_ids(self, request: ChatCompletionRequest) -> bool:
"""Whether prompt (input) token ids should be returned via sglext."""
return (
request.return_input_ids_in_sglext
or self.tokenizer_manager.server_args.return_input_ids
)
def _should_return_output_ids(self, request: ChatCompletionRequest) -> bool:
"""Whether sampled output token ids should be returned via sglext."""
return (
request.return_output_ids_in_sglext
or self.tokenizer_manager.server_args.return_output_ids
)
def _continuous_usage_cached_details(
self, content: Dict[str, Any]
) -> Optional[PromptTokensDetails]:
@@ -1034,6 +1048,15 @@ class OpenAIServingChat(OpenAIServingBase):
request: ChatCompletionRequest,
raw_request: Request = None,
) -> tuple[GenerateReqInput, ChatCompletionRequest]:
# Header-based opt-in (same rationale as request_headers.py).
if raw_request is not None and not request.return_input_ids_in_sglext:
if raw_request.headers.get("x-sglext-return-input-ids") == "1":
request.return_input_ids_in_sglext = True
if raw_request is not None and not request.return_output_ids_in_sglext:
if raw_request.headers.get("x-sglext-return-output-ids") == "1":
request.return_output_ids_in_sglext = True
reasoning_effort = (
request.chat_template_kwargs.pop("reasoning_effort", None)
if request.chat_template_kwargs
@@ -1155,8 +1178,11 @@ 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
or request.return_token_ids,
return_prompt_token_ids=(
request.return_prompt_token_ids
or request.return_token_ids
or self._should_return_input_ids(request)
),
)
if (
raw_request is not None
@@ -1666,14 +1692,25 @@ class OpenAIServingChat(OpenAIServingBase):
image_tokens = {}
audio_tokens = {}
video_tokens = {}
input_ids: Optional[List[int]] = None
output_ids: Dict[int, List[int]] = {}
stream_started = False
error_aborted = False
try:
include_usage, continuous_usage_stats = should_include_usage(
request.stream_options,
self.tokenizer_manager.server_args.stream_response_default_include_usage,
)
return_input_ids = self._should_return_input_ids(request)
return_output_ids = self._should_return_output_ids(request)
ids_framed = (
raw_request is not None
and raw_request.headers.get("x-sglext-ids-framed") == "1"
)
async for content in self.tokenizer_manager.generate_request(
adapted_request, raw_request
):
@@ -1702,6 +1739,32 @@ class OpenAIServingChat(OpenAIServingBase):
audio_tokens[index] = content["meta_info"].get("audio_tokens", 0)
video_tokens[index] = content["meta_info"].get("video_tokens", 0)
finish_reason = content["meta_info"].get("finish_reason", None)
finish_reason_type = finish_reason["type"] if finish_reason else None
if return_input_ids and input_ids is None:
# The prompt is the full, shared prompt (same across choices
# and constant across chunks), so capture it once.
chunk_input_ids = content.get("prompt_token_ids")
if chunk_input_ids is not None:
input_ids = list(chunk_input_ids)
if return_output_ids:
chunk_output_ids = content.get("output_ids")
if chunk_output_ids is not None:
if self.tokenizer_manager.server_args.incremental_streaming_output:
accumulated = output_ids.setdefault(index, [])
if finish_reason_type == "abort":
# The abort chunk re-sends the last token plus any coalesced deltas;
# keep only what brings the total up to completion_tokens.
keep = completion_tokens[index] - len(accumulated)
chunk_output_ids = chunk_output_ids[: max(keep, 0)]
accumulated.extend(chunk_output_ids)
else:
# Intermediate chunks share the live state.output_ids
# list; the final chunk is a stable copy.
output_ids[index] = chunk_output_ids
# Handle logprobs
choice_logprobs = None
if request.logprobs:
@@ -1715,9 +1778,6 @@ class OpenAIServingChat(OpenAIServingBase):
).model_dump()
n_prev_tokens[index] = total_output_logprobs
finish_reason = content["meta_info"].get("finish_reason", None)
finish_reason_type = finish_reason["type"] if finish_reason else None
# Track finish_reason for each index
if finish_reason_type:
# Abort with an explicit error status_code is a system error
@@ -1736,6 +1796,7 @@ class OpenAIServingChat(OpenAIServingBase):
code.value,
)
yield f"data: {error}\n\n"
error_aborted = True
break
finish_reasons[index] = finish_reason
@@ -1840,24 +1901,54 @@ class OpenAIServingChat(OpenAIServingBase):
spec_details if request.n > 1 else spec_details[0]
)
if any(
obj is not None
for obj in [
sglext_routed,
sglext_cached_tokens_details,
sglext_spec_tokens_details,
# Omit token ids after an error abort.
sglext_input_ids = None
if return_input_ids and input_ids and not error_aborted:
sglext_input_ids = list(input_ids)
sglext_output_ids = None
if return_output_ids and output_ids and not error_aborted:
sglext_output_ids = [
list(output_ids.get(i, [])) for i in range(request.n)
]
):
sglext_full = SglExt(
routed_experts=sglext_routed,
cached_tokens_details=sglext_cached_tokens_details,
spec_tokens_details=sglext_spec_tokens_details,
input_ids=sglext_input_ids,
output_ids=sglext_output_ids,
)
sglext_non_ids, sglext_ids = sglext_full.split_ids()
if ids_framed:
# A named SSE event lets transit hops pick the ids out without parsing JSON;
# the other sglext fields keep the plain data-chunk shape.
if sglext_non_ids is not None:
sglext_chunk = ChatCompletionStreamResponse(
id=content["meta_info"]["id"],
created=int(time.time()),
choices=[],
model=request.model,
sglext=sglext_non_ids,
)
yield f"data: {sglext_chunk.model_dump_json()}\n\n"
if sglext_ids is not None:
sglext_ids_chunk = ChatCompletionStreamResponse(
id=content["meta_info"]["id"],
created=int(time.time()),
choices=[],
model=request.model,
sglext=sglext_ids,
)
yield f"event: sglext_ids\ndata: {sglext_ids_chunk.model_dump_json()}\n\n"
elif sglext_non_ids is not None or sglext_ids is not None:
sglext_chunk = ChatCompletionStreamResponse(
id=content["meta_info"]["id"],
created=int(time.time()),
choices=[], # sglext is at response level
model=request.model,
sglext=SglExt(
routed_experts=sglext_routed,
cached_tokens_details=sglext_cached_tokens_details,
spec_tokens_details=sglext_spec_tokens_details,
),
sglext=sglext_full,
)
yield f"data: {sglext_chunk.model_dump_json()}\n\n"
@@ -1972,12 +2063,26 @@ class OpenAIServingChat(OpenAIServingBase):
if request.n > 1
else (spec_details[0] if spec_details else None)
)
input_ids = None
if self._should_return_input_ids(request) and "prompt_token_ids" in ret[0]:
input_ids = list(ret[0]["prompt_token_ids"])
output_ids = None
if self._should_return_output_ids(request):
output_ids = [list(ret_item["output_ids"]) for ret_item in ret]
response_sglext = None
if routed_experts or cached_tokens_details or spec_tokens_details:
if (
routed_experts
or cached_tokens_details
or spec_tokens_details
or input_ids is not None
or output_ids is not None
):
response_sglext = SglExt(
routed_experts=routed_experts,
cached_tokens_details=cached_tokens_details,
spec_tokens_details=spec_tokens_details,
input_ids=input_ids,
output_ids=output_ids,
)
for idx, ret_item in enumerate(ret):
python/sglang/srt/managers/tokenizer_manager.py
+3 -1
View File
@@ -3277,13 +3277,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
output_ids = state.output_ids
meta_info["completion_tokens"] = len(output_ids)
if is_stream:
if is_stream and self.incremental_streaming_output:
output_ids = [output_ids[-1]] if len(output_ids) > 0 else []
out = {
"text": state.get_text(),
"output_ids": output_ids,
"meta_info": meta_info,
}
if state.prompt_token_ids is not None:
out["prompt_token_ids"] = state.prompt_token_ids
del self.rid_to_state[recv_obj.rid]
state.out_list.append(out)
python/sglang/srt/server_args.py
+10
View File
@@ -1352,6 +1352,16 @@ class ServerArgs:
"Return number of cached tokens in usage.prompt_tokens_details for each openai request.",
NS("serving"),
] = False
return_input_ids: A[
bool,
"Return prompt (input) token ids on the response-level sglext extension for every chat completion request, as if return_input_ids_in_sglext were set on the request.",
NS("serving"),
] = False
return_output_ids: A[
bool,
"Return sampled output token ids on the response-level sglext extension for every chat completion request, as if return_output_ids_in_sglext were set on the request.",
NS("serving"),
] = False
reasoning_parser: A[Optional[str], NS("serving")] = None
default_chat_template_kwargs: A[
Optional[Dict[str, Any]],