From 8a98f11078b0fc7beff6f7ce2b1ea59682d83171 Mon Sep 17 00:00:00 2001 From: Amy Chang <78667707+amykchang@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:45:13 -0500 Subject: [PATCH] [feature] Add response-level input/output token ids to chat completions via SglExt (#34488) --- .../sglang/srt/entrypoints/openai/protocol.py | 18 + .../srt/entrypoints/openai/serving_chat.py | 141 ++++- .../sglang/srt/managers/tokenizer_manager.py | 4 +- python/sglang/srt/server_args.py | 10 + .../entrypoints/openai/test_serving_chat.py | 484 ++++++++++++++++++ .../test_tokenizer_manager_rid_cleanup.py | 45 ++ 6 files changed, 683 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 346c7c02e..3b99baea6 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -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, diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 8aaa07e34..43eb943ce 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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): diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index cc7d03d4b..e9d98ecb6 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0e9597b6d..40524e411 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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]], diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 969b44e1e..204c34ad2 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -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, + [], + [""], + 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): diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py index 4f54c6e6a..1fab0d801 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -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."""