From a450316f26fcff8f5715e51c444cd272b9cad12a Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Fri, 16 Oct 2026 14:07:58 -0700 Subject: [PATCH 1/2] [Performance][MiniMax-M3] Tune SM100 sparse decode Tune KV-block splitting for SM100 and mask padded merge lanes. Keep BF16 dot operands by default and expose optional FP16 operands through AttentionConfig. Preserve batch-invariant behavior and existing fallback operand math. Derived from xinli-sw/vllm commits 97e0532d95bf0863376b0e4d57a23f3d9cc2c597 and 5c353c3c0a3e7b0bf8134a162dada99ebd648601. Co-authored-by: Xin Li Co-authored-by: Suyog Gupta Co-authored-by: Codex Signed-off-by: Hanlin Bi --- tests/kernels/attention/test_minimax_m3.py | 2 +- .../test_minimax_m3_sparse_attn_fp8_scale.py | 101 ++++++++++++++++++ vllm/config/attention.py | 5 + .../minimax_m3/common/ops/sparse_attn.py | 44 ++++++-- .../minimax_m3/common/sparse_attention.py | 3 + vllm/models/minimax_m3/nvidia/model.py | 3 + .../minimax_m3/nvidia/sparse_attention_msa.py | 3 + 7 files changed, 154 insertions(+), 7 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 918e4562cea3..f5e10faebfd0 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/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]) diff --git a/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py b/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py index 9cd481b5e336..a0587d76af95 100644 --- a/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py +++ b/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)) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index fc3ecb0dfc53..2d5aa45b0e1f 100644 --- a/vllm/config/attention.py +++ b/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", diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 6426115ae356..186758382288 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/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, + 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 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 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 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, + FP16_DOT=fp16_dot, **pdl_launch, ) merge_grid = (total_q, num_heads) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 1346765f84cb..f39e8048a119 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/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. diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index d67e9ea19d90..85a0b8840227 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/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( diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index a5598c719bfa..2c2ecbf919d5 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/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. From 4df6ebd262af62f984626f7b9bb5c5179b7def28 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Fri, 16 Oct 2026 15:45:28 -0700 Subject: [PATCH 2/2] Rename MiniMax sparse decode FP16 constexpr Co-authored-by: Codex Signed-off-by: Hanlin Bi --- vllm/models/minimax_m3/common/ops/sparse_attn.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 186758382288..fd81176e5a65 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -285,7 +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, - FP16_DOT: 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). @@ -333,7 +333,7 @@ def _gqa_sparse_decode_kernel( mask=hd_mask, other=0.0, ) - if FP16_DOT: + 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 @@ -362,7 +362,7 @@ 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 FP16_DOT: + if USE_FP16_DOT: k = tl.clamp(k, -65504.0, 65504.0).to(q.dtype) elif KV_SCALE_MODE == 2: k_scale = tl.load( @@ -395,7 +395,7 @@ 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 FP16_DOT: + if USE_FP16_DOT: v = tl.clamp(v, -65504.0, 65504.0).to(q.dtype) elif KV_SCALE_MODE == 2: v_scale = tl.load( @@ -744,7 +744,7 @@ def minimax_m3_sparse_attn_decode( USE_FP8=use_fp8, KV_SCALE_MODE=kv_scale_mode, USE_PDL=use_pdl, - FP16_DOT=fp16_dot, + USE_FP16_DOT=fp16_dot, **pdl_launch, ) merge_grid = (total_q, num_heads)