[MiniMax-M3 Performance] Prefer CuTe for MiniMax M3 FP8 indexer prefill

.buildkite/test_areas/kernels.yaml

@@ -562,6 +562,7 @@ steps:
   num_devices: 1
   source_file_dependencies:
   - vllm/models/minimax_m3/
+  - vllm/config/attention.py
   - vllm/config/cache.py
   - vllm/v1/kv_cache_interface.py
   - tests/v1/attention/utils.py

tests/kernels/attention/test_minimax_m3.py

@@ -822,20 +822,35 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
     _assert_topk_indices_equal_unordered(actual, expected)
 
 
-# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha/CuteDSL score
-# + unified top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
-# builders on the SAME CommonAttentionMetadata + index cache, and assert the
-# selected blocks agree. This exercises all the metadata the impl/kernels consume
-# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
-# decode_pages split) -- a metadata bug on either side shifts the causal window
-# or the block->page mapping and breaks the comparison.
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(100),
     reason="fmha_sm100 indexer requires SM100 (Blackwell).",
 )
-@pytest.mark.parametrize("topk", [16])
-@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
-def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
+@pytest.mark.parametrize(
+    (
+        "prefill_backend",
+        "index_dtype",
+        "num_idx_heads",
+        "init_blocks",
+        "local_blocks",
+        "expect_cute",
+    ),
+    [
+        pytest.param(None, torch.float8_e4m3fn, 2, 1, 2, True, id="default-cute"),
+        pytest.param("auto", torch.float8_e4m3fn, 4, 2, 0, False, id="local0-fallback"),
+        pytest.param("auto", torch.bfloat16, 4, 0, 1, False, id="bf16-fallback"),
+    ],
+)
+def test_msa_indexer_impl_matches_triton(
+    prefill_backend,
+    index_dtype,
+    num_idx_heads,
+    init_blocks,
+    local_blocks,
+    expect_cute,
+    monkeypatch,
+):
+    """Real builders and native Top16 must preserve ragged, padded batch selection."""
     import vllm.models.minimax_m3.common.indexer as indexer_mod
     from tests.v1.attention.utils import (
         BatchSpec,
@@ -850,25 +865,40 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
     )
     from vllm.models.minimax_m3.nvidia.indexer_msa import (
         MiniMaxM3IndexerMSAImpl,
+        MiniMaxM3IndexerMSAMetadata,
         MiniMaxM3IndexerMSAMetadataBuilder,
     )
 
+    if expect_cute:
+        from vllm.third_party import fmha_sm100
+
+        wrapper = getattr(fmha_sm100, "BatchPrefillIndexerQ8KV8Wrapper", None)
+        if not hasattr(wrapper, "run_scores"):
+            pytest.skip("Installed MSA lacks the prefill run_scores API.")
+
     torch.manual_seed(0)
     device = torch.device("cuda")
-    num_idx_heads, head_dim = 4, HEAD_DIM
+    head_dim, topk = HEAD_DIM, TOPK
     # TP=1: avoid requiring an initialized distributed group in a unit test.
     monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
 
     vllm_config = create_vllm_config(
-        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
+        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=512
     )
+    if prefill_backend is not None:
+        vllm_config.attention_config.minimax_m3_indexer_prefill_backend = (
+            prefill_backend
+        )
     vllm_config.model_config.hf_config.sparse_attention_config = {
-        "sparse_num_index_heads": num_idx_heads
+        "sparse_num_index_heads": num_idx_heads,
+        "sparse_init_block": init_blocks,
+        "sparse_local_block": local_blocks,
     }
 
-    # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
-    # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
-    batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
+    # Decode-first batch; fresh and cached prefills cross page and Q-tile tails.
+    batch = BatchSpec(
+        seq_lens=[2305, 2561, 129, 2831, 4097], query_lens=[1, 1, 129, 255, 31]
+    )
     common = create_common_attn_metadata(
         batch, BLOCK_SIZE, device, arange_block_indices=True
     )
@@ -882,22 +912,31 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
         ]
     )
 
-    # Deterministic index cache: distinct, monotonic per-logical-block values so
-    # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
     block_table = common.block_table_tensor
     num_pages = int(block_table.max().item()) + 1
