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]],
test/registered/unit/entrypoints/openai/test_serving_chat.py
+484
View File
@@ -127,6 +127,9 @@ class _MockTokenizerManager:
reasoning_parser=None,
stream_response_default_include_usage=False,
default_chat_template_kwargs=None,
return_input_ids=False,
return_output_ids=False,
incremental_streaming_output=False,
)
self.model_path = self.server_args.model_path
# The manager tracks the served name itself; a weight update rewrites it.
@@ -2525,6 +2528,162 @@ class ServingChatTestCase(unittest.TestCase):
self.assertEqual(len(chunks), 2)
self.assertIn("error", chunks[0])
def test_streaming_abort_with_ids_enabled(self):
"""Test that a terminal abort with input/output ids enabled yields only an error and [DONE]."""
err_msg = "Aborted by scheduler"
err_code = HTTPStatus.INTERNAL_SERVER_ERROR
async def _mock_generate_abort():
yield {
"text": "Partial ",
"prompt_token_ids": [4, 5, 6],
"output_ids": [1, 2, 3],
"meta_info": {
"id": "chatcmpl-test",
"prompt_tokens": 10,
"completion_tokens": 2,
"cached_tokens": 0,
"finish_reason": {
"type": "abort",
"status_code": err_code,
"message": err_msg,
},
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 0,
}
self.tm.generate_request.return_value = _mock_generate_abort()
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
temperature=0.7,
max_tokens=100,
stream=True,
return_input_ids_in_sglext=True,
return_output_ids_in_sglext=True,
)
with patch(
"sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
) as conv_mock:
conv_ins = Mock()
conv_ins.get_prompt.return_value = "Test prompt"
conv_mock.return_value = conv_ins
adapted_request, _ = self.chat._convert_to_internal_request(
req, self.fastapi_request
)
async def run_stream():
chunks = []
try:
async for chunk in self.chat._generate_chat_stream(
adapted_request, req, self.fastapi_request
):
chunks.append(chunk)
except Exception as e:
print(f"Error during stream iteration: {e}")
return chunks
loop = get_or_create_event_loop()
chunks = loop.run_until_complete(run_stream())
# Exactly one error chunk followed by [DONE]; no sglext ids leak.
self.assertIn("error", chunks[0])
self.assertEqual(chunks[1], "data: [DONE]\n\n")
self.assertFalse(
any("input_ids" in c or "output_ids" in c for c in chunks),
"sglext ids event leaked after abort error",
)
error_chunk_data = json.loads(chunks[0][len("data: ") :])
self.assertEqual(error_chunk_data["error"]["message"], err_msg)
self.assertEqual(error_chunk_data["error"]["code"], err_code.value)
def test_streaming_error_abort_still_finalizes_other_choices(self):
"""An error abort ends generation but still runs finalization, so a
sibling choice's finish chunk and the usage chunk follow the error."""
err_code = HTTPStatus.SERVICE_UNAVAILABLE
async def _mock_generate():
yield {
"text": "hello",
"prompt_token_ids": [4, 5, 6],
"output_ids": [1, 2],
"meta_info": {
"id": "chatcmpl-multi-abort",
"prompt_tokens": 3,
"completion_tokens": 2,
"cached_tokens": 0,
"finish_reason": {"type": "stop"},
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 0,
}
yield {
"text": "partial",
"prompt_token_ids": [4, 5, 6],
"output_ids": [7],
"meta_info": {
"id": "chatcmpl-multi-abort",
"prompt_tokens": 3,
"completion_tokens": 1,
"cached_tokens": 0,
"finish_reason": {
"type": "abort",
"status_code": err_code,
"message": "Aborted by scheduler",
},
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 1,
}
self.tm.generate_request.return_value = _mock_generate()
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
max_tokens=100,
n=2,
stream=True,
stream_options={"include_usage": True},
return_input_ids_in_sglext=True,
return_output_ids_in_sglext=True,
)
with patch(
"sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
) as conv_mock:
conv_ins = Mock()
conv_ins.get_prompt.return_value = "Test prompt"
conv_mock.return_value = conv_ins
adapted_request, _ = self.chat._convert_to_internal_request(
req, self.fastapi_request
)
chunks = self._run_chat_stream(adapted_request, req)
error_idx = next(i for i, c in enumerate(chunks) if "error" in c)
after_error = self._parse_chunks(chunks[error_idx + 1 :])
finish_reasons = [
choice["finish_reason"]
for c in after_error
for choice in c.get("choices", [])
if choice.get("finish_reason") is not None
]
self.assertEqual(finish_reasons, ["stop"])
self.assertTrue(
any(c.get("usage") is not None for c in after_error),
"usage chunk dropped after error abort",
)
def _run_chat_stream(self, adapted_request, req):
async def run_stream():
chunks = []
@@ -3008,6 +3167,331 @@ class ServingChatTestCase(unittest.TestCase):
},
)
def _output_ids_ret(self, *output_ids_by_choice):
return [
{
"text": "Answer",
"prompt_token_ids": [1, 2, 3],
"output_ids": list(output_ids),
"meta_info": {
"id": "chatcmpl-output-ids",
"prompt_tokens": 3,
"completion_tokens": len(output_ids),
"cached_tokens": 0,
"finish_reason": {"type": "stop", "matched": None},
"weight_version": "default",
},
}
for output_ids in output_ids_by_choice
]
def test_non_streaming_ids_emit_sglext(self):
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
n=2,
return_input_ids_in_sglext=True,
return_output_ids_in_sglext=True,
)
response = self.chat._build_chat_response(
req, self._output_ids_ret([5, 6, 7], [8, 9]), 123
)
self.assertIsNotNone(response.sglext)
self.assertEqual(response.sglext.input_ids, [1, 2, 3])
self.assertEqual(response.sglext.output_ids, [[5, 6, 7], [8, 9]])
dumped = json.loads(response.model_dump_json())
self.assertEqual(dumped["sglext"]["input_ids"], [1, 2, 3])
self.assertEqual(dumped["sglext"]["output_ids"], [[5, 6, 7], [8, 9]])
# None sglext fields stay stripped from the serialized response.
self.assertNotIn("routed_experts", dumped["sglext"])
def test_non_streaming_ids_not_returned_by_default(self):
req = ChatCompletionRequest(
model="x", messages=[{"role": "user", "content": "Hi?"}]
)
response = self.chat._build_chat_response(
req, self._output_ids_ret([5, 6, 7]), 123
)
self.assertIsNone(response.sglext)
def test_non_streaming_ids_server_default_enables_flag(self):
self.tm.server_args.return_input_ids = True
self.tm.server_args.return_output_ids = True
req = ChatCompletionRequest(
model="x", messages=[{"role": "user", "content": "Hi?"}]
)
response = self.chat._build_chat_response(
req, self._output_ids_ret([5, 6, 7]), 123
)
self.assertIsNotNone(response.sglext)
self.assertEqual(response.sglext.input_ids, [1, 2, 3])
self.assertEqual(response.sglext.output_ids, [[5, 6, 7]])
def test_ids_headers_enable_flags(self):
self.fastapi_request.headers = {
"x-sglext-return-input-ids": "1",
"x-sglext-return-output-ids": "1",
}
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
)
processed_messages = MessageProcessingResult(
"Test prompt",
[1, 2, 3],
None,
None,
[],
["</s>"],
None,
)
with patch.object(
self.chat, "_process_messages", return_value=processed_messages
):
_, processed_request = self.chat._convert_to_internal_request(
req, self.fastapi_request
)
self.assertTrue(processed_request.return_input_ids_in_sglext)
self.assertTrue(processed_request.return_output_ids_in_sglext)
def _run_output_ids_stream(
self,
chunk_output_ids,
incremental,
framed=False,
return_raw=False,
cached_tokens_details=None,
):
"""Stream chunks with incremental or cumulative output_ids;
return parsed sglext chunks, or raw SSE strings when return_raw."""
self.tm.server_args.incremental_streaming_output = incremental
if framed:
self.fastapi_request.headers["x-sglext-ids-framed"] = "1"
async def _mock_generate():
for i, ids in enumerate(chunk_output_ids):
finished = i == len(chunk_output_ids) - 1
yield {
"text": "chunk",
"prompt_token_ids": [1, 2, 3],
"output_ids": list(ids),
"meta_info": {
"id": "chatcmpl-output-ids-stream",
"prompt_tokens": 3,
"completion_tokens": 1 + i,
"cached_tokens": 0,
"cached_tokens_details": cached_tokens_details,
"finish_reason": (
{"type": "stop", "matched": None} if finished else None
),
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 0,
}
self.tm.generate_request.return_value = _mock_generate()
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
max_tokens=100,
stream=True,
return_input_ids_in_sglext=True,
return_output_ids_in_sglext=True,
return_cached_tokens_details=cached_tokens_details is not None,
)
with patch(
"sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
) as conv_mock:
conv_ins = Mock()
conv_ins.get_prompt.return_value = "Test prompt"
conv_mock.return_value = conv_ins
adapted_request, _ = self.chat._convert_to_internal_request(
req, self.fastapi_request
)
chunks = self._run_chat_stream(adapted_request, req)
if return_raw:
return chunks
return [c for c in self._parse_chunks(chunks) if "sglext" in c]
def test_streaming_output_ids_incremental_accumulates_deltas(self):
sglext_chunks = self._run_output_ids_stream([[5, 6], [7]], incremental=True)
self.assertEqual(len(sglext_chunks), 1)
self.assertEqual(sglext_chunks[0]["choices"], [])
# input_ids is captured once as a flat shared prompt, not accumulated.
self.assertEqual(sglext_chunks[0]["sglext"]["input_ids"], [1, 2, 3])
self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
def test_streaming_output_ids_non_incremental_keeps_latest_full_list(self):
sglext_chunks = self._run_output_ids_stream(
[[5, 6], [5, 6, 7]], incremental=False
)
self.assertEqual(len(sglext_chunks), 1)
self.assertEqual(sglext_chunks[0]["sglext"]["input_ids"], [1, 2, 3])
self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
_SGLEXT_IDS_EVENT_PREFIX = "event: sglext_ids\ndata: "
def test_streaming_framed_ids_emits_named_event(self):
raw = self._run_output_ids_stream(
[[5, 6], [5, 6, 7]], incremental=False, framed=True, return_raw=True
)
named = [c for c in raw if c.startswith(self._SGLEXT_IDS_EVENT_PREFIX)]
self.assertEqual(len(named), 1)
payload = json.loads(named[0][len(self._SGLEXT_IDS_EVENT_PREFIX) :])
self.assertEqual(payload["choices"], [])
# The named event carries ONLY the id fields.
self.assertEqual(set(payload["sglext"]), {"input_ids", "output_ids"})
self.assertEqual(payload["sglext"]["input_ids"], [1, 2, 3])
self.assertEqual(payload["sglext"]["output_ids"], [[5, 6, 7]])
# All requested sglext fields were ids, so no plain sglext chunk.
self.assertEqual([c for c in self._parse_chunks(raw) if "sglext" in c], [])
def test_streaming_framed_splits_non_id_fields_into_plain_chunk(self):
raw = self._run_output_ids_stream(
[[5, 6], [5, 6, 7]],
incremental=False,
framed=True,
return_raw=True,
cached_tokens_details={"device": 2, "host": 1},
)
named = [c for c in raw if c.startswith(self._SGLEXT_IDS_EVENT_PREFIX)]
self.assertEqual(len(named), 1)
named_payload = json.loads(named[0][len(self._SGLEXT_IDS_EVENT_PREFIX) :])
self.assertEqual(set(named_payload["sglext"]), {"input_ids", "output_ids"})
plain_sglext = [c for c in self._parse_chunks(raw) if "sglext" in c]
self.assertEqual(len(plain_sglext), 1)
self.assertEqual(
plain_sglext[0]["sglext"]["cached_tokens_details"],
{"device": 2, "host": 1},
)
self.assertNotIn("input_ids", plain_sglext[0]["sglext"])
self.assertNotIn("output_ids", plain_sglext[0]["sglext"])
# The plain non-id chunk precedes the named ids event.
plain_raw_idx = next(
i
for i, c in enumerate(raw)
if c.startswith("data: ") and "cached_tokens_details" in c
)
self.assertLess(plain_raw_idx, raw.index(named[0]))
def _run_output_ids_stream_with_graceful_abort(
self, normal_chunks, abort_output_ids, incremental, abort_completion_tokens
):
"""Stream normal chunks then a graceful-abort chunk; return parsed sglext chunks."""
self.tm.server_args.incremental_streaming_output = incremental
async def _mock_generate():
generated = 0
for ids in normal_chunks:
generated = generated + len(ids) if incremental else len(ids)
yield {
"text": "chunk",
"output_ids": list(ids),
"meta_info": {
"id": "chatcmpl-abort-ids-stream",
"prompt_tokens": 3,
"completion_tokens": generated,
"cached_tokens": 0,
"finish_reason": None,
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 0,
}
# Graceful abort terminal chunk (no status_code): falls through to
# the normal finalization path.
yield {
"text": "chunk",
"output_ids": list(abort_output_ids),
"meta_info": {
"id": "chatcmpl-abort-ids-stream",
"prompt_tokens": 3,
"completion_tokens": abort_completion_tokens,
"cached_tokens": 0,
"finish_reason": {"type": "abort", "message": "Aborted."},
"output_token_logprobs": None,
"output_top_logprobs": None,
},
"index": 0,
}
self.tm.generate_request.return_value = _mock_generate()
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
max_tokens=100,
stream=True,
return_output_ids_in_sglext=True,
)
with patch(
"sglang.srt.entrypoints.openai.serving_chat.generate_chat_conv"
) as conv_mock:
conv_ins = Mock()
conv_ins.get_prompt.return_value = "Test prompt"
conv_mock.return_value = conv_ins
adapted_request, _ = self.chat._convert_to_internal_request(
req, self.fastapi_request
)
chunks = self._run_chat_stream(adapted_request, req)
return [c for c in self._parse_chunks(chunks) if "sglext" in c]
def test_streaming_output_ids_incremental_ignores_abort_chunk(self):
# Abort re-sends the last token; it must not be appended again.
sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
normal_chunks=[[5, 6], [7]],
abort_output_ids=[7],
incremental=True,
abort_completion_tokens=3,
)
self.assertEqual(len(sglext_chunks), 1)
self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
def test_streaming_output_ids_incremental_keeps_coalesced_abort_deltas(self):
# Coalesced abort chunk: keep real deltas, drop the repeated last token.
sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
normal_chunks=[[5, 6]],
abort_output_ids=[7, 8, 8],
incremental=True,
abort_completion_tokens=4,
)
self.assertEqual(len(sglext_chunks), 1)
self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7, 8]])
def test_streaming_output_ids_non_incremental_abort_supersedes_earlier(self):
sglext_chunks = self._run_output_ids_stream_with_graceful_abort(
normal_chunks=[[5, 6]],
abort_output_ids=[5, 6, 7],
incremental=False,
abort_completion_tokens=3,
)
self.assertEqual(len(sglext_chunks), 1)
self.assertEqual(sglext_chunks[0]["sglext"]["output_ids"], [[5, 6, 7]])
def test_streaming_parallel_sampling_orders_spec_details_by_choice(self):
async def mock_generate():
for index in (1, 0):
test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+45
View File
@@ -270,6 +270,51 @@ class TestRidToStateCleanupOnAbort(CustomTestCase):
)
class TestAbortOutputPayload(CustomTestCase):
"""An abort chunk is often the only thing a client sees;
it must carry the same optional fields as a normal finish chunk."""
def test_abort_includes_prompt_token_ids_only_when_requested(self):
"""The abort chunk carries prompt_token_ids captured at tokenization,
and omits the field when the request did not ask for them."""
tm = _make_tokenizer_manager(self)
with_ids = _make_req_state("abort_prompt_ids_rid")
with_ids.prompt_token_ids = [1, 2, 3]
without_ids = _make_req_state("abort_no_prompt_ids_rid")
for state in (with_ids, without_ids):
tm.rid_to_state[state.obj.rid] = state
tm._handle_abort_req(_make_abort_req(state.obj.rid))
self.assertEqual(with_ids.out_list[0]["prompt_token_ids"], [1, 2, 3])
self.assertNotIn("prompt_token_ids", without_ids.out_list[0])
def test_abort_output_ids_match_the_streaming_mode(self):
"""Only incremental streaming collapses the abort chunk to the last
token; cumulative chunks supersede, so they carry the whole generation.
"""
cases = [
("incremental stream", True, True, [7]),
("cumulative stream", True, False, [5, 6, 7]),
("non-stream", False, False, [5, 6, 7]),
]
for name, is_stream, incremental, expected in cases:
with self.subTest(name):
tm = _make_tokenizer_manager(self)
tm.incremental_streaming_output = incremental
rid = f"abort_output_ids_{name}"
state = _make_req_state(rid)
state.obj.stream = is_stream
state.output_ids = [5, 6, 7]
tm.rid_to_state[rid] = state
tm._handle_abort_req(_make_abort_req(rid))
out = state.out_list[0]
self.assertEqual(out["output_ids"], expected)
self.assertEqual(out["meta_info"]["completion_tokens"], 3)
class TestRidToStateCleanupOnBatchOutput(CustomTestCase):
"""Test that _handle_batch_output removes rid from rid_to_state on completion."""