[Bugfix][Frontend] Check reused prompt token ids against the vocab before streaming (#59555)

Signed-off-by: Shijie Lyu <lshjhf@gmail.com>
This commit is contained in:
Shijie Lyu
2026-10-01 09:26:55 -07:00
committed by GitHub
parent 83cadd65d9
commit 9aaef4296f
4 changed files with 63 additions and 38 deletions
@@ -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(
+37
View File
@@ -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.
+7 -4
View File
@@ -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
+3 -31
View File
@@ -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,