[Perf][MiniMax-M3] Tune SM100 sparse decode

tests/kernels/attention/test_minimax_m3.py

@@ -2584,7 +2584,7 @@ def test_amd_decode_fused_topk_emits_sparse_table(
 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
 @pytest.mark.parametrize(
     "seq_lens_list",
-    [(130, 257), (129, 200, 384)],
+    [(130, 257), (129, 200, 384), (4096, 4096, 4096)],
 )
 @pytest.mark.parametrize("decode_query_len", [1, 4])
 @pytest.mark.parametrize("num_padded_reqs", [0, 2])

tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py

@@ -187,3 +187,104 @@ def test_minimax_m3_sparse_decode_fp8_kv_scales(scale_mode: str):
 
     torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
     assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)
+
+
+@pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
+@pytest.mark.parametrize(
+    "fp16_dot,batch_invariant", [(False, False), (True, False), (True, True)]
+)
+@torch.inference_mode()
+def test_minimax_m3_sparse_decode_fp8_split_scales(
+    scale_mode: str, fp16_dot: bool, batch_invariant: bool, monkeypatch
+):
+    if not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ):
+        pytest.skip("Sparse decode tuning requires SM100-family GPUs")
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", str(int(batch_invariant)))
+    total_q, num_blocks = 40, 16
+    torch.manual_seed(2)
+    q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
+    kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=3)
+    k_scale = torch.tensor(0.3, dtype=torch.float32, device=DEVICE)
+    v_scale = torch.tensor(0.7, dtype=torch.float32, device=DEVICE)
+    kv_dequant[..., :HEAD_DIM] = (kv_fp8[..., :HEAD_DIM].float() * k_scale).to(DTYPE)
+    kv_dequant[..., HEAD_DIM:] = (kv_fp8[..., HEAD_DIM:].float() * v_scale).to(DTYPE)
+    if scale_mode == "per_token_head":
+        shape = (NUM_KV_HEADS, num_blocks * SPARSE_BLOCK_SIZE)
+        k_scale = k_scale.expand(shape).contiguous()
+        v_scale = v_scale.expand(shape).contiguous()
+
+    blocks = torch.arange(num_blocks, dtype=torch.int32, device=DEVICE)
+    topk = blocks.expand(NUM_KV_HEADS, total_q, -1).contiguous()
+    block_table = blocks.expand(total_q, -1).contiguous()
+    seq_lens = torch.full(
+        (total_q,), num_blocks * SPARSE_BLOCK_SIZE, dtype=torch.int32, device=DEVICE
+    )
+    seq_lens[0] -= SPARSE_BLOCK_SIZE // 2
+    got, ref, unscaled = (torch.empty_like(q) for _ in range(3))
+    for output, cache, ks, vs in (
+        (got, kv_fp8, k_scale, v_scale),
+        (ref, kv_dequant, None, None),
+        (unscaled, kv_fp8, None, None),
+    ):
+        minimax_m3_sparse_attn_decode(
+            q,
+            cache,
+            topk,
+            block_table,
+            seq_lens,
+            NUM_KV_HEADS,
+            HEAD_DIM**-0.5,
+            output,
+            decode_query_len=1,
+            k_scale=ks,
+            v_scale=vs,
+            fp16_dot=fp16_dot,
+        )
+
+    torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
+    assert (got.float() - ref.float()).abs().mean() < (
+        unscaled.float() - ref.float()
+    ).abs().mean() * 0.1
+
+
+@pytest.mark.parametrize(
+    "cache_dtype,scale", [(torch.float8_e4m3fn, 200.0), (torch.float8_e5m2, 2.0)]
+)
+@pytest.mark.parametrize("fp16_dot", [False, True])
+@torch.inference_mode()
+def test_minimax_m3_sparse_decode_fp16_saturates_scaled_kv(
+    cache_dtype: torch.dtype, scale: float, fp16_dot: bool, monkeypatch
+):
+    if not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ):
+        pytest.skip("Sparse decode tuning requires SM100-family GPUs")
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
+    q = torch.full((1, NUM_HEADS, HEAD_DIM), 0.001, dtype=DTYPE, device=DEVICE)
+    cache_max = torch.finfo(cache_dtype).max
+    kv = torch.full(
+        (1, NUM_KV_HEADS, SPARSE_BLOCK_SIZE, 2 * HEAD_DIM),
+        cache_max,
+        dtype=DTYPE,
+        device=DEVICE,
+    ).to(cache_dtype)
+    scale_tensor = torch.tensor(scale, dtype=torch.float32, device=DEVICE)
+    output = torch.empty_like(q)
+    minimax_m3_sparse_attn_decode(
+        q,
+        kv,
+        torch.zeros((NUM_KV_HEADS, 1, 1), dtype=torch.int32, device=DEVICE),
+        torch.zeros((1, 1), dtype=torch.int32, device=DEVICE),
+        torch.tensor([SPARSE_BLOCK_SIZE], dtype=torch.int32, device=DEVICE),
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        output,
+        decode_query_len=1,
+        k_scale=scale_tensor,
+        v_scale=scale_tensor,
+        fp16_dot=fp16_dot,
+    )
+    expected = min(cache_max * scale, 65504.0) if fp16_dot else cache_max * scale
+    torch.testing.assert_close(output, torch.full_like(output, expected))

