[Bugfix][MiniMax-M3] Keep FP8 query allocation stable across CUDA graph replay (#51203)
Signed-off-by: Kyle Liang <kylliang@nvidia.com>
kyleliang-nv committed
Aug 20, 2026 at 10:45 UTC
7c8b68b9ce30c45ea17063cc040b2473e152fa7d
2 files changed
+91
-11
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+81
-1
@@ -9,12 +9,14 @@ import pytest
9
import torch
10
11
from vllm import _custom_ops as ops
12
-from vllm.config import AttentionConfig
12
+from vllm.config import AttentionConfig, CUDAGraphMode
13
+from vllm.forward_context import ForwardContext, override_forward_context
14
from vllm.models.minimax_m3.common.ops.sparse_attn import (
15
minimax_m3_sparse_attn_decode,
16
)
17
from vllm.models.minimax_m3.common.sparse_attention import (
18
MiniMaxM3SparseBackend,
19
+ MiniMaxM3SparseMetadata,
20
MiniMaxM3SparseMetadataBuilder,
21
MiniMaxM3SparseTritonImpl,
22
select_main_backend_and_impl_cls,
@@ -22,6 +24,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
24
from vllm.models.minimax_m3.nvidia import (
25
sparse_attention_msa as sparse_attention_msa_module,
26
)
27
+from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
28
from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
29
MSACutlassDecodePlanCache,
30
msa_cutlass_sparse_decode,
@@ -316,6 +319,83 @@ def test_msa_cutlass_plan_cache_keys_query_len(
319
assert built_query_lens == [1, 2]
320
321
322
+def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay(
323
+ monkeypatch: pytest.MonkeyPatch,
324
+) -> None:
325
+ import vllm.envs as envs
326
+ from vllm.compilation.breakable_cudagraph import (
327
+ BreakableCUDAGraphCapture,
328
+ eager_break_during_capture,
329
+ )
330
+
331
+ monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
332
+ envs.disable_envs_cache()
333
+
334
+ num_tokens = 16
335
+ layer_name = "model.layers.0.self_attn.attn"
336
+ seq_lens = torch.full((num_tokens,), 257, dtype=torch.int32, device="cuda")
337
+ block_table = torch.zeros(num_tokens, 3, dtype=torch.int32, device="cuda")
338
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
339
+ decode = MiniMaxM3SparseMSADecodeMetadata(
340
+ seq_lens=seq_lens,
341
+ block_table=block_table,
342
+ decode_query_len=1,
343
+ msa_cutlass=None,
344
+ )
345
+ metadata = MiniMaxM3SparseMetadata(
346
+ seq_lens=seq_lens,
347
+ max_seq_len=257,
348
+ slot_mapping=slot_mapping,
349
+ num_actual_tokens=num_tokens,
350
+ num_decodes=num_tokens,
351
+ num_decode_tokens=num_tokens,
352
+ num_prefills=0,
353
+ num_prefill_tokens=0,
354
+ decode=decode,
355
+ )
356
+ forward_context = ForwardContext(
357
+ no_compile_layers={},
358
+ attn_metadata={layer_name: metadata},
359
+ slot_mapping={},
360
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
361
+ )
362
+
363
+ impl = object.__new__(MiniMaxM3SparseMSAImpl)
364
+ impl.use_cutlass_decode = True
365
+ attention = SimpleNamespace(impl=impl, q_size=HEAD_DIM)
366
+ qkv = torch.empty(num_tokens, HEAD_DIM, device="cuda")
367
+ observed_ptrs: list[int] = []
368
+
369
+ @eager_break_during_capture
370
+ def run_attention(query_fp8: torch.Tensor | None) -> None:
371
+ if impl.should_use_msa_decode(layer_name):
372
+ assert query_fp8 is not None
373
+ observed_ptrs.append(query_fp8.data_ptr())
374
+
375
+ stream = torch.cuda.Stream()
376
+
377
+ with torch.cuda.stream(stream), override_forward_context(forward_context):
378
+ capture = BreakableCUDAGraphCapture()
379
+ with capture:
380
+ query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv)
381
+ if query_fp8 is not None:
382
+ query_fp8.zero_()
383
+ run_attention(query_fp8)
384
+ qkv.zero_()
385
+
386
+ assert capture.num_graphs == 2
387
+ assert capture.num_eager_breaks == 1
388
+ assert observed_ptrs == []
389
+
390
+ decode.msa_cutlass = object() # type: ignore[assignment]
391
+ for _ in range(3):
392
+ capture.replay()
393
+ stream.synchronize()
394
+
395
+ assert len(observed_ptrs) == 3
396
+ assert len(set(observed_ptrs)) == 1
397
+
398
+
399
def _make_topk(
400
seq_lens: list[int],
401
num_kv_heads: int,
vllm/models/minimax_m3/nvidia/model.py
+10
-10
@@ -564,6 +564,15 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
564
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
565
)
566
567
+ def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None:
568
+ if not getattr(self.impl, "use_cutlass_decode", False):
569
+ return None
570
+ return torch.empty(
571
+ (qkv.shape[0], self.q_size),
572
+ dtype=torch.float8_e4m3fn,
573
+ device=qkv.device,
574
+ )
575
+
576
def forward(
577
self,
578
positions: torch.Tensor,
@@ -598,16 +607,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
607
main_slot_mapping = fwd_slot_mapping[self.layer_name]
608
index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
609
q = qkv.new_empty((num_tokens, self.q_size))
601
- use_msa_decode = self.impl.should_use_msa_decode(self.layer_name)
602
- query_fp8 = (
603
- torch.empty(
604
- (num_tokens, self.q_size),
605
- dtype=torch.float8_e4m3fn,
606
- device=qkv.device,
607
- )
608
- if use_msa_decode
609
- else None
610
- )
610
+ query_fp8 = self._allocate_query_fp8(qkv)
611
# index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
612
# the fused kernel emits fp8 directly when this buffer is e4m3.
613
index_q = qkv.new_empty(