diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 7ecfdc0d6..0189195eb 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -366,9 +366,42 @@ class OpenAIServingChat(OpenAIServingBase): ) except Exception: self._tokenizer_auto_adds_specials = True - self._chat_template_cache: OrderedDict[ - bytes, tuple[str, tuple[int, ...], str] - ] = OrderedDict() + self._prompt_text_round_trip_is_lossy = self._probe_prompt_text_round_trip() + self._chat_template_cache: OrderedDict[bytes, tuple[tuple[int, ...], str]] = ( + OrderedDict() + ) + + def _probe_prompt_text_round_trip(self) -> bool: + """Does rendering the chat template to text and re-encoding lose anything? + + mistral_common tokenizers emit control tokens ([INST], + [AVAILABLE_TOOLS], ...) that have no text form. Rendering to a string + turns them into literal characters and re-encoding also prepends a + second BOS, so the model sees the letters "AVAILABLE_TOOLS" instead of + the control token that frames the tool block. Encoding straight to ids + is the only faithful route on such tokenizers, so compare the two here + once and remember which to trust. + """ + probe = [{"role": "user", "content": "x"}] + try: + tokenizer = self.tokenizer_manager.tokenizer + rendered = tokenizer.apply_chat_template( + probe, tokenize=False, add_generation_prompt=True, return_dict=False + ) + encode_kwargs = ( + {"add_special_tokens": False} + if self._tokenizer_auto_adds_specials + else {} + ) + via_text = tokenizer.encode(rendered, **encode_kwargs) + via_ids = tokenizer.apply_chat_template( + probe, tokenize=True, add_generation_prompt=True, return_dict=False + ) + return list(via_text) != list(via_ids) + except Exception: + # A template that needs kwargs this probe does not supply tells us + # nothing; keep the long-standing text path. + return False def _handle_last_assistant_message( self, @@ -1084,8 +1117,23 @@ class OpenAIServingChat(OpenAIServingBase): pre-rendered input_ids with single placeholder ids and leave the text empty; pass those through rather than re-tokenizing an empty prompt. """ - if is_multimodal and not chat_encoding.spec_renders_prompt_ids( - self.chat_encoding_spec + # A lossy text round-trip makes the rendered prompt unusable, so send the + # ids instead. Only when nothing needs placeholder expansion: with media + # attached the MM processor still has to tokenize the text itself. + prefers_prompt_ids = ( + self._prompt_text_round_trip_is_lossy + and isinstance(processed_messages.prompt_ids, list) + and processed_messages.prompt_ids + and not ( + processed_messages.image_data + or processed_messages.video_data + or processed_messages.audio_data + ) + ) + if ( + is_multimodal + and not chat_encoding.spec_renders_prompt_ids(self.chat_encoding_spec) + and not prefers_prompt_ids ): return "text", processed_messages.prompt if isinstance(processed_messages.prompt_ids, str): @@ -1603,14 +1651,12 @@ class OpenAIServingChat(OpenAIServingBase): else {} ) try: - rendered_prompt, prompt_ids, decoded_prompt = ( - self._render_and_encode_chat_template( - openai_compatible_messages, - tools=tools, - template_kwargs=extra_template_kwargs, - encode_kwargs=encode_kwargs, - use_cache=is_multimodal, - ) + prompt_ids, decoded_prompt = self._render_and_encode_chat_template( + openai_compatible_messages, + tools=tools, + template_kwargs=extra_template_kwargs, + encode_kwargs=encode_kwargs, + use_cache=is_multimodal, ) except Exception: # If the first attempt fails, try with flat function-only format. @@ -1621,14 +1667,12 @@ class OpenAIServingChat(OpenAIServingBase): else None ) try: - rendered_prompt, prompt_ids, decoded_prompt = ( - self._render_and_encode_chat_template( - openai_compatible_messages, - tools=tools, - template_kwargs=extra_template_kwargs, - encode_kwargs=encode_kwargs, - use_cache=is_multimodal, - ) + prompt_ids, decoded_prompt = self._render_and_encode_chat_template( + openai_compatible_messages, + tools=tools, + template_kwargs=extra_template_kwargs, + encode_kwargs=encode_kwargs, + use_cache=is_multimodal, ) except _CHAT_TEMPLATE_CLIENT_ERRORS as template_error: # Template errors (e.g., from raise_exception in Jinja templates) @@ -1674,7 +1718,7 @@ class OpenAIServingChat(OpenAIServingBase): template_kwargs: Dict[str, Any], encode_kwargs: Dict[str, Any], use_cache: bool, - ) -> tuple[str, List[int], Optional[str]]: + ) -> tuple[List[int], Optional[str]]: cache_key = None if use_cache: try: @@ -1699,20 +1743,31 @@ class OpenAIServingChat(OpenAIServingBase): cached = self._chat_template_cache.get(cache_key) if cached is not None: self._chat_template_cache.move_to_end(cache_key) - rendered_prompt, prompt_ids, decoded_prompt = cached - return rendered_prompt, list(prompt_ids), decoded_prompt + prompt_ids, decoded_prompt = cached + return list(prompt_ids), decoded_prompt - rendered_prompt = self.tokenizer_manager.tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - tools=tools, - return_dict=False, - **template_kwargs, - ) - prompt_ids = self.tokenizer_manager.tokenizer.encode( - rendered_prompt, **encode_kwargs - ) + if self._prompt_text_round_trip_is_lossy: + # Re-encoding rendered text would drop the template's control tokens. + prompt_ids = self.tokenizer_manager.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + tools=tools, + return_dict=False, + **template_kwargs, + ) + else: + rendered_prompt = self.tokenizer_manager.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + tools=tools, + return_dict=False, + **template_kwargs, + ) + prompt_ids = self.tokenizer_manager.tokenizer.encode( + rendered_prompt, **encode_kwargs + ) decoded_prompt = ( self.tokenizer_manager.tokenizer.decode(prompt_ids) if cache_key is not None @@ -1720,15 +1775,11 @@ class OpenAIServingChat(OpenAIServingBase): ) if cache_key is not None: - self._chat_template_cache[cache_key] = ( - rendered_prompt, - tuple(prompt_ids), - decoded_prompt, - ) + self._chat_template_cache[cache_key] = (tuple(prompt_ids), decoded_prompt) if len(self._chat_template_cache) > _CHAT_TEMPLATE_CACHE_MAX_SIZE: self._chat_template_cache.popitem(last=False) - return rendered_prompt, prompt_ids, decoded_prompt + return prompt_ids, decoded_prompt def _apply_conversation_template( self, diff --git a/test/registered/unit/entrypoints/openai/test_chat_prompt_round_trip.py b/test/registered/unit/entrypoints/openai/test_chat_prompt_round_trip.py new file mode 100644 index 000000000..bb398dd37 --- /dev/null +++ b/test/registered/unit/entrypoints/openai/test_chat_prompt_round_trip.py @@ -0,0 +1,136 @@ +"""Chat prompts must not go through text on tokenizers that can't round-trip. + +mistral_common tokenizers emit control tokens that have no text form, so +rendering the template to a string and re-encoding silently replaces them with +their literal characters (and adds a second BOS). These tests pin the probe that +detects such a tokenizer and the prompt dispatch that reacts to it. +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +probe = OpenAIServingChat._probe_prompt_text_round_trip +engine_prompt = OpenAIServingChat._engine_prompt +render_and_encode = OpenAIServingChat._render_and_encode_chat_template + + +class FakeTokenizer: + """Renders to text and encodes to ids independently, like the real thing.""" + + def __init__(self, text_ids, template_ids, raises=False): + self._text_ids = text_ids + self._template_ids = template_ids + self._raises = raises + self.template_calls = [] + + def encode(self, text, **kwargs): + return [] if text == "" else list(self._text_ids) + + def apply_chat_template(self, messages, tokenize=False, **kwargs): + self.template_calls.append(tokenize) + if self._raises: + raise ValueError("template needs kwargs this probe does not pass") + return list(self._template_ids) if tokenize else "[INST]x[/INST]" + + +def _server(tokenizer, auto_adds_specials=False): + return SimpleNamespace( + tokenizer_manager=SimpleNamespace(tokenizer=tokenizer), + _tokenizer_auto_adds_specials=auto_adds_specials, + ) + + +def _messages(prompt_ids, prompt="rendered", **media): + return SimpleNamespace( + prompt=prompt, + prompt_ids=prompt_ids, + image_data=media.get("image_data"), + video_data=media.get("video_data"), + audio_data=media.get("audio_data"), + ) + + +class TestProbe(unittest.TestCase): + def test_divergent_encodings_are_lossy(self): + tok = FakeTokenizer(text_ids=[9, 9, 9, 9], template_ids=[1, 3, 4]) + self.assertTrue(probe(_server(tok))) + + def test_matching_encodings_are_not_lossy(self): + tok = FakeTokenizer(text_ids=[1, 3, 4], template_ids=[1, 3, 4]) + self.assertFalse(probe(_server(tok))) + + def test_a_template_that_raises_keeps_the_text_path(self): + tok = FakeTokenizer(text_ids=[1], template_ids=[1], raises=True) + self.assertFalse(probe(_server(tok))) + + +class TestRenderAndEncode(unittest.TestCase): + def test_lossy_tokenizer_encodes_straight_to_ids_without_a_text_render(self): + tok = FakeTokenizer(text_ids=[9, 9, 9, 9], template_ids=[1, 3, 4]) + server = SimpleNamespace( + tokenizer_manager=SimpleNamespace(tokenizer=tok), + _prompt_text_round_trip_is_lossy=True, + ) + + prompt_ids, _ = render_and_encode( + server, + [{"role": "user", "content": "x"}], + tools=None, + template_kwargs={}, + encode_kwargs={}, + use_cache=False, + ) + + self.assertEqual(prompt_ids, [1, 3, 4]) + self.assertEqual(tok.template_calls, [True]) + + +class TestEnginePrompt(unittest.TestCase): + def _server(self, lossy): + return SimpleNamespace( + _prompt_text_round_trip_is_lossy=lossy, chat_encoding_spec=None + ) + + def test_lossy_text_only_sends_ids(self): + key, value = engine_prompt( + self._server(True), _messages([1, 3, 4]), is_multimodal=True + ) + self.assertEqual(key, "input_ids") + self.assertEqual(value, [1, 3, 4]) + + def test_lossy_with_an_image_still_sends_text(self): + # The MM processor has to tokenize the text itself to expand placeholders. + key, _ = engine_prompt( + self._server(True), + _messages([1, 3, 4], image_data=["img"]), + is_multimodal=True, + ) + self.assertEqual(key, "text") + + def test_non_lossy_multimodal_is_unchanged(self): + key, value = engine_prompt( + self._server(False), _messages([1, 3, 4]), is_multimodal=True + ) + self.assertEqual(key, "text") + self.assertEqual(value, "rendered") + + def test_lossy_without_ids_falls_back_to_text(self): + key, _ = engine_prompt(self._server(True), _messages([]), is_multimodal=True) + self.assertEqual(key, "text") + + def test_text_model_sends_ids_either_way(self): + for lossy in (True, False): + key, value = engine_prompt( + self._server(lossy), _messages([1, 3, 4]), is_multimodal=False + ) + self.assertEqual(key, "input_ids") + self.assertEqual(value, [1, 3, 4]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 7fca45146..18a87f036 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -227,10 +227,10 @@ class TestChatTemplateCache(CustomTestCase): def test_cache_hit_reuses_render_encode_and_returns_an_owned_id_list(self): first = self._render() - first[1].append(99) + first[0].append(99) second = self._render() - self.assertEqual(second, ("rendered", [11, 12], "decoded")) + self.assertEqual(second, ([11, 12], "decoded")) self.tokenizer_manager.tokenizer.apply_chat_template.assert_called_once() self.tokenizer_manager.tokenizer.encode.assert_called_once() self.tokenizer_manager.tokenizer.decode.assert_called_once() @@ -292,6 +292,7 @@ class ServingChatTestCase(unittest.TestCase): self.tm = _MockTokenizerManager() self.template_manager = _MockTemplateManager() self.chat = OpenAIServingChat(self.tm, self.template_manager) + self.tm.tokenizer.reset_mock() # frequently reused requests self.basic_req = ChatCompletionRequest(