Overview

With CUTLASS MSA sparse decode on (an NVFP4 KV cache, or an FP8 E4M3 cache with the CUTLASS_MSA decode backend), the fused QK-norm/RoPE/insert kernel writes MiniMax-M3's query only in E4M3. MiniMaxM3SparseMSAImpl then dequantized it back into the BF16 query buffer before every prefill. Both fmha_sm100 prefill kernels take an E4M3 query natively, so prefill now passes it through as is. The PR also bumps fmha_sm100 (vllm-project/MSA) from be4a52a to 968061a. The new pin includes MSA's Blackwell prefill port, which now serves an E4M3 query over vLLM's packed K|V pages (vllm-project/MSA#32).

Claims

• MSA prefill attention is 26–37% faster with an FP8 KV cache and 8–13% faster with an NVFP4 KV cache than on main (per layer on GB300; table below).
• Prefill no longer makes a BF16 copy of the query. Only the Triton decode fallback still reads the dequantized query.
• BF16 caches and the default Triton decode backend for FP8 caches still pass a BF16 query, so their path is unchanged.

Validation

Tests: on GB300 (SM103, CUDA 13.4, fmha_sm100 built from the new pin), pytest tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py tests/kernels/attention/test_minimax_m3.py tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py gives 240 passed, 45 skipped. All skips are ROCm-only tests. The MSA test file runs in CI on B200.

• New test_msa_prefill_reads_query_fp8_over_fp8_cache[tp1|tp4] covers FP8-cache prefill, including chunks over cached context, with CUTLASS decode on. It checks that sparse_atten_func receives the E4M3 query and that the output matches Triton on the dequantized query.
• test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache now checks that sparse_atten_nvfp4_kv_func receives the E4M3 query. Its reference is now the cache as an E4M3-query kernel reads it (see Details).
• Both prefill tests require all but 1e-4 of the elements within atol/rtol 0.05 and bound the rest at 0.25. The E4M3 probabilities are rounded per split, and splits follow the SM count, so a rare element lands just past 0.05: 1 of 4.1M on B200 CI. On GB300 over 16 seeds the worst case is 2.4e-7 of the elements and 0.07 max abs error.
• On the old code, all four prefill cases fail with [torch.bfloat16] == [torch.float8_e4m3fn].

Performance: MiniMaxM3SparseMSAImpl.forward on prefill-only batches (build_k2q_csr + attention + combine, plus the dequant where the query is BF16). GB300, one request, top-k 16, page 128. Each value is the median of 3×50 iterations after warmup; spread ≤1%, except ±5% on the first FP8 row of the main run.

Percentages are relative to main. With the new pin, a BF16 query over the FP8 cache also reaches the Blackwell port, but with BF16 MMAs. The E4M3 query gives the bulk of the FP8 gain at either pin: on main's pin alone it is 13–24% (411 → 356, 462 → 367, 801 → 623, 1423 → 1083, 1498 → 1162 µs).

"""Per-layer MSA prefill attention (impl.forward) with an E4M3 query."""
import math
from types import SimpleNamespace

import torch

import tests.kernels.attention.test_minimax_m3_msa_cutlass_sparse_decode as t
from vllm.forward_context import ForwardContext, override_forward_context
from vllm.models.minimax_m3.common.sparse_attention import (
    MiniMaxM3SparseMetadata, MiniMaxM3SparsePrefillMetadata)

BLK, TOPK = 128, 16


def case(q_len, seq_len, hkv, g):
    pos = torch.arange(seq_len - q_len, seq_len, device="cuda")
    visible = pos // BLK + 1
    vmax = int(visible.max())
    r = torch.rand(q_len, hkv, vmax, device="cuda", generator=g)
    ar = torch.arange(vmax, device="cuda")
    r = torch.where(ar < visible[:, None, None] - 1, r, -1.0)
    r = torch.where(ar == visible[:, None, None] - 1, 2.0, r)
    k = min(TOPK, vmax)
    val, idx = r.topk(k, dim=-1)
    topk = torch.where(val >= 0, idx, -1).int()
    if k < TOPK:
        topk = torch.cat([topk, topk.new_full((q_len, hkv, TOPK - k), -1)], -1)
    pages = math.ceil(seq_len / BLK)
    block_table = torch.randperm(pages, dtype=torch.int32, device="cuda")[None]
    seq = torch.tensor([seq_len], dtype=torch.int32, device="cuda")
    cu_q = torch.tensor([0, q_len], dtype=torch.int32, device="cuda")
    cu_k = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
    p = MiniMaxM3SparsePrefillMetadata(
        cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, seq_lens=seq, context_lens=seq - q_len,
        block_table=block_table, max_query_len=q_len, max_seq_len=seq_len,
        total_kv_blocks=pages)
    md = MiniMaxM3SparseMetadata(
        seq_lens=seq, max_seq_len=seq_len,
        slot_mapping=torch.empty(q_len, dtype=torch.int64, device="cuda"),
        num_actual_tokens=q_len, num_decodes=0, num_decode_tokens=0, num_prefills=1,
        num_prefill_tokens=q_len, prefill=p)
    return md, topk, pages


