diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index e2b59dc334..f6ea80cca0 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/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, diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index a75eca2b63..99e6d1f7ae 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/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(