From 9aaef4296f4eb9df7392de6920b6b2ed79f354ba Mon Sep 17 00:00:00 2001 From: Shijie Lyu Date: Thu, 1 Oct 2026 09:26:55 -0700 Subject: [PATCH] [Bugfix][Frontend] Check reused prompt token ids against the vocab before streaming (#59555) Signed-off-by: Shijie Lyu --- .../chat_completion/test_serving_chat.py | 19 ++++++++-- vllm/renderers/base.py | 37 +++++++++++++++++++ vllm/renderers/online_renderer.py | 11 ++++-- vllm/v1/engine/input_processor.py | 34 ++--------------- 4 files changed, 63 insertions(+), 38 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 241920d1e3..9c80f504f7 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -564,6 +564,9 @@ class MockModelConfig: def get_diff_sampling_param(self): return self.diff_sampling_param or {} + def get_vocab_size(self) -> int: + return 50257 + @dataclass class MockParallelConfig: @@ -2860,7 +2863,16 @@ async def test_chat_kv_transfer_prompt_token_ids_rejects_invalid_ids(ids): @pytest.mark.asyncio -async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector(): +@pytest.mark.parametrize( + ("ids", "stream", "match"), + [ + ([1.5], False, "non-negative integers"), + ([999999], True, r"out of vocabulary.*kv_transfer_params\.prompt_token_ids"), + ], +) +async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector( + ids, stream, match +): """A render-time 400 on a decode request notifies the KV connector.""" mock_engine = _build_mock_engine() mock_engine.notify_kv_transfer_request_rejected = AsyncMock() @@ -2870,9 +2882,10 @@ async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector request = ChatCompletionRequest( model=MODEL_NAME, messages=[{"role": "user", "content": "hi"}], - kv_transfer_params={"prompt_token_ids": [1.5], "do_remote_prefill": True}, + stream=stream, + kv_transfer_params={"prompt_token_ids": ids, "do_remote_prefill": True}, ) - with pytest.raises(VLLMValidationError, match="non-negative integers"): + with pytest.raises(VLLMValidationError, match=match): await serving_chat.create_chat_completion(request) mock_engine.notify_kv_transfer_request_rejected.assert_awaited_once_with( diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 0936dd021e..8c1d91e032 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -417,6 +417,43 @@ class BaseRenderer(ABC, Generic[_T]): return self.tokenizer.eos_token_id + def validate_token_ids( + self, token_ids: Sequence[int], *, parameter: str | None = None + ) -> None: + """Raise VLLMValidationError if a token id is out of vocabulary. + + Skipped when the tokenizer is not initialized. + """ + tokenizer = self.tokenizer + if not token_ids or tokenizer is None: + return + + max_input_id = max(token_ids) + min_input_id = min(token_ids) + + # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while + # self.model_config.get_vocab_size() is the model’s vocab size. + # For Qwen3 models, the language model has extra tokens that do + # not exist in the tokenizer, and vice versa for multimodal + # placeholder tokens in some multimodal models. + # See https://github.com/QwenLM/Qwen3/issues/29#issuecomment-1933720399 # noqa: E501 + # and https://github.com/vllm-project/vllm/pull/22471#discussion_r2312251421 # noqa: E501 + + # Here we take the max of the two to determine if a token id is + # truly out-of-vocabulary. + model_vocab_size = self.model_config.get_vocab_size() + # A negative id is out of vocabulary just like an over-large one, + # but is not caught by the upper-bound check below. Reject it here + # so it is not used as an embedding index downstream. + if min_input_id < 0: + raise VLLMValidationError( + f"Token id {min_input_id} is out of vocabulary", parameter=parameter + ) + if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1): + raise VLLMValidationError( + f"Token id {max_input_id} is out of vocabulary", parameter=parameter + ) + def get_dec_start_token_id(self) -> int: """Obtain the decoder start token id employed by an encoder/decoder model, raising an error if it is not available. diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index bbe24f7df3..ec8b58f73d 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -103,7 +103,7 @@ def _extract_allowed_tools_from_mcp_requests( def _reused_prompt_token_ids( - request: Any, messages: list[Any] | None = None + request: Any, renderer: BaseRenderer, messages: list[Any] | None = None ) -> list[int] | None: """Pop prompt token ids forwarded for decode-side reuse, if any. @@ -113,7 +113,8 @@ def _reused_prompt_token_ids( Returns None without checking the ids when ``echo`` is set or ``messages`` has non-text content, since both need ``messages`` to be rendered. - Otherwise raises VLLMValidationError if the ids are malformed. + Otherwise raises VLLMValidationError if the ids are malformed or out of + vocabulary. """ kv = getattr(request, "kv_transfer_params", None) if not isinstance(kv, dict): @@ -144,6 +145,8 @@ def _reused_prompt_token_ids( "of non-negative integers.", parameter="kv_transfer_params.prompt_token_ids", ) + # The engine checks this too, but only after a streamed response starts. + renderer.validate_token_ids(ids, parameter="kv_transfer_params.prompt_token_ids") return ids @@ -530,7 +533,7 @@ class OnlineRenderer: should_include_tools: bool = True, ): """Build Harmony (GPT-OSS) messages and engine prompt from a chat request.""" - reuse_ids = _reused_prompt_token_ids(request) + reuse_ids = _reused_prompt_token_ids(request, self.renderer) if reuse_ids: # Decode-side token reuse: feed the forwarded ids straight to the # engine. Harmony has no adjust_request hook to preserve. @@ -781,7 +784,7 @@ class OnlineRenderer: default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) - reuse_ids = _reused_prompt_token_ids(request, messages) + reuse_ids = _reused_prompt_token_ids(request, renderer, messages) if reuse_ids: # Decode-side token reuse: feed the forwarded ids straight to the # engine, skipping templating and tokenization. ``messages`` are not diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 3113b90f3e..fa68eb8ec2 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -556,9 +556,6 @@ class InputProcessor: prompt_input: SingletonInput, prompt_type: Literal["encoder", "decoder"], ) -> None: - model_config = self.model_config - tokenizer = self.tokenizer - prompt_ids = ( None if prompt_input["type"] == "embeds" @@ -595,34 +592,9 @@ class InputProcessor: f"by setting --limit-mm-per-prompt at startup." ) - if prompt_ids and tokenizer is not None: - max_input_id = max(prompt_ids, default=0) - min_input_id = min(prompt_ids, default=0) - - # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while - # self.model_config.get_vocab_size() is the model’s vocab size. - # For Qwen3 models, the language model has extra tokens that do - # not exist in the tokenizer, and vice versa for multimodal - # placeholder tokens in some multimodal models. - # See https://github.com/QwenLM/Qwen3/issues/29#issuecomment-1933720399 # noqa: E501 - # and https://github.com/vllm-project/vllm/pull/22471#discussion_r2312251421 # noqa: E501 - - # Here we take the max of the two to determine if a token id is - # truly out-of-vocabulary. - model_vocab_size = model_config.get_vocab_size() - # A negative id is out of vocabulary just like an over-large one, - # but is not caught by the upper-bound check below. Reject it here - # so it is not used as an embedding index downstream. This - # validation path is shared by generate, embedding and pooling - # requests, so the check covers all three. - if min_input_id < 0: - raise VLLMValidationError( - f"Token id {min_input_id} is out of vocabulary" - ) - if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1): - raise VLLMValidationError( - f"Token id {max_input_id} is out of vocabulary" - ) + # Shared by generate, embedding and pooling requests. + if prompt_ids: + self.renderer.validate_token_ids(prompt_ids) def _validate_model_inputs( self,