+    block_table.copy_(
+        torch.randperm(num_pages, device=device, dtype=torch.int32).view_as(block_table)
+    )
     index_cache = torch.zeros(
         num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
     )
+    assert max(batch.seq_lens) <= BLOCK_SIZE * len(_E4M3_EXACT_VALUES)
     for r, seq_len in enumerate(batch.seq_lens):
-        for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
-            index_cache[block_table[r, b]] = float(b + 1)
+        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+        for b in range(num_blocks):
+            # Distinct request rankings expose crossed block-table rows.
+            rank = (b + 7 * r) % num_blocks
+            index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[rank])
+    padded_tokens = num_tokens + 7
     index_q = torch.ones(
-        num_tokens, num_idx_heads * head_dim, device=device, dtype=index_dtype
+        padded_tokens, num_idx_heads, head_dim, device=device, dtype=torch.float32
     )
+    # Different heads select opposite ends of the ranking, exposing head strides.
+    index_q[:, 1::2].neg_()
+    index_q = index_q.to(index_dtype).view(padded_tokens, -1)
 
     spec = MLAAttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=index_dtype
     )
     impl_kwargs = dict(
         num_kv_heads=num_idx_heads,
@@ -906,8 +945,9 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
         sparse_block_size=BLOCK_SIZE,
         num_index_heads=num_idx_heads,
         index_head_dim=head_dim,
-        init_blocks=0,
-        local_blocks=0,
+        init_blocks=init_blocks,
+        local_blocks=local_blocks,
+        indexer_kv_dtype="fp8" if index_dtype == torch.float8_e4m3fn else "bf16",
     )
 
     with set_current_vllm_config(vllm_config):
