[Bugfix][MiniMax-M3] Keep FP8 query allocation stable across CUDA graph replay (#51203)

Signed-off-by: Kyle Liang <kylliang@nvidia.com>

tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py

@@ -9,12 +9,14 @@ import pytest
 import torch
 
 from vllm import _custom_ops as ops
-from vllm.config import AttentionConfig
+from vllm.config import AttentionConfig, CUDAGraphMode
+from vllm.forward_context import ForwardContext, override_forward_context
 from vllm.models.minimax_m3.common.ops.sparse_attn import (
     minimax_m3_sparse_attn_decode,
 )
 from vllm.models.minimax_m3.common.sparse_attention import (
     MiniMaxM3SparseBackend,
+    MiniMaxM3SparseMetadata,
     MiniMaxM3SparseMetadataBuilder,
     MiniMaxM3SparseTritonImpl,
     select_main_backend_and_impl_cls,
@@ -22,6 +24,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
 from vllm.models.minimax_m3.nvidia import (
     sparse_attention_msa as sparse_attention_msa_module,
 )
+from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
 from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
     MSACutlassDecodePlanCache,
     msa_cutlass_sparse_decode,
@@ -316,6 +319,83 @@ def test_msa_cutlass_plan_cache_keys_query_len(
     assert built_query_lens == [1, 2]
 
 
+def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import vllm.envs as envs
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
+    envs.disable_envs_cache()
+
+    num_tokens = 16
+    layer_name = "model.layers.0.self_attn.attn"
+    seq_lens = torch.full((num_tokens,), 257, dtype=torch.int32, device="cuda")
+    block_table = torch.zeros(num_tokens, 3, dtype=torch.int32, device="cuda")
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
+    decode = MiniMaxM3SparseMSADecodeMetadata(
+        seq_lens=seq_lens,
+        block_table=block_table,
+        decode_query_len=1,
+        msa_cutlass=None,
+    )
+    metadata = MiniMaxM3SparseMetadata(
+        seq_lens=seq_lens,
+        max_seq_len=257,
+        slot_mapping=slot_mapping,
+        num_actual_tokens=num_tokens,
+        num_decodes=num_tokens,
+        num_decode_tokens=num_tokens,
+        num_prefills=0,
+        num_prefill_tokens=0,
+        decode=decode,
+    )
+    forward_context = ForwardContext(
+        no_compile_layers={},
+        attn_metadata={layer_name: metadata},
+        slot_mapping={},
+        cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+    )
+
+    impl = object.__new__(MiniMaxM3SparseMSAImpl)
+    impl.use_cutlass_decode = True
+    attention = SimpleNamespace(impl=impl, q_size=HEAD_DIM)
+    qkv = torch.empty(num_tokens, HEAD_DIM, device="cuda")
+    observed_ptrs: list[int] = []
+
+    @eager_break_during_capture
+    def run_attention(query_fp8: torch.Tensor | None) -> None:
+        if impl.should_use_msa_decode(layer_name):
+            assert query_fp8 is not None
+            observed_ptrs.append(query_fp8.data_ptr())
+
+    stream = torch.cuda.Stream()
+
+    with torch.cuda.stream(stream), override_forward_context(forward_context):
+        capture = BreakableCUDAGraphCapture()
+        with capture:
+            query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv)
+            if query_fp8 is not None:
+                query_fp8.zero_()
+            run_attention(query_fp8)
+            qkv.zero_()
+
+        assert capture.num_graphs == 2
+        assert capture.num_eager_breaks == 1
+        assert observed_ptrs == []
+
+        decode.msa_cutlass = object()  # type: ignore[assignment]
+        for _ in range(3):
+            capture.replay()
+        stream.synchronize()
+
+    assert len(observed_ptrs) == 3
+    assert len(set(observed_ptrs)) == 1
+
+
 def _make_topk(
     seq_lens: list[int],
     num_kv_heads: int,

vllm/models/minimax_m3/nvidia/model.py

@@ -564,6 +564,15 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
         )
 
+    def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None:
+        if not getattr(self.impl, "use_cutlass_decode", False):
+            return None
+        return torch.empty(
+            (qkv.shape[0], self.q_size),
+            dtype=torch.float8_e4m3fn,
+            device=qkv.device,
+        )
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -598,16 +607,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
         index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
         q = qkv.new_empty((num_tokens, self.q_size))
-        use_msa_decode = self.impl.should_use_msa_decode(self.layer_name)
-        query_fp8 = (
-            torch.empty(
-                (num_tokens, self.q_size),
-                dtype=torch.float8_e4m3fn,
-                device=qkv.device,
-            )
-            if use_msa_decode
-            else None
-        )
+        query_fp8 = self._allocate_query_fp8(qkv)
         # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
         # the fused kernel emits fp8 directly when this buffer is e4m3.
         index_q = qkv.new_empty(