[CI] Add M3 MSA tests to CI (#49143)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Codex <noreply@openai.com>

.buildkite/test_areas/kernels.yaml

@@ -287,8 +287,11 @@ steps:
   - vllm/cute_utils/
   - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
   - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
+  - vllm/models/minimax_m3/nvidia/
+  - cmake/external_projects/fmha_sm100.cmake
   - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
   - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
+  - tests/kernels/attention/test_minimax_m3.py
   - tests/kernels/test_ll_bf16_gemm.py
   - tests/kernels/test_top_k_per_row.py
   commands:
@@ -320,6 +323,7 @@ steps:
     - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
     - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
     - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
+    - pytest -v -s tests/kernels/attention/test_minimax_m3.py
     - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
     # e2e
     - pytest -v -s tests/models/quantization/test_nvfp4.py

tests/kernels/attention/test_minimax_m3.py

@@ -431,7 +431,7 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
     not current_platform.is_device_capability_family(100),
     reason="fmha_sm100 indexer requires SM100 (Blackwell).",
 )
-@pytest.mark.parametrize("topk", [8, 16])
+@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):
     import vllm.models.minimax_m3.common.indexer as indexer_mod
@@ -471,6 +471,14 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
         batch, BLOCK_SIZE, device, arange_block_indices=True
     )
     num_tokens = batch.compute_num_tokens()
+    # Absolute token positions; the MSA builder derives per-token causal page
+    # counts from them.
+    common.positions = torch.cat(
+        [
+            torch.arange(s - q, s, device=device, dtype=torch.int64)
+            for s, q in zip(batch.seq_lens, batch.query_lens)
+        ]
+    )
 
     # Deterministic index cache: distinct, monotonic per-logical-block values so
     # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
@@ -515,14 +523,14 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
     triton_impl.index_cache.kv_cache = index_cache
 
     # Exercise the shared persistent top-k buffer for BOTH impls: each must write
-    # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
-    # Separate buffers so the two forwards don't clobber each other.
+    # decode ([:nd]) and prefill ([nd:]) into its token-major buffer and (Triton
+    # 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_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+        (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
     )
     triton_impl.topk_indices_buffer = torch.full(
-        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+        (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
     )
 
     attn_metadata = {
@@ -533,18 +541,17 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
         msa_decode, msa_prefill = msa_impl(index_q)
         tri_decode, tri_prefill = triton_impl(index_q)
 
-    assert msa_decode is not None and tri_decode is not None
-    assert msa_prefill is not None and tri_prefill is not None
-    _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
-    _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
-    # decode/prefill outputs are views into each impl's persistent buffer.
-    for impl, dec, pre in (
-        (msa_impl, msa_decode, msa_prefill),
-        (triton_impl, tri_decode, tri_prefill),
-    ):
-        buf = impl.topk_indices_buffer
-        assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
-        assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
+    # 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()
 
 
 @pytest.mark.parametrize(

vllm/models/minimax_m3/common/indexer.py

@@ -480,10 +480,10 @@ def select_indexer_impl_cls(
 ) -> type[MiniMaxM3IndexerImpl]:
     """Pick the indexer impl off the platform, top-k count, and cache dtype.
 
-    On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
-    main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
-    bf16 and fp8 index caches. Everything else falls back to the Triton indexer
-    (bf16 only).
+    On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's
+    ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is
+    used for both bf16 and fp8 index caches. Everything else falls back to the
+    Triton indexer (bf16 only).
     """
     if indexer_kv_dtype in ("mxfp4", "nvfp4"):
         raise NotImplementedError(
@@ -495,7 +495,7 @@ def select_indexer_impl_cls(
     )
     use_msa = (
         is_sm100
-        and topk_blocks in (4, 8, 16, 32)
+        and topk_blocks == 16
         and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
     )
     if use_msa:
@@ -505,7 +505,7 @@ def select_indexer_impl_cls(
         )
 
         logger.info_once(
-            "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) "
+            "MiniMax M3 indexer: selected MSA (fmha_sm100 score + top-k) "
             "[topk_blocks=%d, indexer_kv_dtype=%s]",
             topk_blocks,
             indexer_kv_dtype,