@@ -929,31 +969,40 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
     # only) return views. Separate buffers so the two forwards don't clobber.
     nd = sum(q for q in batch.query_lens if q <= 1)
     msa_impl.topk_indices_buffer = torch.full(
-        (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
+        (padded_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
     )
     triton_impl.topk_indices_buffer = torch.full(
-        (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
+        (padded_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
     )
 
     attn_metadata = {
         msa_impl.index_cache.prefix: msa_builder.build(0, common),
         triton_impl.index_cache.prefix: triton_builder.build(0, common),
     }
-    with set_forward_context(attn_metadata, vllm_config):
-        msa_decode, msa_prefill = msa_impl(index_q)
-        tri_decode, tri_prefill = triton_impl(index_q)
+    msa_metadata = attn_metadata[msa_impl.index_cache.prefix]
+    assert isinstance(msa_metadata, MiniMaxM3IndexerMSAMetadata)
+    assert (msa_metadata.prefill_cute is not None) == expect_cute
+
+    # Reuse one plan/arena for two layers with opposite score rankings.
+    for _ in range(2):
+        with set_forward_context(attn_metadata, vllm_config):
+            msa_decode, msa_prefill = msa_impl(index_q)
+            tri_decode, tri_prefill = triton_impl(index_q)
+
+        assert msa_decode is None and msa_prefill is None
+        assert tri_decode is not None and tri_prefill is not None
+        _assert_topk_indices_equal_unordered(
+            msa_impl.topk_indices_buffer[:num_tokens],
+            triton_impl.topk_indices_buffer[:num_tokens],
+        )
+        assert torch.all(msa_impl.topk_indices_buffer[num_tokens:] == -2)
+        assert torch.all(triton_impl.topk_indices_buffer[num_tokens:] == -2)
+        index_q.copy_(index_q.float().neg().to(index_dtype))
 
-    # MSA's return is vestigial; the attend reads its buffer directly.
-    assert msa_decode is None and msa_prefill is None
-    assert tri_decode is not None and tri_prefill is not None
-    _assert_topk_indices_equal_unordered(
-        msa_impl.topk_indices_buffer[:num_tokens],
-        triton_impl.topk_indices_buffer[:num_tokens],
-    )
     # Triton's decode/prefill outputs are views into its persistent buffer.
     buf_htk = triton_impl.topk_indices_buffer.transpose(0, 1)
     assert tri_decode.data_ptr() == buf_htk[:, :nd, :].data_ptr()
-    assert tri_prefill.data_ptr() == buf_htk[:, nd:, :].data_ptr()
+    assert tri_prefill.data_ptr() == buf_htk[:, nd:num_tokens, :].data_ptr()
 
 
 @pytest.mark.parametrize(

vllm/config/attention.py

@@ -12,6 +12,7 @@
 
 IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"]
 MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
+MiniMaxM3IndexerPrefillBackend = Literal["auto", "fmha", "cute"]
 
 
 @config
@@ -43,6 +44,11 @@ class AttentionConfig:
     minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton"
     """Sparse decode kernel used by the MiniMax M3 MSA backend."""
 
+    minimax_m3_indexer_prefill_backend: MiniMaxM3IndexerPrefillBackend = "auto"
+    """Prefill scorer for the MiniMax M3 MSA indexer. "auto" uses the CuTe
+    scorer for supported FP8 indexer Q/K and FMHA otherwise; "fmha" and "cute"
+    force a backend without changing indexer precision."""
+
     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/nvidia/indexer_msa.py

@@ -9,9 +9,8 @@
 causal page count and force-includes the init/local blocks, so the unwritten
 tail of the buffer is pre-filled with ``-inf``.
 
-Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faster
-than Triton for the wide prefill score, benchmarked ~3-5x), writing its
-``max_score`` straight into the buffer's prefill region (stride-aware, no copy).
+Prefill defaults to MSA's dedicated CuTe scorer for supported FP8 indexer Q/K
+and FMHA ``OnlyScore`` for BF16. Both write into the buffer's prefill region.
 
 Decode scores with CuteDSL when the flattened query tile is supported and fall
 back to Triton otherwise, writing into the decode region. Only the top-k is
@@ -22,12 +21,13 @@
 """
 
 from dataclasses import dataclass
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 import torch
 
 from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
 from vllm.models.minimax_m3.common.indexer import (
     MiniMaxM3IndexerBackend,
     MiniMaxM3IndexerDecodeMetadata,
@@ -39,6 +39,7 @@
     minimax_m3_index_decode_score,
 )
 from vllm.models.minimax_m3.nvidia.ops import minimax_m3_index_decode_score_cutedsl
+from vllm.platforms import current_platform
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -47,6 +48,11 @@
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.kv_cache_interface import AttentionSpec
 
+if TYPE_CHECKING:
+    from vllm.third_party.fmha_sm100 import BatchPrefillIndexerQ8KV8Wrapper
+
+logger = init_logger(__name__)
+
 # Page size == sparse block size == index-K block; fmha tile id == M3 block id.
 PAGE_SIZE = 128
 
@@ -84,10 +90,12 @@ class MiniMaxM3IndexerMSAPrefillMetadata:
 @dataclass
 class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
     """Decode reuses the inherited base ``decode`` field (the Triton decode
-    metadata); ``prefill_msa`` carries the fmha score plan for the prefill side
+    metadata); prefill uses either ``prefill_msa`` or ``prefill_cute``
     (the base ``prefill`` field is unused on this path)."""
 
     prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None
+    # Planned once per batch; every layer writes into unified_scores.
+    prefill_cute: "BatchPrefillIndexerQ8KV8Wrapper | None" = None
     # Per-forward view (``[:num_tokens]``) of the builder's persistent unified
     # score buffer ``[total_q, H, MAX_K_TILES]``, shared by decode and prefill
     # and reused across all layers. Pre-filled with the -inf sentinel in
@@ -123,6 +131,50 @@ def __init__(
         device: torch.device,
     ) -> None:
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        prefill_backend = (
+            vllm_config.attention_config.minimax_m3_indexer_prefill_backend
+        )
+        hf_config = vllm_config.model_config.hf_config
+        text_config = getattr(hf_config, "text_config", hf_config)
+        local_blocks = text_config.sparse_attention_config.get("sparse_local_block", 0)
+        capability = current_platform.get_device_capability(device.index or 0)
+        # The CuTe scorer skips the current page; top-k must force it in.
+        supports_cute = (
+            capability is not None
+            and capability.to_int() in (100, 103, 107)
+            and kv_cache_spec.dtype == torch.float8_e4m3fn
+            and kv_cache_spec.head_size == 128
+            and kv_cache_spec.block_size == PAGE_SIZE
+            and self.num_index_heads in (1, 2, 4)
+            and local_blocks >= 1
+        )
+        if prefill_backend == "cute" and not supports_cute:
+            raise ValueError(
+                "MiniMax M3 CuTe prefill scoring requires SM100/SM103/SM107, "
+                "FP8 E4M3 indexer Q/K, head dimension 128, page size 128, "
+                "1/2/4 index heads per rank, and sparse_local_block >= 1."
+            )
+        self.use_cute_prefill = supports_cute and prefill_backend != "fmha"
+        if self.use_cute_prefill:
+            from vllm.third_party import fmha_sm100
+
+            wrapper_cls = getattr(fmha_sm100, "BatchPrefillIndexerQ8KV8Wrapper", None)
+            if not hasattr(wrapper_cls, "run_scores"):
+                if prefill_backend == "cute":
+                    raise ValueError(
+                        "The installed MSA lacks the prefill run_scores API. "
+                        "Update MSA or select minimax_m3_indexer_prefill_backend=fmha."
+                    )
+                self.use_cute_prefill = False
+                logger.warning_once(
+                    "MSA lacks the prefill run_scores API; using FMHA OnlyScore."
+                )
+        logger.info_once(
+            "MiniMax M3 indexer prefill scorer: %s (Q/K dtype=%s, policy=%s)",
+            "CuTe" if self.use_cute_prefill else "FMHA OnlyScore",
+            kv_cache_spec.dtype,
+            prefill_backend,
+        )
         # Persistent unified score buffer [T, H, MAX_K_TILES] shared by all
         # indexer layers and reused across forwards. Stable address (required
         # for the captured decode path) + fixed tile dim so the decode score
@@ -196,7 +248,26 @@ def build(
             )
 
         prefill: MiniMaxM3IndexerMSAPrefillMetadata | None = None
-        if num_prefills > 0:
+        prefill_cute = None
+        if num_prefills > 0 and self.use_cute_prefill:
+            from vllm.third_party.fmha_sm100 import BatchPrefillIndexerQ8KV8Wrapper
+
+            lo, hi = num_decodes, num_reqs
+            prefill_cute = BatchPrefillIndexerQ8KV8Wrapper(
+                num_heads=self.num_index_heads
+            )
+            prefill_cute.plan(
+                query_start_loc[lo : hi + 1] - query_start_loc[lo],
+                seq_lens[lo:hi],
+                block_table[lo:hi],
+                total_q=num_prefill_tokens,
+                max_seqlen_q=common_attn_metadata.max_query_len,
+                max_seqlen_k=min(
+                    common_attn_metadata.max_seq_len, block_table.shape[1] * PAGE_SIZE
+                ),
+                score_out=unified_scores[num_decode_tokens:].view(-1, MAX_K_TILES),
+            )
+        elif num_prefills > 0:
             # Prefill is eager (not captured); the host lengths it needs (and the
             # _fmha_sm100_plan .tolist() inside) make the D->H sync acceptable.
             from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan
@@ -247,6 +318,7 @@ def build(
             num_prefill_tokens=num_prefill_tokens,
             decode=decode,
             prefill_msa=prefill,
+            prefill_cute=prefill_cute,
             unified_scores=unified_scores,
             max_k_tiles=max_k_tiles,
             topk_cu_seqlens_q=query_start_loc[: num_reqs + 1],
@@ -257,7 +329,7 @@ def build(
 
 
 class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
-    """Decode: CuteDSL/Triton score. Prefill: fmha_sm100 OnlyScore + top-k."""
+    """Decode: CuteDSL/Triton score. Prefill: CuTe or fmha_sm100 OnlyScore + top-k."""
 
     indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend
 
@@ -317,7 +389,9 @@ def forward(
                 score_out=unified_scores[:nd].transpose(0, 1),
             )
 
-        if md.prefill_msa is not None:
+        if md.prefill_cute is not None:
+            md.prefill_cute.run_scores(index_q[nd:], kv)
+        elif md.prefill_msa is not None:
             from vllm.third_party.fmha_sm100.api import _fmha_sm100
 
             p = md.prefill_msa