vllm/config/attention.py

@@ -43,6 +43,11 @@ class AttentionConfig:
     minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton"
     """Sparse decode kernel used by the MiniMax M3 MSA backend."""
 
+    minimax_m3_msa_decode_fp16_dot: bool = False
+    """Use FP16 dot operands with FP32 accumulation for SM100 Triton sparse
+    decode with BF16 queries and FP8 KV. Requires absent or scalar scales
+    and is disabled in batch-invariant mode."""
+
     backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict)
     """Per-KV-cache-group attention backend overrides, keyed by
     `KVCacheSpecKind` (e.g. `{"mla_attention": "FLASHINFER_MLA",

vllm/models/minimax_m3/common/ops/sparse_attn.py

@@ -19,8 +19,11 @@
 
 import torch
 
+import vllm.envs as envs
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+from vllm.utils.platform_utils import num_compute_units
 
 # One sparse block == one KV page.
 SPARSE_BLOCK_SIZE = 128
@@ -282,6 +285,7 @@ def _gqa_sparse_decode_kernel(
     USE_FP8: tl.constexpr,  # fp8 KV cache: dequantize K/V to q.dtype on load
     KV_SCALE_MODE: tl.constexpr,  # 0: none, 1: scalar, 2: [kv_head, token]
     USE_PDL: tl.constexpr,
+    USE_FP16_DOT: tl.constexpr,
 ):
     sm_scale_log2e = sm_scale * 1.4426950409
     # split-K over the topk dimension: pid(0) folds (query-token, chunk).
@@ -329,6 +333,8 @@ def _gqa_sparse_decode_kernel(
         mask=hd_mask,
         other=0.0,
     )
+    if USE_FP16_DOT:
+        q = tl.clamp(q.to(tl.float32), -65504.0, 65504.0).to(tl.float16)
 
     cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
     for _ in tl.range(chunk_start_topk, chunk_end_topk):
@@ -356,6 +362,8 @@ def _gqa_sparse_decode_kernel(
             k = k.to(q.dtype)
             if KV_SCALE_MODE == 1:
                 k = (k * tl.load(k_scale_ptr)).to(q.dtype)
+                if USE_FP16_DOT:
+                    k = tl.clamp(k, -65504.0, 65504.0).to(q.dtype)
             elif KV_SCALE_MODE == 2:
                 k_scale = tl.load(
                     k_scale_ptr
@@ -387,6 +395,8 @@ def _gqa_sparse_decode_kernel(
             v = v.to(q.dtype)
             if KV_SCALE_MODE == 1:
                 v = (v * tl.load(v_scale_ptr)).to(q.dtype)
+                if USE_FP16_DOT:
+                    v = tl.clamp(v, -65504.0, 65504.0).to(q.dtype)
             elif KV_SCALE_MODE == 2:
                 v_scale = tl.load(
                     v_scale_ptr
@@ -457,7 +467,7 @@ def _merge_topk_attn_out_kernel(
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
 
-    off_c = tl.arange(0, NUM_TOPK_CHUNKS)
+    off_c = tl.arange(0, triton.next_power_of_2(NUM_TOPK_CHUNKS))
     off_d = tl.arange(0, BLOCK_SIZE_D)
     d_mask = off_d < head_dim
     o_ptrs = (
@@ -468,8 +478,10 @@ def _merge_topk_attn_out_kernel(
         + off_d[None, :] * stride_o_d
     )
     lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
-    o = tl.load(o_ptrs, mask=d_mask[None, :], other=0.0)
-    lse = tl.load(lse_ptrs)  # empty chunks contribute -inf -> weight 0
+    o = tl.load(
+        o_ptrs, mask=(off_c < NUM_TOPK_CHUNKS)[:, None] & d_mask[None, :], other=0.0
+    )
+    lse = tl.load(lse_ptrs, mask=off_c < NUM_TOPK_CHUNKS, other=float("-inf"))
     lse_max = tl.max(lse, axis=0)
     weights = tl.exp2(lse - lse_max)
     weights = weights / tl.sum(weights, axis=0)
@@ -625,6 +637,7 @@ def minimax_m3_sparse_attn_decode(
     decode_query_len: int,
     k_scale: torch.Tensor | None = None,
     v_scale: torch.Tensor | None = None,
+    fp16_dot: bool = False,
 ) -> None:
     """GQA block-sparse attention for decode (split-K over the top-k blocks)."""
     total_q, num_heads, head_dim = q.shape
@@ -659,10 +672,28 @@ def minimax_m3_sparse_attn_decode(
     # launch_pdl was specified but unrecognised"). Only pass it when PDL is
     # actually supported -- on ROCm use_pdl is always False, so it's omitted.
     pdl_launch = {"launch_pdl": True} if use_pdl else {}
+    tune_decode = (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and not envs.VLLM_BATCH_INVARIANT
+    )
+    fp16_dot = (
+        tune_decode
+        and q.dtype == torch.bfloat16
+        and use_fp8
+        and kv_scale_mode != _KV_SCALE_PER_TOKEN_HEAD
+        and fp16_dot
+    )
     # split-K over the selected blocks; chunk count is shape-constant (cuda graph).
-    TARGET_GRID = 256
-    target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads)))
-    num_topk_chunks = 1 << (target.bit_length() - 1)
+    if tune_decode:
+        # Target two CTAs per SM without rounding the split count to a power of two.
+        resident = 2 * num_compute_units(q.device.index or 0)
+        target = max(1, min(max_topk, resident // max(1, total_q * num_kv_heads)))
+        num_topk_chunks = cdiv(max_topk, cdiv(max_topk, target))
+    else:
+        TARGET_GRID = 256
+        target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads)))
+        num_topk_chunks = 1 << (target.bit_length() - 1)
     o_partial = torch.empty(
         num_topk_chunks, total_q, num_heads, head_dim, dtype=q.dtype, device=q.device
     )
@@ -713,6 +744,7 @@ def minimax_m3_sparse_attn_decode(
         USE_FP8=use_fp8,
         KV_SCALE_MODE=kv_scale_mode,
         USE_PDL=use_pdl,
+        USE_FP16_DOT=fp16_dot,
         **pdl_launch,
     )
     merge_grid = (total_q, num_heads)

vllm/models/minimax_m3/common/sparse_attention.py

@@ -405,10 +405,12 @@ def __init__(
         topk_blocks: int,
         sparse_block_size: int,
         msa_decode_backend: str = "triton",
+        msa_decode_fp16_dot: bool = False,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size
         self.scale = scale
+        self.fp16_dot = msa_decode_fp16_dot
         self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
         self.kv_cache_dtype = kv_cache_dtype
         self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
@@ -500,6 +502,7 @@ def forward(
                 d.decode_query_len,
                 k_scale=k_scale,
                 v_scale=v_scale,
+                fp16_dot=self.fp16_dot,
             )
 
         # Prefill [nd:]: cu_seqlens_q already rebased to 0.

vllm/models/minimax_m3/nvidia/model.py

@@ -531,6 +531,9 @@ def __init__(
             msa_decode_backend=(
                 vllm_config.attention_config.minimax_m3_msa_decode_backend
             ),
+            msa_decode_fp16_dot=(
+                vllm_config.attention_config.minimax_m3_msa_decode_fp16_dot
+            ),
         )
         # Self-contained nn.Module: owns its side cache, selects its impl in init.
         self.indexer = MiniMaxM3Indexer(

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@@ -204,6 +204,7 @@ def __init__(
         topk_blocks: int,
         sparse_block_size: int,
         msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton",
+        msa_decode_fp16_dot: bool = False,
     ) -> None:
         super().__init__(
             num_heads,
@@ -213,6 +214,7 @@ def __init__(
             kv_cache_dtype,
             topk_blocks=topk_blocks,
             sparse_block_size=sparse_block_size,
+            msa_decode_fp16_dot=msa_decode_fp16_dot,
         )
         self.use_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype)
         self.use_cutlass_decode = supports_cutlass_sparse_decode(
@@ -332,6 +334,7 @@ def forward(
                     d.decode_query_len,
                     k_scale=k_scale,
                     v_scale=v_scale,
+                    fp16_dot=self.fp16_dot,
                 )
 
         # Prefill [nd:]: MSA sparse FMHA over the selected blocks.