From 34bb795ff3efee6cc08c9dd104deceefff2c4d55 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 31 Jul 2026 16:01:48 +0800 Subject: [PATCH] [CI] Add M3 MSA tests to CI (#49143) Signed-off-by: Thien Tran Co-authored-by: Codex --- .buildkite/test_areas/kernels.yaml | 4 +++ tests/kernels/attention/test_minimax_m3.py | 41 +++++++++++++--------- vllm/models/minimax_m3/common/indexer.py | 12 +++---- 3 files changed, 34 insertions(+), 23 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 3644990674..db1f909598 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 8a2bd5892a..8fe2342808 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/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( diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index fea22dbdac..767aa2d19b 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/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,