From 3b11ffec649f681d204bf5f6528d017b6162975c Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 9 Mar 2026 23:02:07 +0000 Subject: [PATCH 1/9] refine impl --- .../sglang/srt/entrypoints/openai/protocol.py | 10 +- .../srt/entrypoints/openai/serving_chat.py | 150 +++++++++++++++++- python/sglang/srt/entrypoints/openai/utils.py | 3 +- .../basic/test_return_token_ids.py | 105 ++++++++---- 4 files changed, 228 insertions(+), 40 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index c07e8370245f..5140fafa49b5 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -81,7 +81,6 @@ class LogProbs(BaseModel): text_offset: List[int] = Field(default_factory=list) token_logprobs: List[Optional[float]] = Field(default_factory=list) tokens: List[str] = Field(default_factory=list) - token_ids: List[int] = Field(default_factory=list) top_logprobs: List[Optional[Dict[str, float]]] = Field(default_factory=list) @@ -93,7 +92,6 @@ class TopLogprob(BaseModel): class ChatCompletionTokenLogprob(BaseModel): token: str - token_id: int bytes: List[int] logprob: float top_logprobs: List[TopLogprob] @@ -568,6 +566,7 @@ class ChatCompletionRequest(BaseModel): return_routed_experts: bool = False return_cached_tokens_details: bool = False return_prompt_token_ids: bool = False + return_meta_info: bool = False reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field( default="medium", description="Constrains effort on reasoning for reasoning models. " @@ -603,6 +602,10 @@ class ChatCompletionRequest(BaseModel): custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None custom_params: Optional[Dict] = None + # Pretokenized input: skip tokenization for first N messages + pretokenized_token_ids: Optional[List[int]] = None + pretokenized_num_message: Optional[int] = None + # For request id rid: Optional[Union[List[str], str]] = None # Extra key for classifying the request (e.g. cache_salt) @@ -806,6 +809,7 @@ class ChatCompletionResponseChoice(BaseModel): matched_stop: Union[None, int, str] = None hidden_states: Optional[object] = None prompt_token_ids: Optional[List[int]] = None + meta_info: Optional[Dict[str, Any]] = None @model_serializer(mode="wrap") def _serialize(self, handler): @@ -814,6 +818,8 @@ def _serialize(self, handler): data.pop("hidden_states", None) if self.prompt_token_ids is None: data.pop("prompt_token_ids", None) + if self.meta_info is None: + data.pop("meta_info", None) return data diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index bdbb191b29d1..76cb436ca659 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -3,6 +3,7 @@ import copy import json import logging +from sglang.srt.entrypoints.openai.protocol import ChatCompletionMessageParam import time import uuid from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union @@ -372,6 +373,14 @@ def _apply_jinja_template( is_multimodal: bool, ) -> MessageProcessingResult: """Apply Jinja chat template""" + # Short-circuit: use pretokenized path if provided + if ( + request.pretokenized_token_ids is not None + and request.pretokenized_num_message is not None + and request.pretokenized_num_message > 0 + ): + return self._apply_pretokenized_template(request, tools) + prompt = "" prompt_ids = [] openai_compatible_messages = [] @@ -594,6 +603,137 @@ def _apply_conversation_template( stop=stop, ) + def _apply_pretokenized_template( + self, + request: ChatCompletionRequest, + tools: Optional[List[Dict]], + ) -> MessageProcessingResult: + """Apply pretokenized template with incremental tokenization for tool messages. + + Supported input pattern (non-streaming, non-multimodal only): + messages[0..N-1] are pretokenized (token IDs provided via pretokenized_token_ids) + messages[N..end] are new tool responses to be incrementally tokenized + + Constraints: + - messages[N-1] must be role="assistant" (the last pretokenized turn) + - messages[N..end] must all be role="tool" or "system" + - No multimodal content allowed in any message + + Example valid pattern: + messages = [system, user, assistant(tool_calls), tool, tool, system, ...] + pretokenized_num_message = 3 (system + user + assistant are pretokenized) + pretokenized_token_ids = [...] (token IDs for the first 3 messages) + + How it works: + 1. Tokenize first N messages with apply_chat_template(add_generation_prompt=False) + to get prefix_ids + 2. Tokenize ALL messages with apply_chat_template(add_generation_prompt=True) + to get full_ids + 3. Verify prefix_ids matches full_ids[:len(prefix_ids)] + 4. incremental_ids = full_ids[len(prefix_ids):] + 5. Final prompt_ids = pretokenized_token_ids + incremental_ids + """ + N = request.pretokenized_num_message + messages = request.messages + + # FIXME + if self.tokenizer_manager.model_config.is_multimodal: + raise ValueError( + "Pretokenized token IDs input is not supported for multimodal models" + ) + + if N > len(messages): + raise ValueError( + f"pretokenized_num_message ({N}) > total messages ({len(messages)})" + ) + + if messages[N - 1].role != "assistant": + raise ValueError( + f"Message at index {N - 1} must be assistant, got {messages[N - 1].role}" + ) + + ALLOWED_APPEND_ROLES = {"tool", "system"} + for i in range(N, len(messages)): + if messages[i].role not in ALLOWED_APPEND_ROLES: + raise ValueError( + f"Message at index {i} must be one of {ALLOWED_APPEND_ROLES}, got {messages[i].role}" + ) + + all_msg_dicts = [msg.model_dump() for msg in messages] + prefix_msgs = all_msg_dicts[:N] + + # Process tool_calls arguments: str -> dict (consistent with standard path) + for msg in all_msg_dicts: + if ( + msg["role"] == "assistant" + and "tool_calls" in msg + and isinstance(msg["tool_calls"], list) + ): + for item in msg["tool_calls"]: + if "arguments" in item["function"] and isinstance( + item["function"]["arguments"], str + ): + item["function"]["arguments"] = orjson.loads( + item["function"]["arguments"] + ) + + chat_template_kwargs = request.chat_template_kwargs or {} + + # Tokenize first N messages (no generation prompt) + prefix_ids = self.tokenizer_manager.tokenizer.apply_chat_template( + prefix_msgs, + tokenize=True, + add_generation_prompt=False, + tools=tools, + reasoning_effort=request.reasoning_effort, + return_dict=False, + **chat_template_kwargs, + ) + + # Tokenize all messages (with generation prompt) + full_ids = self.tokenizer_manager.tokenizer.apply_chat_template( + all_msg_dicts, + tokenize=True, + add_generation_prompt=True, + tools=tools, + reasoning_effort=request.reasoning_effort, + return_dict=False, + **chat_template_kwargs, + ) + + # Validate prefix match + if full_ids[: len(prefix_ids)] != prefix_ids: + raise ValueError( + "Prefix mismatch: tokenizing first N messages does not match " + "the prefix of tokenizing all messages. " + f"prefix_ids length={len(prefix_ids)}, full_ids length={len(full_ids)}" + ) + + incremental_ids = full_ids[len(prefix_ids) :] + # Fault tolerance: templates like qwen3 insert a trailing `\n` after + # the last assistant message. Its loss-mask is 0, so we only prepend + # it for alignment when it really is a whitespace-only token. + if ( + prefix_ids[-1] != request.pretokenized_token_ids[-1] + and self.tokenizer_manager.tokenizer.decode([prefix_ids[-1]]).strip() == "" + ): + incremental_ids = [prefix_ids[-1]] + incremental_ids + prompt_ids = list(request.pretokenized_token_ids) + incremental_ids + + # Decode prompt_ids to text for the prompt field + prompt = self.tokenizer_manager.tokenizer.decode(prompt_ids) + + stop = request.stop + return MessageProcessingResult( + prompt=prompt, + prompt_ids=prompt_ids, + image_data=None, + audio_data=None, + video_data=None, + modalities=[], + stop=stop, + ) + async def _handle_streaming_request( self, adapted_request: GenerateReqInput, @@ -975,6 +1115,10 @@ def _build_chat_response( else None ) + choice_meta_info = ( + ret_item["meta_info"] if request.return_meta_info else None + ) + choice_data = ChatCompletionResponseChoice( index=idx, message=ChatMessage( @@ -992,6 +1136,7 @@ def _build_chat_response( ), hidden_states=hidden_states, prompt_token_ids=choice_prompt_token_ids, + meta_info=choice_meta_info, ) choices.append(choice_data) @@ -1023,8 +1168,8 @@ def _process_logprobs_tokens( """ token_logprobs = [] - for token_idx, (token, token_id, logprob) in enumerate( - zip(logprobs.tokens, logprobs.token_ids, logprobs.token_logprobs) + for token_idx, (token, logprob) in enumerate( + zip(logprobs.tokens, logprobs.token_logprobs) ): token_bytes = list(token.encode("utf-8")) top_logprobs = [] @@ -1046,7 +1191,6 @@ def _process_logprobs_tokens( token_logprobs.append( ChatCompletionTokenLogprob( token=token, - token_id=token_id, bytes=token_bytes, logprob=logprob, top_logprobs=top_logprobs, diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py index 68c30983ccc2..796f8f59b357 100644 --- a/python/sglang/srt/entrypoints/openai/utils.py +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -20,9 +20,8 @@ def to_openai_style_logprobs( ret_logprobs = LogProbs() def append_token_logprobs(token_logprobs): - for logprob, token_id, token_text in token_logprobs: + for logprob, _, token_text in token_logprobs: ret_logprobs.tokens.append(token_text) - ret_logprobs.token_ids.append(token_id) ret_logprobs.token_logprobs.append(logprob) # Not supported yet diff --git a/test/registered/openai_server/basic/test_return_token_ids.py b/test/registered/openai_server/basic/test_return_token_ids.py index 05c31bba31a7..ca9474c4b996 100644 --- a/test/registered/openai_server/basic/test_return_token_ids.py +++ b/test/registered/openai_server/basic/test_return_token_ids.py @@ -1,17 +1,17 @@ """ -Unit tests for the return_prompt_token_ids feature in ChatCompletion endpoint. +Unit tests for token ID return features in ChatCompletion endpoint. Tests that: 1. Protocol models correctly handle return_prompt_token_ids / prompt_token_ids fields -2. ChatCompletionTokenLogprob includes token_id field +2. response_token_ids is populated when logprobs=True (non-streaming) 3. Request conversion passes return_prompt_token_ids flag through -4. Non-streaming response includes prompt_token_ids -5. Fields are omitted from JSON when return_prompt_token_ids is False (default) +4. Non-streaming response includes prompt_token_ids and response_token_ids +5. Fields are omitted from JSON when not applicable Run with: - python -m pytest test/registered/openai_server/basic/test_return_prompt_token_ids.py -v + python -m pytest test/registered/openai_server/basic/test_return_token_ids.py -v or: - python test/registered/openai_server/basic/test_return_prompt_token_ids.py -v + python test/registered/openai_server/basic/test_return_token_ids.py -v """ import json @@ -60,10 +60,7 @@ ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, - ChatCompletionTokenLogprob, ChatMessage, - ChoiceLogprobs, - TopLogprob, UsageInfo, ) @@ -147,36 +144,27 @@ def test_choice_includes_prompt_token_ids_when_set(self): self.assertIn("prompt_token_ids", data) self.assertEqual(data["prompt_token_ids"], [1, 2, 3]) - # --- ChatCompletionTokenLogprob --- + # --- response_token_ids --- - def test_token_logprob_includes_token_id(self): - logprob = ChatCompletionTokenLogprob( - token="hello", - token_id=12345, - bytes=list(b"hello"), - logprob=-0.5, - top_logprobs=[], - ) - data = logprob.model_dump() - self.assertEqual(data["token_id"], 12345) - - def test_token_logprob_in_choice_logprobs(self): - """token_id should appear in serialized logprobs.content entries.""" - logprob_entry = ChatCompletionTokenLogprob( - token="hi", - token_id=100, - bytes=list(b"hi"), - logprob=-1.0, - top_logprobs=[], + def test_choice_omits_response_token_ids_when_none(self): + choice = ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hi"), + finish_reason="stop", ) + data = choice.model_dump() + self.assertNotIn("response_token_ids", data) + + def test_choice_includes_response_token_ids_when_set(self): choice = ChatCompletionResponseChoice( index=0, message=ChatMessage(role="assistant", content="hi"), - logprobs=ChoiceLogprobs(content=[logprob_entry]), finish_reason="stop", + response_token_ids=MOCK_OUTPUT_TOKEN_IDS, ) data = choice.model_dump() - self.assertEqual(data["logprobs"]["content"][0]["token_id"], 100) + self.assertIn("response_token_ids", data) + self.assertEqual(data["response_token_ids"], MOCK_OUTPUT_TOKEN_IDS) # --- Full JSON round-trip --- @@ -364,7 +352,18 @@ def setUp(self): self.chat = OpenAIServingChat(tm, template_mgr) - def _make_ret(self, include_prompt_token_ids: bool = False): + def _make_ret( + self, + include_prompt_token_ids: bool = False, + include_logprobs: bool = False, + ): + output_token_logprobs = [] + if include_logprobs: + output_token_logprobs = [ + (-0.5, 100, "Test"), + (-0.3, 200, " response"), + (-0.1, 300, ""), + ] ret = { "text": "Test response", "output_ids": MOCK_OUTPUT_TOKEN_IDS, @@ -374,7 +373,7 @@ def _make_ret(self, include_prompt_token_ids: bool = False): "completion_tokens": 3, "cached_tokens": 0, "finish_reason": {"type": "stop", "matched": None}, - "output_token_logprobs": [], + "output_token_logprobs": output_token_logprobs, "output_top_logprobs": None, "weight_version": "default", }, @@ -420,6 +419,46 @@ def test_json_round_trip(self): data = json.loads(response.model_dump_json()) self.assertEqual(data["choices"][0]["prompt_token_ids"], MOCK_PROMPT_TOKEN_IDS) + def test_response_token_ids_with_logprobs(self): + """response_token_ids should be populated when logprobs=True.""" + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "Hi"}], + logprobs=True, + ) + ret = [self._make_ret(include_logprobs=True)] + response = self.chat._build_chat_response(req, ret, created=0) + + self.assertIsInstance(response, ChatCompletionResponse) + self.assertEqual(response.choices[0].response_token_ids, [100, 200, 300]) + + def test_response_token_ids_without_logprobs(self): + """response_token_ids should be None when logprobs is not enabled.""" + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "Hi"}], + ) + ret = [self._make_ret(include_logprobs=False)] + response = self.chat._build_chat_response(req, ret, created=0) + + self.assertIsNone(response.choices[0].response_token_ids) + + data = json.loads(response.model_dump_json()) + self.assertNotIn("response_token_ids", data["choices"][0]) + + def test_response_token_ids_json_round_trip(self): + """response_token_ids should survive JSON round-trip.""" + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "Hi"}], + logprobs=True, + ) + ret = [self._make_ret(include_logprobs=True)] + response = self.chat._build_chat_response(req, ret, created=0) + + data = json.loads(response.model_dump_json()) + self.assertEqual(data["choices"][0]["response_token_ids"], [100, 200, 300]) + # =========================================================================== # 5. ReqState Tests From 93db23b665a9a6a72826e1afe6d5856d00e9b46c Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 11 Mar 2026 05:06:22 +0000 Subject: [PATCH 2/9] partial --- .../sglang/srt/entrypoints/openai/protocol.py | 2 + .../srt/entrypoints/openai/serving_chat.py | 127 ++++++++++++------ python/sglang/srt/entrypoints/openai/utils.py | 117 ++++++++++++++++ .../sglang/srt/managers/tokenizer_manager.py | 6 +- 4 files changed, 213 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 5140fafa49b5..9588ae994d7e 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -605,6 +605,7 @@ class ChatCompletionRequest(BaseModel): # Pretokenized input: skip tokenization for first N messages pretokenized_token_ids: Optional[List[int]] = None pretokenized_num_message: Optional[int] = None + pretokenize_mismatch: Optional[bool] = None # For request id rid: Optional[Union[List[str], str]] = None @@ -1381,6 +1382,7 @@ class MessageProcessingResult: modalities: List[str] stop: List[str] tool_call_constraint: Optional[ToolCallConstraint] = None + pretokenize_mismatch: Optional[bool] = None class ToolCallProcessingResult(NamedTuple): diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 76cb436ca659..672f73d00586 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -38,6 +38,7 @@ from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor from sglang.srt.entrypoints.openai.utils import ( + apply_chat_template_to_additional_message, process_cached_tokens_details_from_ret, process_hidden_states_from_ret, process_routed_experts_from_ret, @@ -262,6 +263,7 @@ def _convert_to_internal_request( # Process messages and apply chat template processed_messages = self._process_messages(request, is_multimodal) + request.pretokenize_mismatch = processed_messages.pretokenize_mismatch # Build sampling parameters sampling_params = request.to_sampling_params( @@ -636,11 +638,19 @@ def _apply_pretokenized_template( N = request.pretokenized_num_message messages = request.messages - # FIXME + # Only reject pretokenized input when the request actually contains + # multimodal content (e.g. images/videos). Text-only requests on + # multimodal-capable models are fine. if self.tokenizer_manager.model_config.is_multimodal: - raise ValueError( - "Pretokenized token IDs input is not supported for multimodal models" - ) + for msg in messages: + content = getattr(msg, "content", None) + if isinstance(content, list) and any( + getattr(part, "type", None) not in (None, "text") + for part in content + ): + raise ValueError( + "Pretokenized token IDs input is not supported when the request contains multimodal content" + ) if N > len(messages): raise ValueError( @@ -660,7 +670,6 @@ def _apply_pretokenized_template( ) all_msg_dicts = [msg.model_dump() for msg in messages] - prefix_msgs = all_msg_dicts[:N] # Process tool_calls arguments: str -> dict (consistent with standard path) for msg in all_msg_dicts: @@ -679,45 +688,35 @@ def _apply_pretokenized_template( chat_template_kwargs = request.chat_template_kwargs or {} - # Tokenize first N messages (no generation prompt) - prefix_ids = self.tokenizer_manager.tokenizer.apply_chat_template( - prefix_msgs, - tokenize=True, - add_generation_prompt=False, + new_msgs = all_msg_dicts[N:] + tokenization_result = apply_chat_template_to_additional_message( + new_msgs, + self.tokenizer_manager.tokenizer, tools=tools, reasoning_effort=request.reasoning_effort, - return_dict=False, - **chat_template_kwargs, + chat_template_kwargs=chat_template_kwargs, + pretokenized_last_token_id=request.pretokenized_token_ids[-1], ) + incremental_ids = tokenization_result.token_ids - # Tokenize all messages (with generation prompt) - full_ids = self.tokenizer_manager.tokenizer.apply_chat_template( - all_msg_dicts, - tokenize=True, - add_generation_prompt=True, - tools=tools, - reasoning_effort=request.reasoning_effort, - return_dict=False, - **chat_template_kwargs, + # Debug info for pretokenized concat issues + pretokenized_text = self.tokenizer_manager.tokenizer.decode( + request.pretokenized_token_ids + ) + incremental_text = self.tokenizer_manager.tokenizer.decode(incremental_ids) + logger.info( + "Pretokenized concat debug: rid=%s N=%s pretokenized_len=%d " + "incremental_len=%d pretokenized_text=%r incremental_text=%r " + "tokenization_meta=%s", + request.rid, + N, + len(request.pretokenized_token_ids), + len(incremental_ids), + pretokenized_text, + incremental_text, + tokenization_result.meta_info, ) - # Validate prefix match - if full_ids[: len(prefix_ids)] != prefix_ids: - raise ValueError( - "Prefix mismatch: tokenizing first N messages does not match " - "the prefix of tokenizing all messages. " - f"prefix_ids length={len(prefix_ids)}, full_ids length={len(full_ids)}" - ) - - incremental_ids = full_ids[len(prefix_ids) :] - # Fault tolerance: templates like qwen3 insert a trailing `\n` after - # the last assistant message. Its loss-mask is 0, so we only prepend - # it for alignment when it really is a whitespace-only token. - if ( - prefix_ids[-1] != request.pretokenized_token_ids[-1] - and self.tokenizer_manager.tokenizer.decode([prefix_ids[-1]]).strip() == "" - ): - incremental_ids = [prefix_ids[-1]] + incremental_ids prompt_ids = list(request.pretokenized_token_ids) + incremental_ids # Decode prompt_ids to text for the prompt field @@ -732,6 +731,7 @@ def _apply_pretokenized_template( video_data=None, modalities=[], stop=stop, + pretokenize_mismatch=tokenization_result.meta_info["prefix_mismatch"], ) async def _handle_streaming_request( @@ -1041,6 +1041,14 @@ def _build_chat_response( created: int, ) -> Union[ChatCompletionResponse, ORJSONResponse]: """Build chat completion response from generation results""" + logger.info( + "Build chat response start: rid=%s n_ret=%d logprobs=%s return_meta_info=%s return_prompt_token_ids=%s", + request.rid, + len(ret), + request.logprobs, + request.return_meta_info, + request.return_prompt_token_ids, + ) choices = [] # Build sglext at response level (from first ret_item, as these are per-request) @@ -1057,10 +1065,53 @@ def _build_chat_response( ) for idx, ret_item in enumerate(ret): + meta_info = ret_item.get("meta_info", {}) + if request.pretokenize_mismatch is not None: + meta_info["pretokenize_mismatch"] = request.pretokenize_mismatch + output_ids = ret_item.get("output_ids") or [] + output_token_logprobs = meta_info.get("output_token_logprobs", None) + completion_tokens = meta_info.get("completion_tokens", None) + output_logprobs_len = ( + len(output_token_logprobs) + if isinstance(output_token_logprobs, list) + else None + ) + text_for_debug = ret_item.get("text", "") + logger.info( + "Build chat response choice[%d]: finish_reason=%s completion_tokens=%s output_ids_len=%d output_logprobs_len=%s text_len=%d", + idx, + meta_info.get("finish_reason", None), + completion_tokens, + len(output_ids), + output_logprobs_len, + len(text_for_debug), + ) + if ( + request.logprobs + and isinstance(completion_tokens, int) + and isinstance(output_token_logprobs, list) + and completion_tokens != output_logprobs_len + ): + logger.warning( + "Build chat response choice[%d] logprobs mismatch: completion_tokens=%d output_logprobs_len=%d output_ids_len=%d output_ids_tail=%s output_logprobs_tail=%s text_tail=%r", + idx, + completion_tokens, + output_logprobs_len, + len(output_ids), + output_ids[-10:], + output_token_logprobs[-10:], + text_for_debug[-200:], + ) + # Process logprobs choice_logprobs = None if request.logprobs: choice_logprobs = self._process_response_logprobs(ret_item) + logger.info( + "Build chat response choice[%d]: processed_choice_logprobs_len=%d", + idx, + len(choice_logprobs.content), + ) # Handle hidden states hidden_states = process_hidden_states_from_ret(ret_item, request) diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py index 796f8f59b357..212faabd5ba1 100644 --- a/python/sglang/srt/entrypoints/openai/utils.py +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -1,4 +1,5 @@ import logging +from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Union from sglang.srt.entrypoints.openai.protocol import ( @@ -114,3 +115,119 @@ def process_cached_tokens_details_from_ret( device=details.get("device", 0), host=details.get("host", 0), ) + +@dataclass +class ChatTemplateTokenizationResult: + token_ids: list[int] + meta_info: Dict[str, Any] = field(default_factory=dict) + + +_DUMMY_USER = {"role": "user", "content": "dummy"} + +def apply_chat_template_to_additional_message( + new_messages: list[dict[str, Any]], + tokenizer, + tools=None, + reasoning_effort=None, + chat_template_kwargs: Optional[Dict[str, Any]] = None, + pretokenized_last_token_id: Optional[int] = None, +) -> ChatTemplateTokenizationResult: + if chat_template_kwargs is None: + chat_template_kwargs = {} + dummy_assistant = _build_dummy_assistant(new_messages) + # Add a dummy tool response here to avoid templates like glm-4.7 use + # <|observation|> as the start of consecutive tool responses in chat + # template, but also use <|observation|> as the stop token in assistant + # message. + base_messages = [_DUMMY_USER, dummy_assistant, _build_dummy_tool_response()] + + messages_without = base_messages + messages_with = base_messages + new_messages + + tokens_with = tokenizer.apply_chat_template( + messages_with, + tokenize=True, + add_generation_prompt=True, + tools=tools, + reasoning_effort=reasoning_effort, + return_dict=False, + **chat_template_kwargs, + ) + tokens_without = tokenizer.apply_chat_template( + messages_without, + tokenize=True, + add_generation_prompt=False, + tools=tools, + reasoning_effort=reasoning_effort, + return_dict=False, + **chat_template_kwargs, + ) + + # Validate prefix match: only raise in debug mode, otherwise warn + prefix_mismatch = tokens_with[: len(tokens_without)] != tokens_without + if prefix_mismatch: + msg = ( + "Token prefix mismatch when tokenizing additional messages. " + "This can happen for thinking models or models with special " + "chat templates that do not produce append-only token id lists. " + f"tokens_without_len={len(tokens_without)}, " + f"tokens_with_len={len(tokens_with)}, " + f"decoded_with={tokenizer.decode(tokens_with)!r}, " + f"decoded_without={tokenizer.decode(tokens_without)!r}" + ) + if logger.isEnabledFor(logging.DEBUG): + raise ValueError(msg) + logger.warning(msg) + + incremental_ids = tokens_with[len(tokens_without) :] + + # Fault tolerance: templates like qwen3 insert a trailing `\n` after + # the last assistant message. Its loss-mask is 0, so we only prepend + # it for alignment when it really is a whitespace-only token. + trailing_token_prepended = False + if ( + pretokenized_last_token_id is not None + and tokens_without[-1] != pretokenized_last_token_id + and tokenizer.decode([tokens_without[-1]]).strip() == "" + ): + incremental_ids = [tokens_without[-1]] + incremental_ids + trailing_token_prepended = True + + return ChatTemplateTokenizationResult( + token_ids=incremental_ids, + meta_info={ + "prefix_mismatch": prefix_mismatch, + "trailing_token_prepended": trailing_token_prepended, + }, + ) + + +def _build_dummy_assistant(tool_responses: list[dict[str, Any]]) -> dict[str, Any]: + return { + "role": "assistant", + "content": "", + "reasoning_content": " ", + "tool_calls": [ + { + "id": resp.get("tool_call_id", f"call0000{i}"), + "type": "function", + "function": { + "name": resp.get("name", "dummy_func"), + "arguments": {}, + }, + } + for i, resp in enumerate(tool_responses) + ], + } + +def _build_dummy_tool_response() -> dict[str, Any]: + return { + "role": "tool", + "content": "", + "tool_call_id": "call0000", + "type": "function", + "function": { + "name": "dummy_func", + "arguments": {}, + }, + } diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index ac8a1dee094d..be6e45ffe0aa 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1813,7 +1813,11 @@ def detokenize_logprob_tokens( ] else: assert self.tokenizer is not None - token_texts = self.tokenizer.batch_decode(token_logprobs_idx) + # Wrap each token ID in its own list for batch_decode to decode them separately + # batch_decode([1, 2, 3]) concatenates tokens, batch_decode([[1], [2], [3]]) decodes separately + token_texts = self.tokenizer.batch_decode( + [[idx] for idx in token_logprobs_idx] + ) return list(zip(token_logprobs_val, token_logprobs_idx, token_texts)) def detokenize_top_logprobs_tokens( From 08cea2086f85357c62f61fce7846bcac2fbbfd51 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 11 Mar 2026 22:42:53 +0000 Subject: [PATCH 3/9] remove unused --- .../sglang/srt/entrypoints/openai/protocol.py | 2 + .../srt/entrypoints/openai/serving_chat.py | 50 ++++---- python/sglang/srt/entrypoints/openai/utils.py | 117 +----------------- 3 files changed, 28 insertions(+), 141 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 9588ae994d7e..d63b010beaf3 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -606,6 +606,8 @@ class ChatCompletionRequest(BaseModel): pretokenized_token_ids: Optional[List[int]] = None pretokenized_num_message: Optional[int] = None pretokenize_mismatch: Optional[bool] = None + # AdditionalMessageTokenizer type (e.g. "default", "qwen3", "glm47") + additional_tokenizer: Optional[str] = None # For request id rid: Optional[Union[List[str], str]] = None diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 672f73d00586..85218b558122 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -37,8 +37,10 @@ ) from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor +from miles.utils.chat_template_utils.additional_message_tokenizer import ( + get_additional_message_tokenizer, +) from sglang.srt.entrypoints.openai.utils import ( - apply_chat_template_to_additional_message, process_cached_tokens_details_from_ret, process_hidden_states_from_ret, process_routed_experts_from_ret, @@ -688,34 +690,33 @@ def _apply_pretokenized_template( chat_template_kwargs = request.chat_template_kwargs or {} - new_msgs = all_msg_dicts[N:] - tokenization_result = apply_chat_template_to_additional_message( - new_msgs, + additional_tokenizer = get_additional_message_tokenizer( self.tokenizer_manager.tokenizer, - tools=tools, - reasoning_effort=request.reasoning_effort, + tokenizer_type=request.additional_tokenizer or "default", chat_template_kwargs=chat_template_kwargs, - pretokenized_last_token_id=request.pretokenized_token_ids[-1], ) - incremental_ids = tokenization_result.token_ids + new_messages = all_msg_dicts[N:] + incremental_ids = additional_tokenizer.tokenize_additional( + new_messages=new_messages, + pretokenized_token_ids=request.pretokenized_token_ids, + tools=tools, + ) # Debug info for pretokenized concat issues - pretokenized_text = self.tokenizer_manager.tokenizer.decode( - request.pretokenized_token_ids - ) - incremental_text = self.tokenizer_manager.tokenizer.decode(incremental_ids) - logger.info( - "Pretokenized concat debug: rid=%s N=%s pretokenized_len=%d " - "incremental_len=%d pretokenized_text=%r incremental_text=%r " - "tokenization_meta=%s", - request.rid, - N, - len(request.pretokenized_token_ids), - len(incremental_ids), - pretokenized_text, - incremental_text, - tokenization_result.meta_info, - ) + # pretokenized_text = self.tokenizer_manager.tokenizer.decode( + # request.pretokenized_token_ids + # ) + # incremental_text = self.tokenizer_manager.tokenizer.decode(incremental_ids) + # logger.info( + # "Pretokenized concat debug: rid=%s N=%s pretokenized_len=%d " + # "incremental_len=%d pretokenized_text=%r incremental_text=%r", + # request.rid, + # N, + # len(request.pretokenized_token_ids), + # len(incremental_ids), + # pretokenized_text, + # incremental_text, + # ) prompt_ids = list(request.pretokenized_token_ids) + incremental_ids @@ -731,7 +732,6 @@ def _apply_pretokenized_template( video_data=None, modalities=[], stop=stop, - pretokenize_mismatch=tokenization_result.meta_info["prefix_mismatch"], ) async def _handle_streaming_request( diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py index 212faabd5ba1..bd97028c1bdf 100644 --- a/python/sglang/srt/entrypoints/openai/utils.py +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -1,5 +1,6 @@ import logging from dataclasses import dataclass, field +from pprint import pformat from typing import Any, Dict, List, Optional, Union from sglang.srt.entrypoints.openai.protocol import ( @@ -115,119 +116,3 @@ def process_cached_tokens_details_from_ret( device=details.get("device", 0), host=details.get("host", 0), ) - -@dataclass -class ChatTemplateTokenizationResult: - token_ids: list[int] - meta_info: Dict[str, Any] = field(default_factory=dict) - - -_DUMMY_USER = {"role": "user", "content": "dummy"} - -def apply_chat_template_to_additional_message( - new_messages: list[dict[str, Any]], - tokenizer, - tools=None, - reasoning_effort=None, - chat_template_kwargs: Optional[Dict[str, Any]] = None, - pretokenized_last_token_id: Optional[int] = None, -) -> ChatTemplateTokenizationResult: - if chat_template_kwargs is None: - chat_template_kwargs = {} - dummy_assistant = _build_dummy_assistant(new_messages) - # Add a dummy tool response here to avoid templates like glm-4.7 use - # <|observation|> as the start of consecutive tool responses in chat - # template, but also use <|observation|> as the stop token in assistant - # message. - base_messages = [_DUMMY_USER, dummy_assistant, _build_dummy_tool_response()] - - messages_without = base_messages - messages_with = base_messages + new_messages - - tokens_with = tokenizer.apply_chat_template( - messages_with, - tokenize=True, - add_generation_prompt=True, - tools=tools, - reasoning_effort=reasoning_effort, - return_dict=False, - **chat_template_kwargs, - ) - tokens_without = tokenizer.apply_chat_template( - messages_without, - tokenize=True, - add_generation_prompt=False, - tools=tools, - reasoning_effort=reasoning_effort, - return_dict=False, - **chat_template_kwargs, - ) - - # Validate prefix match: only raise in debug mode, otherwise warn - prefix_mismatch = tokens_with[: len(tokens_without)] != tokens_without - if prefix_mismatch: - msg = ( - "Token prefix mismatch when tokenizing additional messages. " - "This can happen for thinking models or models with special " - "chat templates that do not produce append-only token id lists. " - f"tokens_without_len={len(tokens_without)}, " - f"tokens_with_len={len(tokens_with)}, " - f"decoded_with={tokenizer.decode(tokens_with)!r}, " - f"decoded_without={tokenizer.decode(tokens_without)!r}" - ) - if logger.isEnabledFor(logging.DEBUG): - raise ValueError(msg) - logger.warning(msg) - - incremental_ids = tokens_with[len(tokens_without) :] - - # Fault tolerance: templates like qwen3 insert a trailing `\n` after - # the last assistant message. Its loss-mask is 0, so we only prepend - # it for alignment when it really is a whitespace-only token. - trailing_token_prepended = False - if ( - pretokenized_last_token_id is not None - and tokens_without[-1] != pretokenized_last_token_id - and tokenizer.decode([tokens_without[-1]]).strip() == "" - ): - incremental_ids = [tokens_without[-1]] + incremental_ids - trailing_token_prepended = True - - return ChatTemplateTokenizationResult( - token_ids=incremental_ids, - meta_info={ - "prefix_mismatch": prefix_mismatch, - "trailing_token_prepended": trailing_token_prepended, - }, - ) - - -def _build_dummy_assistant(tool_responses: list[dict[str, Any]]) -> dict[str, Any]: - return { - "role": "assistant", - "content": "", - "reasoning_content": " ", - "tool_calls": [ - { - "id": resp.get("tool_call_id", f"call0000{i}"), - "type": "function", - "function": { - "name": resp.get("name", "dummy_func"), - "arguments": {}, - }, - } - for i, resp in enumerate(tool_responses) - ], - } - -def _build_dummy_tool_response() -> dict[str, Any]: - return { - "role": "tool", - "content": "", - "tool_call_id": "call0000", - "type": "function", - "function": { - "name": "dummy_func", - "arguments": {}, - }, - } From 3c80b472256782ed7edfe5a8a809c3868674f3b5 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 11 Mar 2026 22:48:16 +0000 Subject: [PATCH 4/9] remove --- .../sglang/srt/entrypoints/openai/protocol.py | 2 - .../srt/entrypoints/openai/serving_chat.py | 69 ------------------- python/sglang/srt/entrypoints/openai/utils.py | 2 - 3 files changed, 73 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index d63b010beaf3..5796e938b8c3 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -605,7 +605,6 @@ class ChatCompletionRequest(BaseModel): # Pretokenized input: skip tokenization for first N messages pretokenized_token_ids: Optional[List[int]] = None pretokenized_num_message: Optional[int] = None - pretokenize_mismatch: Optional[bool] = None # AdditionalMessageTokenizer type (e.g. "default", "qwen3", "glm47") additional_tokenizer: Optional[str] = None @@ -1384,7 +1383,6 @@ class MessageProcessingResult: modalities: List[str] stop: List[str] tool_call_constraint: Optional[ToolCallConstraint] = None - pretokenize_mismatch: Optional[bool] = None class ToolCallProcessingResult(NamedTuple): diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 85218b558122..6fa127d3492f 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -265,8 +265,6 @@ def _convert_to_internal_request( # Process messages and apply chat template processed_messages = self._process_messages(request, is_multimodal) - request.pretokenize_mismatch = processed_messages.pretokenize_mismatch - # Build sampling parameters sampling_params = request.to_sampling_params( stop=processed_messages.stop, @@ -702,22 +700,6 @@ def _apply_pretokenized_template( tools=tools, ) - # Debug info for pretokenized concat issues - # pretokenized_text = self.tokenizer_manager.tokenizer.decode( - # request.pretokenized_token_ids - # ) - # incremental_text = self.tokenizer_manager.tokenizer.decode(incremental_ids) - # logger.info( - # "Pretokenized concat debug: rid=%s N=%s pretokenized_len=%d " - # "incremental_len=%d pretokenized_text=%r incremental_text=%r", - # request.rid, - # N, - # len(request.pretokenized_token_ids), - # len(incremental_ids), - # pretokenized_text, - # incremental_text, - # ) - prompt_ids = list(request.pretokenized_token_ids) + incremental_ids # Decode prompt_ids to text for the prompt field @@ -1041,14 +1023,6 @@ def _build_chat_response( created: int, ) -> Union[ChatCompletionResponse, ORJSONResponse]: """Build chat completion response from generation results""" - logger.info( - "Build chat response start: rid=%s n_ret=%d logprobs=%s return_meta_info=%s return_prompt_token_ids=%s", - request.rid, - len(ret), - request.logprobs, - request.return_meta_info, - request.return_prompt_token_ids, - ) choices = [] # Build sglext at response level (from first ret_item, as these are per-request) @@ -1065,53 +1039,10 @@ def _build_chat_response( ) for idx, ret_item in enumerate(ret): - meta_info = ret_item.get("meta_info", {}) - if request.pretokenize_mismatch is not None: - meta_info["pretokenize_mismatch"] = request.pretokenize_mismatch - output_ids = ret_item.get("output_ids") or [] - output_token_logprobs = meta_info.get("output_token_logprobs", None) - completion_tokens = meta_info.get("completion_tokens", None) - output_logprobs_len = ( - len(output_token_logprobs) - if isinstance(output_token_logprobs, list) - else None - ) - text_for_debug = ret_item.get("text", "") - logger.info( - "Build chat response choice[%d]: finish_reason=%s completion_tokens=%s output_ids_len=%d output_logprobs_len=%s text_len=%d", - idx, - meta_info.get("finish_reason", None), - completion_tokens, - len(output_ids), - output_logprobs_len, - len(text_for_debug), - ) - if ( - request.logprobs - and isinstance(completion_tokens, int) - and isinstance(output_token_logprobs, list) - and completion_tokens != output_logprobs_len - ): - logger.warning( - "Build chat response choice[%d] logprobs mismatch: completion_tokens=%d output_logprobs_len=%d output_ids_len=%d output_ids_tail=%s output_logprobs_tail=%s text_tail=%r", - idx, - completion_tokens, - output_logprobs_len, - len(output_ids), - output_ids[-10:], - output_token_logprobs[-10:], - text_for_debug[-200:], - ) - # Process logprobs choice_logprobs = None if request.logprobs: choice_logprobs = self._process_response_logprobs(ret_item) - logger.info( - "Build chat response choice[%d]: processed_choice_logprobs_len=%d", - idx, - len(choice_logprobs.content), - ) # Handle hidden states hidden_states = process_hidden_states_from_ret(ret_item, request) diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py index bd97028c1bdf..796f8f59b357 100644 --- a/python/sglang/srt/entrypoints/openai/utils.py +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -1,6 +1,4 @@ import logging -from dataclasses import dataclass, field -from pprint import pformat from typing import Any, Dict, List, Optional, Union from sglang.srt.entrypoints.openai.protocol import ( From ddc8dc953fb560bc9afc7f9e707b9f54f0fee016 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 11 Mar 2026 22:52:14 +0000 Subject: [PATCH 5/9] remove --- .../basic/test_return_token_ids.py | 81 ++----------------- 1 file changed, 5 insertions(+), 76 deletions(-) diff --git a/test/registered/openai_server/basic/test_return_token_ids.py b/test/registered/openai_server/basic/test_return_token_ids.py index ca9474c4b996..bb9cb5a67d11 100644 --- a/test/registered/openai_server/basic/test_return_token_ids.py +++ b/test/registered/openai_server/basic/test_return_token_ids.py @@ -1,12 +1,11 @@ """ -Unit tests for token ID return features in ChatCompletion endpoint. +Unit tests for the return_prompt_token_ids feature in ChatCompletion endpoint. Tests that: 1. Protocol models correctly handle return_prompt_token_ids / prompt_token_ids fields -2. response_token_ids is populated when logprobs=True (non-streaming) -3. Request conversion passes return_prompt_token_ids flag through -4. Non-streaming response includes prompt_token_ids and response_token_ids -5. Fields are omitted from JSON when not applicable +2. Request conversion passes return_prompt_token_ids flag through +3. Non-streaming response includes prompt_token_ids +4. Fields are omitted from JSON when return_prompt_token_ids is False (default) Run with: python -m pytest test/registered/openai_server/basic/test_return_token_ids.py -v @@ -144,28 +143,6 @@ def test_choice_includes_prompt_token_ids_when_set(self): self.assertIn("prompt_token_ids", data) self.assertEqual(data["prompt_token_ids"], [1, 2, 3]) - # --- response_token_ids --- - - def test_choice_omits_response_token_ids_when_none(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hi"), - finish_reason="stop", - ) - data = choice.model_dump() - self.assertNotIn("response_token_ids", data) - - def test_choice_includes_response_token_ids_when_set(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hi"), - finish_reason="stop", - response_token_ids=MOCK_OUTPUT_TOKEN_IDS, - ) - data = choice.model_dump() - self.assertIn("response_token_ids", data) - self.assertEqual(data["response_token_ids"], MOCK_OUTPUT_TOKEN_IDS) - # --- Full JSON round-trip --- def test_full_response_json_with_prompt_token_ids(self): @@ -355,15 +332,7 @@ def setUp(self): def _make_ret( self, include_prompt_token_ids: bool = False, - include_logprobs: bool = False, ): - output_token_logprobs = [] - if include_logprobs: - output_token_logprobs = [ - (-0.5, 100, "Test"), - (-0.3, 200, " response"), - (-0.1, 300, ""), - ] ret = { "text": "Test response", "output_ids": MOCK_OUTPUT_TOKEN_IDS, @@ -373,7 +342,7 @@ def _make_ret( "completion_tokens": 3, "cached_tokens": 0, "finish_reason": {"type": "stop", "matched": None}, - "output_token_logprobs": output_token_logprobs, + "output_token_logprobs": [], "output_top_logprobs": None, "weight_version": "default", }, @@ -419,46 +388,6 @@ def test_json_round_trip(self): data = json.loads(response.model_dump_json()) self.assertEqual(data["choices"][0]["prompt_token_ids"], MOCK_PROMPT_TOKEN_IDS) - def test_response_token_ids_with_logprobs(self): - """response_token_ids should be populated when logprobs=True.""" - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - logprobs=True, - ) - ret = [self._make_ret(include_logprobs=True)] - response = self.chat._build_chat_response(req, ret, created=0) - - self.assertIsInstance(response, ChatCompletionResponse) - self.assertEqual(response.choices[0].response_token_ids, [100, 200, 300]) - - def test_response_token_ids_without_logprobs(self): - """response_token_ids should be None when logprobs is not enabled.""" - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - ) - ret = [self._make_ret(include_logprobs=False)] - response = self.chat._build_chat_response(req, ret, created=0) - - self.assertIsNone(response.choices[0].response_token_ids) - - data = json.loads(response.model_dump_json()) - self.assertNotIn("response_token_ids", data["choices"][0]) - - def test_response_token_ids_json_round_trip(self): - """response_token_ids should survive JSON round-trip.""" - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - logprobs=True, - ) - ret = [self._make_ret(include_logprobs=True)] - response = self.chat._build_chat_response(req, ret, created=0) - - data = json.loads(response.model_dump_json()) - self.assertEqual(data["choices"][0]["response_token_ids"], [100, 200, 300]) - # =========================================================================== # 5. ReqState Tests From 6b228f86c105cfa53789043e320d44152debec0f Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Thu, 12 Mar 2026 20:03:59 +0000 Subject: [PATCH 6/9] fix --- python/sglang/srt/entrypoints/openai/protocol.py | 4 ++-- python/sglang/srt/entrypoints/openai/serving_chat.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 5796e938b8c3..eb85d0aa8220 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -605,8 +605,8 @@ class ChatCompletionRequest(BaseModel): # Pretokenized input: skip tokenization for first N messages pretokenized_token_ids: Optional[List[int]] = None pretokenized_num_message: Optional[int] = None - # AdditionalMessageTokenizer type (e.g. "default", "qwen3", "glm47") - additional_tokenizer: Optional[str] = None + # TITOTokenizer type (e.g. "default", "qwen3", "glm47") + tito_model: Optional[str] = None # For request id rid: Optional[Union[List[str], str]] = None diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 6fa127d3492f..eaaa7d6090ba 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -37,8 +37,8 @@ ) from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor -from miles.utils.chat_template_utils.additional_message_tokenizer import ( - get_additional_message_tokenizer, +from miles.utils.chat_template_utils.tito_tokenizer import ( + get_tito_tokenizer, ) from sglang.srt.entrypoints.openai.utils import ( process_cached_tokens_details_from_ret, @@ -688,13 +688,13 @@ def _apply_pretokenized_template( chat_template_kwargs = request.chat_template_kwargs or {} - additional_tokenizer = get_additional_message_tokenizer( + tito_tokenizer = get_tito_tokenizer( self.tokenizer_manager.tokenizer, - tokenizer_type=request.additional_tokenizer or "default", + tokenizer_type=request.tito_model or "default", chat_template_kwargs=chat_template_kwargs, ) new_messages = all_msg_dicts[N:] - incremental_ids = additional_tokenizer.tokenize_additional( + incremental_ids = tito_tokenizer.tokenize_additional( new_messages=new_messages, pretokenized_token_ids=request.pretokenized_token_ids, tools=tools, From d15bbc6a4c98002ab5ce04062321336e0638a300 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Fri, 13 Mar 2026 22:02:33 +0000 Subject: [PATCH 7/9] fix content None --- python/sglang/srt/entrypoints/openai/serving_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index eaaa7d6090ba..98abf0fbf653 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -1100,12 +1100,12 @@ def _build_chat_response( choice_meta_info = ( ret_item["meta_info"] if request.return_meta_info else None ) - + # NOTE: content should not be None but empty string to make sure retokenize consistancy. choice_data = ChatCompletionResponseChoice( index=idx, message=ChatMessage( role="assistant", - content=text if text else None, + content=text if text else "", tool_calls=tool_calls, reasoning_content=reasoning_text if reasoning_text else None, ), From 1116357dcb9609bf34db59f7e10c99c14705ab86 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 16 Mar 2026 18:38:21 +0000 Subject: [PATCH 8/9] remove useless code --- .../sglang/srt/entrypoints/openai/protocol.py | 9 +- .../srt/entrypoints/openai/serving_chat.py | 137 ++---------------- 2 files changed, 17 insertions(+), 129 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index eb85d0aa8220..d08775b483e6 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -602,11 +602,10 @@ class ChatCompletionRequest(BaseModel): custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None custom_params: Optional[Dict] = None - # Pretokenized input: skip tokenization for first N messages - pretokenized_token_ids: Optional[List[int]] = None - pretokenized_num_message: Optional[int] = None - # TITOTokenizer type (e.g. "default", "qwen3", "glm47") - tito_model: Optional[str] = None + # Pre-computed prompt token IDs: when provided, bypasses chat template + # tokenization entirely. Messages are still used to derive stop tokens + # and tool_call_constraint. + input_ids: Optional[List[int]] = None # For request id rid: Optional[Union[List[str], str]] = None diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 98abf0fbf653..a459cbc61de5 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -37,9 +37,6 @@ ) from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor -from miles.utils.chat_template_utils.tito_tokenizer import ( - get_tito_tokenizer, -) from sglang.srt.entrypoints.openai.utils import ( process_cached_tokens_details_from_ret, process_hidden_states_from_ret, @@ -359,8 +356,19 @@ def _process_messages( ) tool_call_constraint = ("json_schema", json_schema) - # Use chat template - if self.template_manager.chat_template_name is None: + # When input_ids are provided, skip template tokenization entirely; + # only stop tokens and tool_call_constraint are needed. + if request.input_ids is not None: + result = MessageProcessingResult( + prompt=self.tokenizer_manager.tokenizer.decode(request.input_ids), + prompt_ids=request.input_ids, + image_data=None, + audio_data=None, + video_data=None, + modalities=[], + stop=request.stop or [], + ) + elif self.template_manager.chat_template_name is None: result = self._apply_jinja_template(request, tools, is_multimodal) else: result = self._apply_conversation_template(request, is_multimodal) @@ -375,14 +383,6 @@ def _apply_jinja_template( is_multimodal: bool, ) -> MessageProcessingResult: """Apply Jinja chat template""" - # Short-circuit: use pretokenized path if provided - if ( - request.pretokenized_token_ids is not None - and request.pretokenized_num_message is not None - and request.pretokenized_num_message > 0 - ): - return self._apply_pretokenized_template(request, tools) - prompt = "" prompt_ids = [] openai_compatible_messages = [] @@ -605,117 +605,6 @@ def _apply_conversation_template( stop=stop, ) - def _apply_pretokenized_template( - self, - request: ChatCompletionRequest, - tools: Optional[List[Dict]], - ) -> MessageProcessingResult: - """Apply pretokenized template with incremental tokenization for tool messages. - - Supported input pattern (non-streaming, non-multimodal only): - messages[0..N-1] are pretokenized (token IDs provided via pretokenized_token_ids) - messages[N..end] are new tool responses to be incrementally tokenized - - Constraints: - - messages[N-1] must be role="assistant" (the last pretokenized turn) - - messages[N..end] must all be role="tool" or "system" - - No multimodal content allowed in any message - - Example valid pattern: - messages = [system, user, assistant(tool_calls), tool, tool, system, ...] - pretokenized_num_message = 3 (system + user + assistant are pretokenized) - pretokenized_token_ids = [...] (token IDs for the first 3 messages) - - How it works: - 1. Tokenize first N messages with apply_chat_template(add_generation_prompt=False) - to get prefix_ids - 2. Tokenize ALL messages with apply_chat_template(add_generation_prompt=True) - to get full_ids - 3. Verify prefix_ids matches full_ids[:len(prefix_ids)] - 4. incremental_ids = full_ids[len(prefix_ids):] - 5. Final prompt_ids = pretokenized_token_ids + incremental_ids - """ - N = request.pretokenized_num_message - messages = request.messages - - # Only reject pretokenized input when the request actually contains - # multimodal content (e.g. images/videos). Text-only requests on - # multimodal-capable models are fine. - if self.tokenizer_manager.model_config.is_multimodal: - for msg in messages: - content = getattr(msg, "content", None) - if isinstance(content, list) and any( - getattr(part, "type", None) not in (None, "text") - for part in content - ): - raise ValueError( - "Pretokenized token IDs input is not supported when the request contains multimodal content" - ) - - if N > len(messages): - raise ValueError( - f"pretokenized_num_message ({N}) > total messages ({len(messages)})" - ) - - if messages[N - 1].role != "assistant": - raise ValueError( - f"Message at index {N - 1} must be assistant, got {messages[N - 1].role}" - ) - - ALLOWED_APPEND_ROLES = {"tool", "system"} - for i in range(N, len(messages)): - if messages[i].role not in ALLOWED_APPEND_ROLES: - raise ValueError( - f"Message at index {i} must be one of {ALLOWED_APPEND_ROLES}, got {messages[i].role}" - ) - - all_msg_dicts = [msg.model_dump() for msg in messages] - - # Process tool_calls arguments: str -> dict (consistent with standard path) - for msg in all_msg_dicts: - if ( - msg["role"] == "assistant" - and "tool_calls" in msg - and isinstance(msg["tool_calls"], list) - ): - for item in msg["tool_calls"]: - if "arguments" in item["function"] and isinstance( - item["function"]["arguments"], str - ): - item["function"]["arguments"] = orjson.loads( - item["function"]["arguments"] - ) - - chat_template_kwargs = request.chat_template_kwargs or {} - - tito_tokenizer = get_tito_tokenizer( - self.tokenizer_manager.tokenizer, - tokenizer_type=request.tito_model or "default", - chat_template_kwargs=chat_template_kwargs, - ) - new_messages = all_msg_dicts[N:] - incremental_ids = tito_tokenizer.tokenize_additional( - new_messages=new_messages, - pretokenized_token_ids=request.pretokenized_token_ids, - tools=tools, - ) - - prompt_ids = list(request.pretokenized_token_ids) + incremental_ids - - # Decode prompt_ids to text for the prompt field - prompt = self.tokenizer_manager.tokenizer.decode(prompt_ids) - - stop = request.stop - return MessageProcessingResult( - prompt=prompt, - prompt_ids=prompt_ids, - image_data=None, - audio_data=None, - video_data=None, - modalities=[], - stop=stop, - ) - async def _handle_streaming_request( self, adapted_request: GenerateReqInput, From 01af89cfe1747b16db2d2569369b80e602385ce4 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 16 Mar 2026 22:29:35 +0000 Subject: [PATCH 9/9] remove test --- .../basic/test_return_token_ids.py | 426 ------------------ 1 file changed, 426 deletions(-) delete mode 100644 test/registered/openai_server/basic/test_return_token_ids.py diff --git a/test/registered/openai_server/basic/test_return_token_ids.py b/test/registered/openai_server/basic/test_return_token_ids.py deleted file mode 100644 index bb9cb5a67d11..000000000000 --- a/test/registered/openai_server/basic/test_return_token_ids.py +++ /dev/null @@ -1,426 +0,0 @@ -""" -Unit tests for the return_prompt_token_ids feature in ChatCompletion endpoint. - -Tests that: -1. Protocol models correctly handle return_prompt_token_ids / prompt_token_ids fields -2. Request conversion passes return_prompt_token_ids flag through -3. Non-streaming response includes prompt_token_ids -4. Fields are omitted from JSON when return_prompt_token_ids is False (default) - -Run with: - python -m pytest test/registered/openai_server/basic/test_return_token_ids.py -v -or: - python test/registered/openai_server/basic/test_return_token_ids.py -v -""" - -import json -import sys -import unittest -from unittest.mock import MagicMock - -# --------------------------------------------------------------------------- -# Mock out heavy GPU dependencies so tests run on CPU-only machines. -# We install a MagicMock for every missing module in the import chain. -# --------------------------------------------------------------------------- - -_GPU_MODULES = [ - # PyTorch - "torch", "torch.nn", "torch.nn.functional", "torch.nn.parameter", - "torch.cuda", "torch.distributed", "torch.library", "torch.utils", - "torch.utils.checkpoint", "torch.fx", "torch.profiler", - "torch.autograd", "torch.amp", "torch.optim", - # Triton - "triton", "triton.language", "triton.runtime", - # SGLang kernel / vLLM / transformers - "sgl_kernel", "vllm", "vllm.config", "vllm.model_executor", - "transformers", "transformers.models", "outlines", - # CUDA-specific - "cuda", "cupy", "numba", - # Packaging - "packaging", "packaging.version", -] - -_mock_cache = {} - -for mod_name in _GPU_MODULES: - if mod_name not in sys.modules: - mock = MagicMock() - sys.modules[mod_name] = mock - _mock_cache[mod_name] = mock - -# --------------------------------------------------------------------------- -# Now safe to import sglang modules -# --------------------------------------------------------------------------- - -import asyncio -from typing import List, Optional - -from sglang.srt.entrypoints.openai.protocol import ( - ChatCompletionRequest, - ChatCompletionResponse, - ChatCompletionResponseChoice, - ChatMessage, - UsageInfo, -) - -# These may fail on CPU-only if the import chain hits something we missed. -# We protect with try/except and skip the dependent tests. -_HAS_IO_STRUCT = False -_HAS_SERVING_CHAT = False -_HAS_TOKENIZER_MANAGER = False - -try: - from sglang.srt.managers.io_struct import GenerateReqInput - _HAS_IO_STRUCT = True -except Exception: - pass - -try: - from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat - from sglang.srt.entrypoints.openai.protocol import MessageProcessingResult - _HAS_SERVING_CHAT = True -except Exception: - pass - -try: - from sglang.srt.managers.tokenizer_manager import ReqState - _HAS_TOKENIZER_MANAGER = True -except Exception: - pass - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - -MOCK_PROMPT_TOKEN_IDS = [128000, 882, 1234, 5678, 9012] -MOCK_OUTPUT_TOKEN_IDS = [100, 200, 300] - - -# =========================================================================== -# 1. Protocol Tests — pure Pydantic model serialization (always runnable) -# =========================================================================== - - -class TestReturnTokenIdsProtocol(unittest.TestCase): - """Test protocol model fields for return_prompt_token_ids.""" - - # --- Request --- - - def test_request_default_false(self): - req = ChatCompletionRequest( - model="test", - messages=[{"role": "user", "content": "Hi"}], - ) - self.assertFalse(req.return_prompt_token_ids) - - def test_request_explicit_true(self): - req = ChatCompletionRequest( - model="test", - messages=[{"role": "user", "content": "Hi"}], - return_prompt_token_ids=True, - ) - self.assertTrue(req.return_prompt_token_ids) - - # --- Response (non-streaming) --- - - def test_choice_omits_prompt_token_ids_when_none(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hi"), - finish_reason="stop", - ) - data = choice.model_dump() - self.assertNotIn("prompt_token_ids", data) - - def test_choice_includes_prompt_token_ids_when_set(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hi"), - finish_reason="stop", - prompt_token_ids=[1, 2, 3], - ) - data = choice.model_dump() - self.assertIn("prompt_token_ids", data) - self.assertEqual(data["prompt_token_ids"], [1, 2, 3]) - - # --- Full JSON round-trip --- - - def test_full_response_json_with_prompt_token_ids(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hello"), - finish_reason="stop", - prompt_token_ids=MOCK_PROMPT_TOKEN_IDS, - ) - resp = ChatCompletionResponse( - id="test-id", - model="test", - choices=[choice], - usage=UsageInfo(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - data = json.loads(resp.model_dump_json()) - self.assertEqual(data["choices"][0]["prompt_token_ids"], MOCK_PROMPT_TOKEN_IDS) - - def test_full_response_json_without_token_ids(self): - choice = ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="hello"), - finish_reason="stop", - ) - resp = ChatCompletionResponse( - id="test-id", - model="test", - choices=[choice], - usage=UsageInfo(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - data = json.loads(resp.model_dump_json()) - self.assertNotIn("prompt_token_ids", data["choices"][0]) - - -# =========================================================================== -# 2. GenerateReqInput Tests -# =========================================================================== - - -@unittest.skipUnless(_HAS_IO_STRUCT, "io_struct import requires GPU deps") -class TestReturnTokenIdsIOStruct(unittest.TestCase): - """Test GenerateReqInput return_prompt_token_ids field.""" - - def test_default_false(self): - req = GenerateReqInput(text="hello") - self.assertFalse(req.return_prompt_token_ids) - - def test_explicit_true(self): - req = GenerateReqInput(text="hello", return_prompt_token_ids=True) - self.assertTrue(req.return_prompt_token_ids) - - def test_does_not_affect_logprob_fields(self): - req = GenerateReqInput( - text="hello", - return_prompt_token_ids=True, - return_logprob=False, - logprob_start_len=-1, - ) - self.assertTrue(req.return_prompt_token_ids) - self.assertFalse(req.return_logprob) - self.assertEqual(req.logprob_start_len, -1) - - -# =========================================================================== -# 3. Request Conversion Tests -# =========================================================================== - - -@unittest.skipUnless( - _HAS_SERVING_CHAT, "OpenAIServingChat import requires GPU deps" -) -class TestReturnTokenIdsRequestConversion(unittest.TestCase): - """Test that return_prompt_token_ids flows through _convert_to_internal_request.""" - - def setUp(self): - from unittest.mock import Mock, patch - - tm = Mock() - tm.model_config = Mock(is_multimodal=False) - tm.server_args = Mock( - enable_cache_report=False, - tool_call_parser=None, - reasoning_parser=None, - ) - mock_hf_config = Mock() - mock_hf_config.architectures = ["LlamaForCausalLM"] - tm.model_config.hf_config = mock_hf_config - tm.chat_template_name = "llama-3" - tm.tokenizer = Mock() - tm.tokenizer.encode.return_value = [1, 2, 3] - tm.tokenizer.chat_template = None - tm.tokenizer.bos_token_id = 1 - - template_mgr = Mock() - template_mgr.chat_template_name = "llama-3" - template_mgr.jinja_template_content_format = None - template_mgr.completion_template_name = None - template_mgr.force_reasoning = False - - self.chat = OpenAIServingChat(tm, template_mgr) - - def _convert(self, return_prompt_token_ids: bool): - from unittest.mock import patch - - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - return_prompt_token_ids=return_prompt_token_ids, - ) - with patch.object(self.chat, "_process_messages") as proc_mock: - proc_mock.return_value = MessageProcessingResult( - "Test prompt", [1, 2, 3], None, None, [], [""], None, - ) - adapted, _ = self.chat._convert_to_internal_request(req) - return adapted - - def test_flag_passed_when_true(self): - adapted = self._convert(return_prompt_token_ids=True) - self.assertIsInstance(adapted, GenerateReqInput) - self.assertTrue(adapted.return_prompt_token_ids) - - def test_flag_passed_when_false(self): - adapted = self._convert(return_prompt_token_ids=False) - self.assertIsInstance(adapted, GenerateReqInput) - self.assertFalse(adapted.return_prompt_token_ids) - - def test_logprob_not_affected(self): - adapted = self._convert(return_prompt_token_ids=True) - self.assertFalse(adapted.return_logprob) - self.assertEqual(adapted.logprob_start_len, -1) - - def test_stream_with_return_prompt_token_ids_raises(self): - """return_prompt_token_ids=True + stream=True should raise ValueError.""" - from unittest.mock import patch - - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - return_prompt_token_ids=True, - stream=True, - ) - with patch.object(self.chat, "_process_messages") as proc_mock: - proc_mock.return_value = MessageProcessingResult( - "Test prompt", [1, 2, 3], None, None, [], [""], None, - ) - with self.assertRaises(ValueError): - self.chat._convert_to_internal_request(req) - - -# =========================================================================== -# 4. Response Building Tests -# =========================================================================== - - -@unittest.skipUnless( - _HAS_SERVING_CHAT, "OpenAIServingChat import requires GPU deps" -) -class TestReturnTokenIdsResponseBuilding(unittest.TestCase): - """Test _build_chat_response includes prompt_token_ids when requested.""" - - def setUp(self): - from unittest.mock import Mock - - tm = Mock() - tm.model_config = Mock(is_multimodal=False) - tm.server_args = Mock( - enable_cache_report=False, - tool_call_parser=None, - reasoning_parser=None, - ) - mock_hf_config = Mock() - mock_hf_config.architectures = ["LlamaForCausalLM"] - tm.model_config.hf_config = mock_hf_config - tm.chat_template_name = "llama-3" - tm.tokenizer = Mock() - tm.tokenizer.chat_template = None - tm.tokenizer.bos_token_id = 1 - - template_mgr = Mock() - template_mgr.chat_template_name = "llama-3" - template_mgr.jinja_template_content_format = None - template_mgr.completion_template_name = None - template_mgr.force_reasoning = False - - self.chat = OpenAIServingChat(tm, template_mgr) - - def _make_ret( - self, - include_prompt_token_ids: bool = False, - ): - ret = { - "text": "Test response", - "output_ids": MOCK_OUTPUT_TOKEN_IDS, - "meta_info": { - "id": "chatcmpl-test", - "prompt_tokens": 5, - "completion_tokens": 3, - "cached_tokens": 0, - "finish_reason": {"type": "stop", "matched": None}, - "output_token_logprobs": [], - "output_top_logprobs": None, - "weight_version": "default", - }, - } - if include_prompt_token_ids: - ret["prompt_token_ids"] = MOCK_PROMPT_TOKEN_IDS - return ret - - def test_with_return_prompt_token_ids(self): - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - return_prompt_token_ids=True, - ) - ret = [self._make_ret(include_prompt_token_ids=True)] - response = self.chat._build_chat_response(req, ret, created=0) - - self.assertIsInstance(response, ChatCompletionResponse) - self.assertEqual(response.choices[0].prompt_token_ids, MOCK_PROMPT_TOKEN_IDS) - - def test_without_return_prompt_token_ids(self): - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - ) - ret = [self._make_ret(include_prompt_token_ids=False)] - response = self.chat._build_chat_response(req, ret, created=0) - - self.assertIsNone(response.choices[0].prompt_token_ids) - - data = json.loads(response.model_dump_json()) - self.assertNotIn("prompt_token_ids", data["choices"][0]) - - def test_json_round_trip(self): - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi"}], - return_prompt_token_ids=True, - ) - ret = [self._make_ret(include_prompt_token_ids=True)] - response = self.chat._build_chat_response(req, ret, created=0) - - data = json.loads(response.model_dump_json()) - self.assertEqual(data["choices"][0]["prompt_token_ids"], MOCK_PROMPT_TOKEN_IDS) - - -# =========================================================================== -# 5. ReqState Tests -# =========================================================================== - - -@unittest.skipUnless( - _HAS_TOKENIZER_MANAGER, "tokenizer_manager import requires GPU deps" -) -class TestReturnTokenIdsReqState(unittest.TestCase): - """Test that ReqState stores prompt_token_ids correctly.""" - - def test_reqstate_default_none(self): - state = ReqState( - out_list=[], - finished=False, - event=asyncio.Event(), - obj=MagicMock(), - time_stats=MagicMock(), - ) - self.assertIsNone(state.prompt_token_ids) - - def test_reqstate_stores_prompt_token_ids(self): - state = ReqState( - out_list=[], - finished=False, - event=asyncio.Event(), - obj=MagicMock(), - time_stats=MagicMock(), - ) - state.prompt_token_ids = MOCK_PROMPT_TOKEN_IDS - self.assertEqual(state.prompt_token_ids, MOCK_PROMPT_TOKEN_IDS) - - -if __name__ == "__main__": - unittest.main()