A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 01:30 UTC
apply _validate_input to MistralTokenizer token-id chat prompts (#32448)
Signed-off-by: Vanshil Shah <vanshilshah@gmail.com>
vanshil shah committed
Jan 16, 2026 at 19:23 UTC
037a6487af3429bb3f3e1adfe3e2f5e5e95aa420
2 files changed
+100
-3
tests/entrypoints/openai/test_serving_chat.py
+95
@@ -731,6 +731,101 @@ async def test_serving_chat_should_set_correct_max_tokens():
731
assert mock_engine.generate.call_args.args[1].max_tokens == 5
732
733
734
+@pytest.mark.asyncio
735
+async def test_serving_chat_mistral_token_ids_prompt_is_validated(monkeypatch_module):
736
+ """Regression test: when the Mistral tokenizer path returns token IDs
737
+ directly, we must still apply input length + max_tokens validation.
738
+ """
739
+
740
+ mock_engine = MagicMock(spec=AsyncLLM)
741
+ mock_engine.errored = False
742
+ mock_engine.model_config = MockModelConfig()
743
+ mock_engine.input_processor = MagicMock()
744
+ mock_engine.io_processor = MagicMock()
745
+
746
+ class DummyMistralTokenizer:
747
+ def decode(self, token_ids):
748
+ # Only used for logging/validation error messages.
749
+ return "dummy"
750
+
751
+ dummy_tokenizer = DummyMistralTokenizer()
752
+ mock_engine.get_tokenizer.return_value = dummy_tokenizer
753
+
754
+ # Patch the OpenAI engine serving module to treat our dummy tokenizer
755
+ # as a MistralTokenizer. This forces the code path where chat template
756
+ # rendering can return a list[int] (token IDs).
757
+ import vllm.entrypoints.openai.engine.serving as engine_serving
758
+
759
+ monkeypatch_module.setattr(
760
+ engine_serving, "MistralTokenizer", DummyMistralTokenizer
761
+ )
762
+
763
+ serving_chat = _build_serving_chat(mock_engine)
764
+
765
+ # Force the Mistral chat template renderer to return token IDs.
766
+ # Choose a prompt length that is < max_model_len, but large enough that
767
+ # adding max_tokens should exceed the model context window.
768
+ serving_chat._apply_mistral_chat_template_async = AsyncMock(
769
+ return_value=list(range(95))
770
+ )
771
+
772
+ req = ChatCompletionRequest(
773
+ model=MODEL_NAME,
774
+ messages=[{"role": "user", "content": "what is 1+1?"}],
775
+ max_tokens=10,
776
+ )
777
+
778
+ resp = await serving_chat.create_chat_completion(req)
779
+ assert isinstance(resp, ErrorResponse)
780
+ assert "max_tokens" in resp.error.message
781
+
782
+
783
+@pytest.mark.asyncio
784
+async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(
785
+ monkeypatch_module,
786
+):
787
+ """Regression test: MistralTokenizer token-id prompts must still enforce
788
+ the max context length for the input itself (token_num >= max_model_len).
789
+ """
790
+
791
+ mock_engine = MagicMock(spec=AsyncLLM)
792
+ mock_engine.errored = False
793
+ mock_engine.model_config = MockModelConfig()
794
+ mock_engine.input_processor = MagicMock()
795
+ mock_engine.io_processor = MagicMock()
796
+
797
+ class DummyMistralTokenizer:
798
+ def decode(self, token_ids):
799
+ return "dummy"
800
+
801
+ dummy_tokenizer = DummyMistralTokenizer()
802
+ mock_engine.get_tokenizer.return_value = dummy_tokenizer
803
+
804
+ import vllm.entrypoints.openai.engine.serving as engine_serving
805
+
806
+ monkeypatch_module.setattr(
807
+ engine_serving, "MistralTokenizer", DummyMistralTokenizer
808
+ )
809
+
810
+ serving_chat = _build_serving_chat(mock_engine)
811
+
812
+ # prompt_token_ids length == max_model_len should be rejected for
813
+ # completion-like requests (ChatCompletionRequest).
814
+ serving_chat._apply_mistral_chat_template_async = AsyncMock(
815
+ return_value=list(range(mock_engine.model_config.max_model_len))
816
+ )
817
+
818
+ req = ChatCompletionRequest(
819
+ model=MODEL_NAME,
820
+ messages=[{"role": "user", "content": "what is 1+1?"}],
821
+ max_tokens=1,
822
+ )
823
+
824
+ resp = await serving_chat.create_chat_completion(req)
825
+ assert isinstance(resp, ErrorResponse)
826
+ assert "maximum context length" in resp.error.message
827
+
828
+
829
@pytest.mark.asyncio
830
async def test_serving_chat_could_load_correct_generation_config():
831
mock_model_config = MockModelConfig()
vllm/entrypoints/openai/engine/serving.py
+5
-3
@@ -1277,9 +1277,11 @@ class OpenAIServing:
1277
assert is_list_of(request_prompt, int), (
1278
"Prompt has to be either a string or a list of token ids"
1279
)
1280
- prompt_inputs = TokensPrompt(
1281
- prompt=tokenizer.decode(request_prompt),
1282
- prompt_token_ids=request_prompt,
1280
+ input_text = tokenizer.decode(request_prompt)
1281
+ prompt_inputs = self._validate_input(
1282
+ request=request,
1283
+ input_ids=request_prompt,
1284
+ input_text=input_text,
1285
)
1286
1287
engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"])