[Bugfix][Frontend] Validate mixed prompt embedding mask lengths (#57006)
Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Co-authored-by: Codex <noreply@openai.com>
shaohuaxi committed
Sep 21, 2026 at 22:50 UTC
5abf85d27ace7998a86b27d2754de239c73d0723
2 files changed
+75
-4
tests/test_request_input_bounds.py
+66
-4
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Bounds on request-controlled inputs that would otherwise amplify work.
3
+"""Bounds and alignment checks for request-controlled inputs.
4
5
-Consolidates the regression tests for the request-input amplification fixes:
6
-stop-string caps, bad-words dedup/tokenization limit, stop-token-id dedup,
7
-beam-width/sequence caps, and the DeepSeek history-scan bound.
5
+Covers prompt-mask alignment, stop-string caps, bad-words dedup/tokenization
6
+limits, stop-token-id dedup, beam-width/sequence caps, and the DeepSeek
7
+history-scan bound.
8
"""
9
10
import os
@@ -15,6 +15,7 @@ from typing import Protocol
15
from unittest.mock import Mock
16
17
import pytest
18
+import torch
19
from pydantic import ValidationError
20
21
import vllm.envs as envs
@@ -26,6 +27,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
27
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
28
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
29
from vllm.exceptions import VLLMValidationError
30
+from vllm.inputs.engine import embeds_input
31
from vllm.sampling_params import BeamSearchParams
32
from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
33
from vllm.v1.engine.input_processor import InputProcessor
@@ -33,6 +35,66 @@ from vllm.v1.engine.input_processor import InputProcessor
35
pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
36
37
38
+@pytest.fixture
39
+def prompt_embeds_processor():
40
+ model_config = Mock(
41
+ max_model_len=16,
42
+ max_logprobs=20,
43
+ logits_processors=None,
44
+ is_diffusion=False,
45
+ return_sampling_mask=False,
46
+ )
47
+ model_config.get_vocab_size.return_value = 32
48
+ model_config.try_get_generation_config.return_value = {}
49
+ config = Mock(
50
+ model_config=model_config,
51
+ speculative_config=None,
52
+ structured_outputs_config=None,
53
+ )
54
+ renderer = Mock(tokenizer=None)
55
+ renderer.get_eos_token_id.return_value = None
56
+ registry = Mock()
57
+ registry.supports_multimodal_inputs.return_value = False
58
+ return InputProcessor(config, renderer, mm_registry=registry)
59
+
60
+
61
+@pytest.mark.parametrize("mask_len", [0, 1, 2, 4])
62
+def test_prompt_embeds_mask_length_rejected_before_engine_submission(
63
+ prompt_embeds_processor, mask_len
64
+):
65
+ """Reject short, long, and broadcastable masks before they reach a worker."""
66
+ prompt = embeds_input(
67
+ torch.zeros(3, 4),
68
+ prompt_token_ids=[1, 2, 3],
69
+ is_token_ids=[False] * mask_len,
70
+ )
71
+
72
+ with pytest.raises(VLLMValidationError, match="prompt_is_token_ids") as exc_info:
73
+ prompt_embeds_processor.process_inputs(
74
+ "invalid", prompt, SamplingParams(max_tokens=1), ("generate",)
75
+ )
76
+
77
+ assert exc_info.value.parameter == "prompt_is_token_ids"
78
+
79
+
80
+@pytest.mark.parametrize(
81
+ ("token_ids", "mask"),
82
+ [(None, None), ([1, 2, 3], None), ([1, 2, 3], [True, False, True])],
83
+ ids=["pure-embeds", "token-ids-without-mask", "mixed"],
84
+)
85
+def test_prompt_embeds_valid_inputs_preserved(prompt_embeds_processor, token_ids, mask):
86
+ embeds = torch.zeros(3, 4)
87
+ prompt = embeds_input(embeds, prompt_token_ids=token_ids, is_token_ids=mask)
88
+
89
+ request = prompt_embeds_processor.process_inputs(
90
+ "valid", prompt, SamplingParams(max_tokens=1), ("generate",)
91
+ )
92
+
93
+ assert request.prompt_embeds is embeds
94
+ assert request.prompt_token_ids == token_ids
95
+ assert request.prompt_is_token_ids == mask
96
+
97
+
98
@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127])
99
def test_routed_experts_offset_validated_before_engine_submission(offset):
100
"""Reject invalid offsets even when routed-expert output is disabled."""
vllm/v1/engine/input_processor.py
+9
@@ -540,6 +540,15 @@ class InputProcessor:
540
prompt_len = length_from_prompt_token_ids_or_embeds(prompt_ids, prompt_embeds)
541
self._validate_prompt_len(prompt_len, prompt_type)
542
543
+ if prompt_input["type"] == "embeds":
544
+ is_token_ids = prompt_input.get("is_token_ids")
545
+ if is_token_ids is not None and len(is_token_ids) != prompt_len:
546
+ raise VLLMValidationError(
547
+ "prompt_is_token_ids must have the same length as prompt_embeds "
548
+ f"(expected {prompt_len}, got {len(is_token_ids)}).",
549
+ parameter="prompt_is_token_ids",
550
+ )
551
+
552
if prompt_input["type"] == "multimodal":
553
decoder_mm_positions = prompt_input["mm_placeholders"]
554
for modality, mm_positions in decoder_mm_positions.items():