def bench(fn, iters=50):
    for _ in range(10):
        fn()
    torch.cuda.synchronize()
    s, e = torch.cuda.Event(True), torch.cuda.Event(True)
    s.record()
    for _ in range(iters):
        fn()
    e.record()
    torch.cuda.synchronize()
    return s.elapsed_time(e) / iters * 1e3


g = torch.Generator(device="cuda").manual_seed(0)
for kv_dtype in ("fp8", "nvfp4"):
    for hq, hkv, q_len, ctx in ((16, 1, 8192, 0), (16, 1, 8192, 32768),
                                (16, 1, 16384, 114688), (64, 4, 8192, 0),
                                (64, 4, 8192, 32768)):
        md, topk, pages = case(q_len, ctx + q_len, hkv, g)
        k_scale = torch.tensor(1.0, device="cuda")
        v_scale = torch.tensor(1.0, device="cuda")
        if kv_dtype == "nvfp4":
            kv_cache, _ = t._make_nvfp4_kv_cache(pages, hkv, "nvfp4", k_scale, v_scale)
        else:
            kv_cache = (torch.randn(pages, hkv, BLK, 256, device="cuda") * 0.5).to(torch.float8_e4m3fn)
        layer = SimpleNamespace(layer_name="l", topk_indices_buffer=topk, _q_scale_float=0.5,
                                _k_scale=k_scale, _v_scale=v_scale)
        qf = torch.randn(q_len, hq * 128, device="cuda").to(torch.float8_e4m3fn)
        query = torch.empty(q_len, hq * 128, dtype=torch.bfloat16, device="cuda")
        out = torch.empty_like(query)
        impl = t.MiniMaxM3SparseMSAImpl(
            num_heads=hq, head_size=128, scale=t.SM_SCALE, num_kv_heads=hkv,
            kv_cache_dtype=kv_dtype, topk_blocks=TOPK, sparse_block_size=BLK,
            msa_decode_backend="cutlass")
        assert impl.use_cutlass_decode
        with override_forward_context(ForwardContext(no_compile_layers={}, attn_metadata={"l": md}, slot_mapping={})):
            bench(lambda: impl.forward(layer, query, kv_cache, out, query_fp8=qf), iters=200)
            ts = [bench(lambda: impl.forward(layer, query, kv_cache, out, query_fp8=qf)) for _ in range(3)]
        print(f"{kv_dtype} H{hq}/{hkv} q{q_len} ctx{ctx}: " + " ".join(f"{x:.1f}" for x in ts) + " us")

No end-to-end accuracy eval or serving benchmark is included; validation is the kernel-level numerics and per-layer timings above.

Details

Kernel capability (MSA 968061a; same at be4a52a except the Blackwell port):

• sparse_atten_func accepts an FP8 E4M3 query over FP8 E4M3 K/V and runs FP8 QK/PV MMAs, on the Blackwell port when eligible. The BF16-Q × FP8-K/V mode stages K/V up to BF16 instead.
• sparse_atten_nvfp4_kv_func accepts a BF16 or E4M3 query. With E4M3 it dequantizes each NVFP4 code × block_scale to E4M3 instead of BF16, and applies the K global scale to the scores and the V global scale in the combine.

Scales: q_scale folds into softmax_scale. The FP8 cache is written with unit scales and the NVFP4 global scales are passed as before, so K/V scale handling is unchanged.

Numerics: the E4M3-query kernels round the softmax probabilities to E4M3. The NVFP4 kernel also rounds K/V to E4M3. On the NVFP4 test (TP1):

• Against the E4M3-staged cache: max abs error 0.06, and 0 of 4.1M elements outside atol/rtol 0.05.
• Against the exact BF16 dequantization: max 0.14, with 122 of 4.1M elements outside that tolerance. This is why the test's reference changed.

On the FP8 test the max abs error is ≈0.03.

Limitations: validated on SM103. SM100 is covered by the CI kernel job. The E4M3-query NVFP4 kernel is untested on SM107 here.

MSA pin bump (be4a52a → 968061a):

• The functions vLLM calls keep their signatures: fmha_sm100, fmha_sm100_plan, _fmha_sm100(_plan), sparse_topk_select, sparse_atten_func, sparse_atten_nvfp4_kv_func and nvfp4_head_slot_views.
• vLLM's decode plans still pin decode_backend="kv_mode3". The new prefill_backend plan option only routes sparse-prefill plans built through fmha_sm100_plan, which vLLM does not use.
• The cmake install list adds the files the package now imports: _jit_cache.py, msa_warmup.py, q8kv4_prefill_adapter.py and prefill_q8kv4/.

Duplicate check: no open PR touches the MSA prefill query or fmha_sm100.cmake. Open MiniMax-M3 MSA PRs (#55838, #56932, #59979) change decode or indexer metadata only.

AI assistance (Claude Code) was used to investigate, implement, test and benchmark this change.

🤖 Generated with Claude Code