[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>
Thien Tran committed
Jul 31, 2026 at 16:01 UTC
34bb795ff3efee6cc08c9dd104deceefff2c4d55
3 files changed
+34
-23
.buildkite/test_areas/kernels.yaml
+4
@@ -287,8 +287,11 @@ steps:
287
- vllm/cute_utils/
288
- vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
289
- vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
290
+ - vllm/models/minimax_m3/nvidia/
291
+ - cmake/external_projects/fmha_sm100.cmake
292
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
293
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
294
+ - tests/kernels/attention/test_minimax_m3.py
295
- tests/kernels/test_ll_bf16_gemm.py
296
- tests/kernels/test_top_k_per_row.py
297
commands:
@@ -320,6 +323,7 @@ steps:
323
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
324
- pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
325
- pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
326
+ - pytest -v -s tests/kernels/attention/test_minimax_m3.py
327
- pytest -v -s tests/kernels/test_ll_bf16_gemm.py
328
# e2e
329
- pytest -v -s tests/models/quantization/test_nvfp4.py
tests/kernels/attention/test_minimax_m3.py
+24
-17
@@ -431,7 +431,7 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
431
not current_platform.is_device_capability_family(100),
432
reason="fmha_sm100 indexer requires SM100 (Blackwell).",
433
)
434
-@pytest.mark.parametrize("topk", [8, 16])
434
+@pytest.mark.parametrize("topk", [16])
435
@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
436
def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
437
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):
471
batch, BLOCK_SIZE, device, arange_block_indices=True
472
)
473
num_tokens = batch.compute_num_tokens()
474
+ # Absolute token positions; the MSA builder derives per-token causal page
475
+ # counts from them.
476
+ common.positions = torch.cat(
477
+ [
478
+ torch.arange(s - q, s, device=device, dtype=torch.int64)
479
+ for s, q in zip(batch.seq_lens, batch.query_lens)
480
+ ]
481
+ )
482
483
# Deterministic index cache: distinct, monotonic per-logical-block values so
484
# 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):
523
triton_impl.index_cache.kv_cache = index_cache
524
525
# Exercise the shared persistent top-k buffer for BOTH impls: each must write
518
- # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
519
- # Separate buffers so the two forwards don't clobber each other.
526
+ # decode ([:nd]) and prefill ([nd:]) into its token-major buffer and (Triton
527
+ # only) return views. Separate buffers so the two forwards don't clobber.
528
nd = sum(q for q in batch.query_lens if q <= 1)
529
msa_impl.topk_indices_buffer = torch.full(
522
- (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
530
+ (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
531
)
532
triton_impl.topk_indices_buffer = torch.full(
525
- (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
533
+ (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
534
)
535
536
attn_metadata = {
@@ -533,18 +541,17 @@ def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
541
msa_decode, msa_prefill = msa_impl(index_q)
542
tri_decode, tri_prefill = triton_impl(index_q)
543
536
- assert msa_decode is not None and tri_decode is not None
537
- assert msa_prefill is not None and tri_prefill is not None
538
- _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
539
- _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
540
- # decode/prefill outputs are views into each impl's persistent buffer.
541
- for impl, dec, pre in (
542
- (msa_impl, msa_decode, msa_prefill),
543
- (triton_impl, tri_decode, tri_prefill),
544
- ):
545
- buf = impl.topk_indices_buffer
546
- assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
547
- assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
544
+ # MSA's return is vestigial; the attend reads its buffer directly.
545
+ assert msa_decode is None and msa_prefill is None
546
+ assert tri_decode is not None and tri_prefill is not None
547
+ _assert_topk_indices_equal_unordered(
548
+ msa_impl.topk_indices_buffer[:num_tokens],
549
+ triton_impl.topk_indices_buffer[:num_tokens],
550
+ )
551
+ # Triton's decode/prefill outputs are views into its persistent buffer.
552
+ buf_htk = triton_impl.topk_indices_buffer.transpose(0, 1)
553
+ assert tri_decode.data_ptr() == buf_htk[:, :nd, :].data_ptr()
554
+ assert tri_prefill.data_ptr() == buf_htk[:, nd:, :].data_ptr()
555
556
557
@pytest.mark.parametrize(
vllm/models/minimax_m3/common/indexer.py
+6
-6
@@ -480,10 +480,10 @@ def select_indexer_impl_cls(
480
) -> type[MiniMaxM3IndexerImpl]:
481
"""Pick the indexer impl off the platform, top-k count, and cache dtype.
482
483
- On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
484
- main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
485
- bf16 and fp8 index caches. Everything else falls back to the Triton indexer
486
- (bf16 only).
483
+ On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's
484
+ ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is
485
+ used for both bf16 and fp8 index caches. Everything else falls back to the
486
+ Triton indexer (bf16 only).
487
"""
488
if indexer_kv_dtype in ("mxfp4", "nvfp4"):
489
raise NotImplementedError(
@@ -495,7 +495,7 @@ def select_indexer_impl_cls(
495
)
496
use_msa = (
497
is_sm100
498
- and topk_blocks in (4, 8, 16, 32)
498
+ and topk_blocks == 16
499
and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
500
)
501
if use_msa:
@@ -505,7 +505,7 @@ def select_indexer_impl_cls(
505
)
506
507
logger.info_once(
508
- "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) "
508
+ "MiniMax M3 indexer: selected MSA (fmha_sm100 score + top-k) "
509
"[topk_blocks=%d, indexer_kv_dtype=%s]",
510
topk_blocks,
511
indexer_kv_dtype,