[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>

tests/v1/worker/test_mixed_warmup_gate.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
+
+
+def _fail(*args, **kwargs):
+    raise AssertionError("worker callback must not run when warmup is skipped")
+
+
+@pytest.mark.parametrize("max_num_reqs", [1, 0])
+def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
+    """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2
+    the warmup must be skipped without touching the worker callbacks."""
+    runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs)
+
+    assert (
+        run_mixed_prefill_decode_warmup(
+            runner,
+            worker_execute_model=_fail,
+            worker_sample_tokens=_fail,
+            num_tokens=128,
+        )
+        is False
+    )

vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py

@@ -130,7 +130,7 @@ def _run_flashinfer_sparse_mla_decode_autotune(
     with torch.inference_mode():
         warmup_executed = True
         if is_leader:
-            if _uses_v2_model_runner(runner):
+            if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
                 v2_runner = cast("V2GPUModelRunner", runner)
                 warmup_executed = run_mixed_prefill_decode_warmup(
                     v2_runner,
@@ -144,7 +144,7 @@ def _run_flashinfer_sparse_mla_decode_autotune(
                 with flashinfer_autotune(True, cache=str(cache_path)):
                     runner._dummy_run(**dummy_run_kwargs)
         else:
-            if _uses_v2_model_runner(runner):
+            if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
                 v2_runner = cast("V2GPUModelRunner", runner)
                 warmup_executed = run_mixed_prefill_decode_warmup(
                     v2_runner,
@@ -236,7 +236,7 @@ def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None:
     )
     mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens)
     if not mixed_warmup_done:
-        if _uses_v2_model_runner(runner):
+        if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
             v2_runner = cast("V2GPUModelRunner", runner)
             run_mixed_prefill_decode_warmup(
                 v2_runner,

vllm/v1/worker/gpu/warmup.py

@@ -34,7 +34,7 @@ def run_mixed_prefill_decode_warmup(
     req_id_prefix: str = "_v2_mixed_warmup",
 ) -> bool:
     """Run a V2 mixed prefill+decode step through normal scheduler inputs."""
-    if model_runner.is_pooling_model or num_tokens < 3:
+    if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3:
         return False
 
     decode_req_id = f"{req_id_prefix}_decode_"