[Bugfix][MRV2] Fix FlashInfer attention and ReplaySSM warmups on MRV2 (#61013)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/model_executor/test_attention_warmup.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.model_executor.warmup import attention_warmup
+from vllm.model_executor.warmup.attention_warmup import mixed_batch_attention_warmup
+
+
+def _worker(*, use_v2: bool, max_num_batched_tokens: int, max_model_len: int = 4096):
+    calls = []
+    worker = SimpleNamespace(
+        model_runner=SimpleNamespace(_dummy_run=lambda **kw: calls.append(kw)),
+        use_v2_model_runner=use_v2,
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens),
+        model_config=SimpleNamespace(max_model_len=max_model_len),
+        execute_model=object(),
+        sample_tokens=object(),
+    )
+    return worker, calls
+
+
+@pytest.mark.parametrize(
+    ("max_num_batched_tokens", "max_model_len", "expected"),
+    [(8192, 4096, 16), (6, 4096, 6), (8192, 8, 8)],
+)
+def test_v1_runs_mixed_dummy_batch_within_limits(
+    max_num_batched_tokens, max_model_len, expected
+):
+    worker, calls = _worker(
+        use_v2=False,
+        max_num_batched_tokens=max_num_batched_tokens,
+        max_model_len=max_model_len,
+    )
+    mixed_batch_attention_warmup(worker)
+    assert calls == [
+        dict(
+            num_tokens=expected,
+            skip_eplb=True,
+            is_profile=True,
+            force_attention=True,
+            create_mixed_batch=True,
+        )
+    ]
+
+
+@pytest.mark.parametrize(("max_num_batched_tokens", "expected"), [(8192, 16), (6, 6)])
+@pytest.mark.parametrize("ran", [True, False])
+def test_v2_runs_mixed_prefill_decode_step(
+    monkeypatch, max_num_batched_tokens, expected, ran
+):
+    worker, dummy_calls = _worker(
+        use_v2=True, max_num_batched_tokens=max_num_batched_tokens
+    )
+    mixed_calls = []
+
+    def fake_mixed_warmup(runner, execute_model, sample_tokens, num_tokens):
+        mixed_calls.append((runner, execute_model, sample_tokens, num_tokens))
+        return ran
+
+    monkeypatch.setattr(
+        attention_warmup, "run_mixed_prefill_decode_warmup", fake_mixed_warmup
+    )
+    mixed_batch_attention_warmup(worker)
+
+    assert mixed_calls == [
+        (worker.model_runner, worker.execute_model, worker.sample_tokens, expected)
+    ]
+    # No dummy-run fallback when the mixed step is skipped.
+    assert dummy_calls == []

vllm/model_executor/warmup/attention_warmup.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING, cast
+
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+    from vllm.v1.worker.gpu_worker import Worker
+
+logger = init_logger(__name__)
+
+
+def mixed_batch_attention_warmup(worker: "Worker", num_tokens: int = 16) -> None:
+    """Run a mixed prefill+decode batch to warm up both attention paths."""
+    num_tokens = min(
+        num_tokens,
+        worker.scheduler_config.max_num_batched_tokens,
+        worker.model_config.max_model_len,
+    )
+    if not worker.use_v2_model_runner:
+        worker.model_runner._dummy_run(
+            num_tokens=num_tokens,
+            skip_eplb=True,
+            is_profile=True,
+            force_attention=True,
+            create_mixed_batch=True,
+        )
+        return
+
+    if not run_mixed_prefill_decode_warmup(
+        cast("GPUModelRunner", worker.model_runner),
+        worker.execute_model,
+        worker.sample_tokens,
+        num_tokens,
+    ):
+        logger.info_once(
+            "Skipping mixed prefill+decode attention warmup; attention kernels "
+            "will be loaded on the first request."
+        )

vllm/model_executor/warmup/kernel_warmup.py

@@ -15,6 +15,9 @@ import torch
 
 import vllm.envs as envs
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.attention_warmup import (
+    mixed_batch_attention_warmup,
+)
 from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
@@ -286,15 +289,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
         )
     ):
         logger.info_once("Warming up FlashInfer attention.")
-        # Warmup with mixed batch containing both prefill and decode tokens
-        # This is to warm up both prefill and decode attention kernels
-        worker.model_runner._dummy_run(
-            num_tokens=16,
-            skip_eplb=True,
-            is_profile=True,
-            force_attention=True,
-            create_mixed_batch=True,
-        )
+        mixed_batch_attention_warmup(worker)
 
 
 def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:

vllm/model_executor/warmup/minimax_m3_msa_warmup.py

@@ -4,6 +4,9 @@
 from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.attention_warmup import (
+    mixed_batch_attention_warmup,
+)
 from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
 from vllm.platforms import current_platform
 from vllm.tracing import instrument
@@ -34,10 +37,4 @@ def minimax_m3_msa_warmup(worker: "Worker") -> None:
     logger.info("Warming up MiniMax M3 MSA kernels.")
 
     # Cover sparse prefill through the normal model path.
-    worker.model_runner._dummy_run(
-        num_tokens=16,
-        skip_eplb=True,
-        is_profile=True,
-        force_attention=True,
-        create_mixed_batch=True,
-    )
+    mixed_batch_attention_warmup(worker)

vllm/v1/worker/gpu/model_runner.py

@@ -857,6 +857,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 - num_reqs * self.speculator.num_extra_query_per_req,
             )
 
+        assert num_tokens <= self.max_num_tokens, (
+            f"Dummy run with {num_tokens} tokens exceeds max_num_batched_tokens "
+            f"({self.max_num_tokens})"
+        )
+
         # Distribute the remainder evenly so no dummy request exceeds
         # ceil(num_tokens / num_reqs) <= max_model_len tokens.
 
@@ -1898,9 +1903,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             )
             assert block_tables is not None
             attn_groups = self.attn_groups
-            if dummy_run and is_profile:
+            if dummy_run and is_profile and not valid_dummy_state_slots:
                 # Mamba layers take a cheap warmup path with no metadata;
                 # attention metadata is still built so those kernels tune.
+                # valid_dummy_state_slots runs the real Mamba kernels on distinct
+                # state slots; the caller must restore that state afterwards.
                 attn_groups = [
                     [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)]
                     for groups in attn_groups