[BugFix] Gate MRV2 mixed sparse-MLA warmup on `max_num_seqs` > 1 (#47050)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: ziminghuang <ziminghuang@inferact.ai>
Nick Hill committed
Jun 30, 2026 at 16:31 UTC
fcaa84efa7a980cee3681976ecfe36133c48251a
3 files changed
+34
-4
tests/v1/worker/test_mixed_warmup_gate.py
new
+30
@@ -0,0 +1,30 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+
9
+from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
10
+
11
+
12
+def _fail(*args, **kwargs):
13
+ raise AssertionError("worker callback must not run when warmup is skipped")
14
+
15
+
16
+@pytest.mark.parametrize("max_num_reqs", [1, 0])
17
+def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
18
+ """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2
19
+ the warmup must be skipped without touching the worker callbacks."""
20
+ runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs)
21
+
22
+ assert (
23
+ run_mixed_prefill_decode_warmup(
24
+ runner,
25
+ worker_execute_model=_fail,
26
+ worker_sample_tokens=_fail,
27
+ num_tokens=128,
28
+ )
29
+ is False
30
+ )
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+3
-3
@@ -130,7 +130,7 @@ def _run_flashinfer_sparse_mla_decode_autotune(
130
with torch.inference_mode():
131
warmup_executed = True
132
if is_leader:
133
- if _uses_v2_model_runner(runner):
133
+ if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
134
v2_runner = cast("V2GPUModelRunner", runner)
135
warmup_executed = run_mixed_prefill_decode_warmup(
136
v2_runner,
@@ -144,7 +144,7 @@ def _run_flashinfer_sparse_mla_decode_autotune(
144
with flashinfer_autotune(True, cache=str(cache_path)):
145
runner._dummy_run(**dummy_run_kwargs)
146
else:
147
- if _uses_v2_model_runner(runner):
147
+ if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
148
v2_runner = cast("V2GPUModelRunner", runner)
149
warmup_executed = run_mixed_prefill_decode_warmup(
150
v2_runner,
@@ -236,7 +236,7 @@ def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None:
236
)
237
mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens)
238
if not mixed_warmup_done:
239
- if _uses_v2_model_runner(runner):
239
+ if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
240
v2_runner = cast("V2GPUModelRunner", runner)
241
run_mixed_prefill_decode_warmup(
242
v2_runner,
vllm/v1/worker/gpu/warmup.py
+1
-1
@@ -34,7 +34,7 @@ def run_mixed_prefill_decode_warmup(
34
req_id_prefix: str = "_v2_mixed_warmup",
35
) -> bool:
36
"""Run a V2 mixed prefill+decode step through normal scheduler inputs."""
37
- if model_runner.is_pooling_model or num_tokens < 3:
37
+ if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3:
38
return False
39
40
decode_req_id = f"{req_id_prefix}_decode_"