[Bugfix] Validate routed-expert prompt offsets before engine submission (#56844)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Bhushan Asati <bhushanasati25@gmail.com> Co-authored-by: Megha Agarwal <19240983+meghaagr13@users.noreply.github.com>
aoshen02 committed
Sep 15, 2026 at 09:17 UTC
84e5f3edac29d4b8cf7266d7882d7cb5f8dd0245
4 files changed
+83
-10
tests/test_request_input_bounds.py
+30
@@ -12,6 +12,7 @@ import subprocess
12
import sys
13
from collections.abc import Callable
14
from typing import Protocol
15
+from unittest.mock import Mock
16
17
import pytest
18
from pydantic import ValidationError
@@ -27,10 +28,39 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
28
from vllm.exceptions import VLLMValidationError
29
from vllm.sampling_params import BeamSearchParams
30
from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
31
+from vllm.v1.engine.input_processor import InputProcessor
32
33
pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
34
35
36
+@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127])
37
+def test_routed_experts_offset_validated_before_engine_submission(offset):
38
+ """Reject invalid offsets even when routed-expert output is disabled."""
39
+ processor = Mock(spec=InputProcessor)
40
+ processor.tokenizer = None
41
+ processor.generation_config_fields = {}
42
+ processor.renderer = Mock()
43
+ processor.renderer.get_eos_token_id.return_value = None
44
+ processor.vllm_config = Mock()
45
+ params = SamplingParams(routed_experts_prompt_start=offset)
46
+ prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]}
47
+
48
+ if not 0 <= offset <= 3:
49
+ with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"):
50
+ InputProcessor.process_inputs(
51
+ processor, "invalid", prompt, params, ("generate",)
52
+ )
53
+ params.routed_experts_prompt_start = 0
54
+
55
+ request = InputProcessor.process_inputs(
56
+ processor, "valid", prompt, params, ("generate",)
57
+ )
58
+ assert (
59
+ request.sampling_params.routed_experts_prompt_start
60
+ == params.routed_experts_prompt_start
61
+ )
62
+
63
+
64
# --- Stop strings: public requests cap the number of stop strings ---------
65
66
tests/v1/core/test_scheduler.py
+39
@@ -51,6 +51,7 @@ from vllm.v1.outputs import (
51
ECConnectorOutput,
52
KVConnectorOutput,
53
ModelRunnerOutput,
54
+ RoutedExpertsLists,
55
SamplingMaskLists,
56
make_empty_encoder_model_runner_output,
57
)
@@ -131,6 +132,44 @@ def test_add_requests():
132
assert len(scheduler.waiting) == i + 1
133
134
135
+def test_routed_experts_prompt_start_at_prompt_end():
136
+ """A prompt-end offset should return empty routing data, not crash."""
137
+ scheduler = create_scheduler()
138
+ (request,) = create_requests(num_requests=1)
139
+ request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens
140
+ scheduler.add_request(request)
141
+ scheduler_output = scheduler.schedule()
142
+ num_tokens = scheduler_output.num_scheduled_tokens[request.request_id]
143
+
144
+ manager = Mock()
145
+ manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8)
146
+ manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8)
147
+ scheduler.enable_return_routed_experts = True
148
+ scheduler.routed_experts_mgr = manager
149
+ scheduler._re_block_ids = {request.request_id: []}
150
+
151
+ outputs = scheduler.update_from_output(
152
+ scheduler_output,
153
+ ModelRunnerOutput(
154
+ req_ids=[request.request_id],
155
+ req_id_to_index={request.request_id: 0},
156
+ sampled_token_ids=[[0]],
157
+ logprobs=None,
158
+ prompt_logprobs_dict={},
159
+ pooler_output=[],
160
+ routed_experts=RoutedExpertsLists(
161
+ routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8),
162
+ slot_mapping=np.arange(num_tokens),
163
+ ),
164
+ ),
165
+ )
166
+
167
+ manager.get.assert_called_once_with(
168
+ [], request.num_prompt_tokens, token_start=request.num_prompt_tokens
169
+ )
170
+ assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1)
171
+
172
+
173
def test_finish_request():
174
scheduler = create_scheduler()
175
requests = create_requests(num_requests=10)
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -2101,7 +2101,7 @@ class Scheduler(SchedulerInterface):
2101
prompt_start = (
2102
request.sampling_params.routed_experts_prompt_start
2103
)
2104
- assert prompt_start < request.num_prompt_tokens
2104
+ assert prompt_start <= request.num_prompt_tokens
2105
else:
2106
prompt_start = 0
2107
routed_experts = self.routed_experts_mgr.get(
vllm/v1/engine/input_processor.py
+13
-9
@@ -354,12 +354,21 @@ class InputProcessor:
354
if isinstance(params, SamplingParams):
355
# TODO: can we avoid cloning here in multiproc case?
356
sampling_params = params.clone()
357
+ prompt_len = length_from_prompt_token_ids_or_embeds(
358
+ prompt_token_ids, prompt_embeds
359
+ )
360
+ if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len:
361
+ raise VLLMValidationError(
362
+ f"routed_experts_prompt_start must be between 0 and "
363
+ f"the prompt length ({prompt_len}), inclusive.",
364
+ parameter="routed_experts_prompt_start",
365
+ value=sampling_params.routed_experts_prompt_start,
366
+ )
367
# If unset max tokens, then generate up to the max_model_len.
368
if sampling_params.max_tokens is None:
359
- seq_len = length_from_prompt_token_ids_or_embeds(
360
- prompt_token_ids, prompt_embeds
369
+ sampling_params.max_tokens = (
370
+ self.model_config.max_model_len - prompt_len
371
)
362
- sampling_params.max_tokens = self.model_config.max_model_len - seq_len
372
373
sampling_params.update_from_generation_config(
374
self.generation_config_fields,
@@ -368,12 +377,7 @@ class InputProcessor:
377
if self.tokenizer is not None:
378
sampling_params.update_from_tokenizer(self.tokenizer)
379
if sampling_params.trace_decode_token_ids:
371
- self._normalize_trace_replay_params(
372
- sampling_params,
373
- length_from_prompt_token_ids_or_embeds(
374
- prompt_token_ids, prompt_embeds
375
- ),
376
- )
380
+ self._normalize_trace_replay_params(sampling_params, prompt_len)
381
else:
382
pooling_params = params.clone()
383