[Bugfix][Frontend] Check reused prompt token ids against the vocab before streaming (#59555)
Signed-off-by: Shijie Lyu <lshjhf@gmail.com>
Shijie Lyu committed
Oct 1, 2026 at 09:26 UTC
9aaef4296f4eb9df7392de6920b6b2ed79f354ba
4 files changed
+63
-38
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+16
-3
@@ -564,6 +564,9 @@ class MockModelConfig:
564
def get_diff_sampling_param(self):
565
return self.diff_sampling_param or {}
566
567
+ def get_vocab_size(self) -> int:
568
+ return 50257
569
+
570
571
@dataclass
572
class MockParallelConfig:
@@ -2860,7 +2863,16 @@ async def test_chat_kv_transfer_prompt_token_ids_rejects_invalid_ids(ids):
2863
2864
2865
@pytest.mark.asyncio
2863
-async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector():
2866
+@pytest.mark.parametrize(
2867
+ ("ids", "stream", "match"),
2868
+ [
2869
+ ([1.5], False, "non-negative integers"),
2870
+ ([999999], True, r"out of vocabulary.*kv_transfer_params\.prompt_token_ids"),
2871
+ ],
2872
+)
2873
+async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector(
2874
+ ids, stream, match
2875
+):
2876
"""A render-time 400 on a decode request notifies the KV connector."""
2877
mock_engine = _build_mock_engine()
2878
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
2882
request = ChatCompletionRequest(
2883
model=MODEL_NAME,
2884
messages=[{"role": "user", "content": "hi"}],
2873
- kv_transfer_params={"prompt_token_ids": [1.5], "do_remote_prefill": True},
2885
+ stream=stream,
2886
+ kv_transfer_params={"prompt_token_ids": ids, "do_remote_prefill": True},
2887
)
2875
- with pytest.raises(VLLMValidationError, match="non-negative integers"):
2888
+ with pytest.raises(VLLMValidationError, match=match):
2889
await serving_chat.create_chat_completion(request)
2890
2891
mock_engine.notify_kv_transfer_request_rejected.assert_awaited_once_with(
vllm/renderers/base.py
+37
@@ -417,6 +417,43 @@ class BaseRenderer(ABC, Generic[_T]):
417
418
return self.tokenizer.eos_token_id
419
420
+ def validate_token_ids(
421
+ self, token_ids: Sequence[int], *, parameter: str | None = None
422
+ ) -> None:
423
+ """Raise VLLMValidationError if a token id is out of vocabulary.
424
+
425
+ Skipped when the tokenizer is not initialized.
426
+ """
427
+ tokenizer = self.tokenizer
428
+ if not token_ids or tokenizer is None:
429
+ return
430
+
431
+ max_input_id = max(token_ids)
432
+ min_input_id = min(token_ids)
433
+
434
+ # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while
435
+ # self.model_config.get_vocab_size() is the model’s vocab size.
436
+ # For Qwen3 models, the language model has extra tokens that do
437
+ # not exist in the tokenizer, and vice versa for multimodal
438
+ # placeholder tokens in some multimodal models.
439
+ # See https://github.com/QwenLM/Qwen3/issues/29#issuecomment-1933720399 # noqa: E501
440
+ # and https://github.com/vllm-project/vllm/pull/22471#discussion_r2312251421 # noqa: E501
441
+
442
+ # Here we take the max of the two to determine if a token id is
443
+ # truly out-of-vocabulary.
444
+ model_vocab_size = self.model_config.get_vocab_size()
445
+ # A negative id is out of vocabulary just like an over-large one,
446
+ # but is not caught by the upper-bound check below. Reject it here
447
+ # so it is not used as an embedding index downstream.
448
+ if min_input_id < 0:
449
+ raise VLLMValidationError(
450
+ f"Token id {min_input_id} is out of vocabulary", parameter=parameter
451
+ )
452
+ if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1):
453
+ raise VLLMValidationError(
454
+ f"Token id {max_input_id} is out of vocabulary", parameter=parameter
455
+ )
456
+
457
def get_dec_start_token_id(self) -> int:
458
"""Obtain the decoder start token id employed by an encoder/decoder model,
459
raising an error if it is not available.
vllm/renderers/online_renderer.py
+7
-4
@@ -103,7 +103,7 @@ def _extract_allowed_tools_from_mcp_requests(
103
104
105
def _reused_prompt_token_ids(
106
- request: Any, messages: list[Any] | None = None
106
+ request: Any, renderer: BaseRenderer, messages: list[Any] | None = None
107
) -> list[int] | None:
108
"""Pop prompt token ids forwarded for decode-side reuse, if any.
109
@@ -113,7 +113,8 @@ def _reused_prompt_token_ids(
113
114
Returns None without checking the ids when ``echo`` is set or ``messages``
115
has non-text content, since both need ``messages`` to be rendered.
116
- Otherwise raises VLLMValidationError if the ids are malformed.
116
+ Otherwise raises VLLMValidationError if the ids are malformed or out of
117
+ vocabulary.
118
"""
119
kv = getattr(request, "kv_transfer_params", None)
120
if not isinstance(kv, dict):
@@ -144,6 +145,8 @@ def _reused_prompt_token_ids(
145
"of non-negative integers.",
146
parameter="kv_transfer_params.prompt_token_ids",
147
)
148
+ # The engine checks this too, but only after a streamed response starts.
149
+ renderer.validate_token_ids(ids, parameter="kv_transfer_params.prompt_token_ids")
150
return ids
151
152
@@ -530,7 +533,7 @@ class OnlineRenderer:
533
should_include_tools: bool = True,
534
):
535
"""Build Harmony (GPT-OSS) messages and engine prompt from a chat request."""
533
- reuse_ids = _reused_prompt_token_ids(request)
536
+ reuse_ids = _reused_prompt_token_ids(request, self.renderer)
537
if reuse_ids:
538
# Decode-side token reuse: feed the forwarded ids straight to the
539
# engine. Harmony has no adjust_request hook to preserve.
@@ -781,7 +784,7 @@ class OnlineRenderer:
784
default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
785
)
786
784
- reuse_ids = _reused_prompt_token_ids(request, messages)
787
+ reuse_ids = _reused_prompt_token_ids(request, renderer, messages)
788
if reuse_ids:
789
# Decode-side token reuse: feed the forwarded ids straight to the
790
# engine, skipping templating and tokenization. ``messages`` are not
vllm/v1/engine/input_processor.py
+3
-31
@@ -556,9 +556,6 @@ class InputProcessor:
556
prompt_input: SingletonInput,
557
prompt_type: Literal["encoder", "decoder"],
558
) -> None:
559
- model_config = self.model_config
560
- tokenizer = self.tokenizer
561
-
559
prompt_ids = (
560
None
561
if prompt_input["type"] == "embeds"
@@ -595,34 +592,9 @@ class InputProcessor:
592
f"by setting --limit-mm-per-prompt at startup."
593
)
594
598
- if prompt_ids and tokenizer is not None:
599
- max_input_id = max(prompt_ids, default=0)
600
- min_input_id = min(prompt_ids, default=0)
601
-
602
- # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while
603
- # self.model_config.get_vocab_size() is the model’s vocab size.
604
- # For Qwen3 models, the language model has extra tokens that do
605
- # not exist in the tokenizer, and vice versa for multimodal
606
- # placeholder tokens in some multimodal models.
607
- # See https://github.com/QwenLM/Qwen3/issues/29#issuecomment-1933720399 # noqa: E501
608
- # and https://github.com/vllm-project/vllm/pull/22471#discussion_r2312251421 # noqa: E501
609
-
610
- # Here we take the max of the two to determine if a token id is
611
- # truly out-of-vocabulary.
612
- model_vocab_size = model_config.get_vocab_size()
613
- # A negative id is out of vocabulary just like an over-large one,
614
- # but is not caught by the upper-bound check below. Reject it here
615
- # so it is not used as an embedding index downstream. This
616
- # validation path is shared by generate, embedding and pooling
617
- # requests, so the check covers all three.
618
- if min_input_id < 0:
619
- raise VLLMValidationError(
620
- f"Token id {min_input_id} is out of vocabulary"
621
- )
622
- if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1):
623
- raise VLLMValidationError(
624
- f"Token id {max_input_id} is out of vocabulary"
625
- )
595
+ # Shared by generate, embedding and pooling requests.
596
+ if prompt_ids:
597
+ self.renderer.validate_token_ids(prompt_ids)
598
599
def _validate_model_inputs(
600
self,