Fix corrupted chat prompts on mistral_common tokenizers (tool_choice auto never fires) (#39773)

Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>

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,

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 "<s>[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()

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(