@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0055b8bfa3

[Attention][MiniMax-M3] Add MSA speculative decode verification (#50032)

Jason Li committed Aug 2, 2026 at 16:46 UTC 0055b8bfa3a2c7c49c51ff5962c162ba85f6de38
15 files changed +1275 -59
.buildkite/test_areas/kernels.yaml
+12 -1
@@ -287,8 +287,17 @@ 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/
290 - cmake/external_projects/fmha_sm100.cmake
291 + - vllm/models/minimax_m3/common/sparse_attention.py
292 + - vllm/models/minimax_m3/nvidia/
293 + - vllm/config/attention.py
294 + - vllm/v1/attention/backends/registry.py
295 + - vllm/_custom_ops.py
296 + - csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
297 + - csrc/libtorch_stable/ops.h
298 + - csrc/libtorch_stable/torch_bindings.cpp
299 + - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
300 + - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
301 - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
302 - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
303 - tests/kernels/attention/test_minimax_m3.py
@@ -304,6 +313,8 @@ steps:
313 - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
314 - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
315 - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
316 + - pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
317 + - pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
318 - pytest -v -s tests/kernels/test_top_k_per_row.py
319 # Quantization
320 - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
cmake/external_projects/fmha_sm100.cmake
+1 -1
@@ -17,7 +17,7 @@ else()
17 FetchContent_Declare(
18 fmha_sm100
19 GIT_REPOSITORY https://github.com/vllm-project/MSA.git
20 - GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
20 + GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa
21 GIT_PROGRESS TRUE
22 CONFIGURE_COMMAND ""
23 BUILD_COMMAND ""
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+66 -23
@@ -269,6 +269,22 @@ __device__ __forceinline__ void storeElemsFp8(
269 #endif
270 }
271
272 +// Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the
273 +// inverse dequantization scale and converting it to E4M3.
274 +template <typename scalar_t>
275 +__device__ __forceinline__ void storeScaledQElemsFp8(
276 + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane],
277 + float const inv_scale) {
278 + using Converter = vllm::_typeConvert<scalar_t>;
279 + float scaled[kElemsPerLane];
280 +#pragma unroll
281 + for (int i = 0; i < kElemsPerLane; i++) {
282 + auto const rounded = Converter::convert(elems[i]);
283 + scaled[i] = static_cast<float>(rounded) * inv_scale;
284 + }
285 + storeElemsFp8(dst, scaled);
286 +}
287 +
288 // ────────────────────────────────────────────────────────────────────────────
289 // Kernel
290 // ────────────────────────────────────────────────────────────────────────────
@@ -297,6 +313,7 @@ template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
313 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
314 scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
315 scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
316 + uint8_t* __restrict__ q_fp8_out, // [N, nq*128] E4M3, or nullptr
317 out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte
318 scalar_t const* __restrict__ q_norm_w,
319 scalar_t const* __restrict__ k_norm_w,
@@ -308,8 +325,9 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
325 int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
326 cache_t* __restrict__ kv_cache, // [nb,nkv,bs,2*128] or nullptr
327 out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
311 - float const eps, int const rotary_dim, int const num_tokens, int const nq,
312 - int const nkv, int const niq, int const block_size,
328 + float const eps, float const q_fp8_inv_scale, int const rotary_dim,
329 + int const num_tokens, int const nq, int const nkv, int const niq,
330 + int const block_size,
331 // kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128].
332 // The content (last) dim is always innermost-contiguous (stride 1), so the
333 // NHD/HND layout choice is captured by the head/token strides.
@@ -445,6 +463,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
463 } else {
464 storeElems<scalar_t>(store_ptr + dim_base, elems);
465 }
466 + if (isQ && q_fp8_out != nullptr) {
467 + storeScaledQElemsFp8<scalar_t>(
468 + q_fp8_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
469 + slot * kHeadDim + dim_base,
470 + elems, q_fp8_inv_scale);
471 + }
472 }
473
474 // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -493,13 +517,14 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
517 // ────────────────────────────────────────────────────────────────────────────
518 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
519 void launchFusedMiniMaxM3(
496 - scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
497 - scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
498 - scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
499 - int64_t const* positions, int64_t const* slot_mapping,
500 - int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
501 - float const eps, int const rotary_dim, int const num_tokens, int const nq,
502 - int const nkv, int const niq, int const block_size,
520 + scalar_t* qkv, scalar_t* q_out, uint8_t* q_fp8_out, void* index_q_out,
521 + scalar_t const* q_norm_w, scalar_t const* k_norm_w,
522 + scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
523 + scalar_t const* cos_sin_cache, int64_t const* positions,
524 + int64_t const* slot_mapping, int64_t const* index_slot_mapping,
525 + cache_t* kv_cache, void* index_cache, float const eps,
526 + float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens,
527 + int const nq, int const nkv, int const niq, int const block_size,
528 int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
529 int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
530 bool const process_index, bool const fp8_idx, cudaStream_t stream) {
@@ -540,10 +565,11 @@ void launchFusedMiniMaxM3(
565 fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
566 HAS_INDEX, INSERT, \
567 PROCESS_INDEX, FP8>, \
543 - qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
544 - iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
545 - index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
546 - eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
568 + qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
569 + q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
570 + slot_mapping, index_slot_mapping, kv_cache, \
571 + reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
572 + rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
573 kv_s_head, kv_s_token, kv_s_dim)
574 #else
575 // ROCm: standard kernel launch syntax (no PDL/stream serialization).
@@ -552,12 +578,12 @@ void launchFusedMiniMaxM3(
578 fusedMiniMaxM3QNormRopeKVInsertKernel< \
579 scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \
580 FP8><<<grid, kBlockSize, 0, stream>>>( \
555 - qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
556 - k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
581 + qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
582 + q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
583 slot_mapping, index_slot_mapping, kv_cache, \
558 - reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, num_tokens, \
559 - nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token, \
560 - kv_s_dim)
584 + reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
585 + rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
586 + kv_s_head, kv_s_token, kv_s_dim)
587 // clang-format on
588 #endif
589
@@ -599,6 +625,9 @@ void launchFusedMiniMaxM3(
625 vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
626 reinterpret_cast<st*>(qkv.data_ptr()), \
627 q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
628 + q_fp8_out.has_value() \
629 + ? reinterpret_cast<uint8_t*>(q_fp8_out->data_ptr()) \
630 + : nullptr, \
631 index_q_out.has_value() \
632 ? reinterpret_cast<void*>(index_q_out->data_ptr()) \
633 : nullptr, \
@@ -622,10 +651,10 @@ void launchFusedMiniMaxM3(
651 (insert_kv && process_index) \
652 ? reinterpret_cast<void*>(index_cache->data_ptr()) \
653 : nullptr, \
625 - static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
626 - nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_head, \
627 - kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx, \
628 - stream)
654 + static_cast<float>(eps), 1.0f / static_cast<float>(q_fp8_scale), \
655 + static_cast<int>(rotary_dim), num_tokens, nq, nkv, niq, \
656 + static_cast<int>(block_size), kv_s_block, kv_s_head, kv_s_token, \
657 + kv_s_dim, has_index, insert_kv, process_index, fp8_idx, stream)
658 // clang-format on
659
660 // ────────────────────────────────────────────────────────────────────────────
@@ -649,7 +678,10 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
678 std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
679 std::optional<torch::stable::Tensor>
680 index_q_out, // [N, niq*128] contiguous
652 - const std::string& kv_cache_dtype, bool skip_index_branch) {
681 + const std::string& kv_cache_dtype, bool skip_index_branch,
682 + std::optional<torch::stable::Tensor>
683 + q_fp8_out, // [N, nq*128] contiguous E4M3
684 + double q_fp8_scale) {
685 STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
686 "qkv must be contiguous CUDA");
687 STD_TORCH_CHECK(
@@ -780,6 +812,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
812 q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
813 "q_out must have num_tokens * num_heads * 128 elements");
814 }
815 + if (q_fp8_out.has_value()) {
816 + STD_TORCH_CHECK(q_fp8_out->is_cuda() && q_fp8_out->is_contiguous() &&
817 + q_fp8_out->scalar_type() ==
818 + torch::headeronly::ScalarType::Float8_e4m3fn,
819 + "q_fp8_out must be a contiguous CUDA fp8 e4m3 tensor");
820 + STD_TORCH_CHECK(
821 + q_fp8_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
822 + "q_fp8_out must have num_tokens * num_heads * 128 elements");
823 + STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0,
824 + "q_fp8_scale must be finite and positive");
825 + }
826 if (index_q_out.has_value()) {
827 STD_TORCH_CHECK(process_index,
828 "index_q_out requires index branch processing");
csrc/libtorch_stable/ops.h
+2 -1
@@ -376,7 +376,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
376 std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
377 std::optional<torch::stable::Tensor> q_out,
378 std::optional<torch::stable::Tensor> index_q_out,
379 - const std::string& kv_cache_dtype, bool skip_index_branch);
379 + const std::string& kv_cache_dtype, bool skip_index_branch,
380 + std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale);
381
382 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
383 void fused_kda_decode(
csrc/libtorch_stable/torch_bindings.cpp
+2 -1
@@ -514,7 +514,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
514 "Tensor? slot_mapping, Tensor? index_slot_mapping, "
515 "Tensor!? kv_cache, Tensor!? index_cache, "
516 "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
517 - "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
517 + "str kv_cache_dtype, bool skip_index_branch=False, "
518 + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0) -> ()");
519
520 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
521 ops.def(
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py new
+561
@@ -0,0 +1,561 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for MiniMax M3 CUTLASS sparse decode."""
4 +
5 +import math
6 +from types import SimpleNamespace
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm import _custom_ops as ops
12 +from vllm.config import AttentionConfig
13 +from vllm.models.minimax_m3.common.ops.sparse_attn import (
14 + minimax_m3_sparse_attn_decode,
15 +)
16 +from vllm.models.minimax_m3.common.sparse_attention import (
17 + MiniMaxM3SparseBackend,
18 + MiniMaxM3SparseMetadataBuilder,
19 + MiniMaxM3SparseTritonImpl,
20 + select_main_backend_and_impl_cls,
21 +)
22 +from vllm.models.minimax_m3.nvidia import (
23 + sparse_attention_msa as sparse_attention_msa_module,
24 +)
25 +from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
26 + MSACutlassDecodePlanCache,
27 + msa_cutlass_sparse_decode,
28 + prepare_decode_metadata,
29 + should_prepare_decode_metadata,
30 +)
31 +from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
32 + MiniMaxM3SparseMSABackend,
33 + MiniMaxM3SparseMSADecodeMetadata,
34 + MiniMaxM3SparseMSAImpl,
35 + MiniMaxM3SparseMSAMetadataBuilder,
36 +)
37 +from vllm.platforms import current_platform
38 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
39 +
40 +if not current_platform.is_device_capability_family(100):
41 + pytest.skip(
42 + "fmha_sm100 sparse decode requires SM100 (Blackwell).",
43 + allow_module_level=True,
44 + )
45 +
46 +
47 +HEAD_DIM = 128
48 +BLOCK_SIZE = 128
49 +TOPK = 16
50 +DEFAULT_QUERY_LEN = 4
51 +SM_SCALE = HEAD_DIM**-0.5
52 +
53 +
54 +@pytest.mark.parametrize(
55 + (
56 + "batch_size",
57 + "decode_query_len",
58 + "num_q_heads",
59 + "num_kv_heads",
60 + "expected",
61 + ),
62 + [
63 + pytest.param(8, 4, 64, 4, False, id="tp1-below-min-batch"),
64 + pytest.param(16, 4, 64, 4, True, id="tp1-supported"),
65 + pytest.param(16, 4, 16, 1, True, id="tp4-min-batch"),
66 + pytest.param(24, 4, 16, 1, True, id="tp4-intermediate-batch"),
67 + pytest.param(32, 4, 16, 1, True, id="tp4-supported"),
68 + pytest.param(16, 1, 64, 4, True, id="tp1-query-len-1"),
69 + pytest.param(16, 1, 16, 1, True, id="tp4-query-len-1"),
70 + pytest.param(16, 2, 64, 4, True, id="tp1-query-len-2"),
71 + pytest.param(16, 2, 16, 1, True, id="tp4-query-len-2"),
72 + pytest.param(16, 32, 64, 4, True, id="query-len-upper-bound"),
73 + pytest.param(16, 0, 64, 4, False, id="query-len-zero"),
74 + pytest.param(16, 33, 64, 4, False, id="query-len-above-bound"),
75 + ],
76 +)
77 +def test_msa_cutlass_decode_static_dispatch(
78 + batch_size: int,
79 + decode_query_len: int,
80 + expected: bool,
81 + num_q_heads: int,
82 + num_kv_heads: int,
83 +) -> None:
84 + assert (
85 + should_prepare_decode_metadata(
86 + batch_size,
87 + decode_query_len,
88 + decode_backend="cutlass",
89 + num_q_heads=num_q_heads,
90 + num_kv_heads=num_kv_heads,
91 + kv_cache_dtype="fp8_e4m3",
92 + page_size=BLOCK_SIZE,
93 + topk_blocks=TOPK,
94 + )
95 + is expected
96 + )
97 +
98 +
99 +def test_msa_cutlass_decode_static_dispatch_requires_opt_in() -> None:
100 + assert not should_prepare_decode_metadata(
101 + 32,
102 + DEFAULT_QUERY_LEN,
103 + decode_backend="triton",
104 + num_q_heads=16,
105 + num_kv_heads=1,
106 + kv_cache_dtype="fp8_e4m3",
107 + page_size=BLOCK_SIZE,
108 + topk_blocks=TOPK,
109 + )
110 +
111 +
112 +def test_msa_cutlass_decode_static_dispatch_accepts_fp8_alias() -> None:
113 + assert should_prepare_decode_metadata(
114 + 32,
115 + DEFAULT_QUERY_LEN,
116 + decode_backend="cutlass",
117 + num_q_heads=16,
118 + num_kv_heads=1,
119 + kv_cache_dtype="fp8",
120 + page_size=BLOCK_SIZE,
121 + topk_blocks=TOPK,
122 + )
123 +
124 +
125 +@pytest.mark.parametrize(
126 + "kv_cache_dtype",
127 + ["auto", "bfloat16", "float16", "fp8_e5m2"],
128 +)
129 +def test_msa_cutlass_decode_static_dispatch_requires_fp8_e4m3(
130 + kv_cache_dtype: str,
131 +) -> None:
132 + assert not should_prepare_decode_metadata(
133 + 32,
134 + DEFAULT_QUERY_LEN,
135 + decode_backend="cutlass",
136 + num_q_heads=16,
137 + num_kv_heads=1,
138 + kv_cache_dtype=kv_cache_dtype,
139 + page_size=BLOCK_SIZE,
140 + topk_blocks=TOPK,
141 + )
142 +
143 +
144 +def test_msa_cutlass_decode_static_dispatch_requires_sm100(
145 + monkeypatch: pytest.MonkeyPatch,
146 +) -> None:
147 + monkeypatch.setattr(
148 + current_platform,
149 + "is_device_capability_family",
150 + lambda _: False,
151 + )
152 + assert not should_prepare_decode_metadata(
153 + 32,
154 + DEFAULT_QUERY_LEN,
155 + decode_backend="cutlass",
156 + num_q_heads=16,
157 + num_kv_heads=1,
158 + kv_cache_dtype="fp8",
159 + page_size=BLOCK_SIZE,
160 + topk_blocks=TOPK,
161 + )
162 +
163 +
164 +@pytest.mark.parametrize(
165 + ("backend", "expected"),
166 + [
167 + (AttentionBackendEnum.CUTLASS_MSA, "cutlass"),
168 + (AttentionBackendEnum.TRITON_MSA, "triton"),
169 + ],
170 +)
171 +def test_msa_attention_backend_alias(
172 + backend: AttentionBackendEnum,
173 + expected: str,
174 +) -> None:
175 + config = AttentionConfig(backend=backend)
176 + assert config.backend is None
177 + assert config.minimax_m3_msa_decode_backend == expected
178 +
179 +
180 +def test_msa_backend_owns_msa_metadata_builder(
181 + monkeypatch: pytest.MonkeyPatch,
182 +) -> None:
183 + backend_cls, impl_cls = select_main_backend_and_impl_cls(
184 + topk_blocks=TOPK,
185 + kv_cache_dtype="fp8_e4m3",
186 + num_kv_heads=1,
187 + )
188 + assert backend_cls is MiniMaxM3SparseMSABackend
189 + assert impl_cls is MiniMaxM3SparseMSAImpl
190 +
191 + monkeypatch.setattr(
192 + current_platform,
193 + "is_device_capability_family",
194 + lambda _: False,
195 + )
196 + backend_cls, impl_cls = select_main_backend_and_impl_cls(
197 + topk_blocks=TOPK,
198 + kv_cache_dtype="fp8_e4m3",
199 + num_kv_heads=1,
200 + )
201 + assert backend_cls is MiniMaxM3SparseBackend
202 + assert impl_cls is MiniMaxM3SparseTritonImpl
203 +
204 +
205 +def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
206 + monkeypatch: pytest.MonkeyPatch,
207 +) -> None:
208 + batch = 16
209 + block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda")
210 + seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda")
211 + base_metadata = SimpleNamespace(
212 + num_decodes=batch,
213 + decode=SimpleNamespace(
214 + block_table=block_table,
215 + seq_lens=seq_lens,
216 + decode_query_len=1,
217 + ),
218 + )
219 + monkeypatch.setattr(
220 + MiniMaxM3SparseMetadataBuilder,
221 + "build",
222 + lambda *args, **kwargs: base_metadata,
223 + )
224 + expected_metadata = object()
225 + monkeypatch.setattr(
226 + sparse_attention_msa_module,
227 + "prepare_decode_metadata",
228 + lambda *args, **kwargs: expected_metadata,
229 + )
230 +
231 + builder = object.__new__(MiniMaxM3SparseMSAMetadataBuilder)
232 + builder.num_q_heads = 64
233 + builder.num_kv_heads = 4
234 + builder.topk_blocks = TOPK
235 + builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4)
236 + builder.kv_cache_dtype = "fp8_e4m3"
237 + builder.decode_backend = "cutlass"
238 + builder.msa_cutlass_plan_cache = object()
239 +
240 + metadata = builder.build(
241 + 0,
242 + SimpleNamespace(
243 + seq_lens_cpu_upper_bound=torch.full((batch,), 257, dtype=torch.int32)
244 + ),
245 + )
246 +
247 + assert isinstance(metadata.decode, MiniMaxM3SparseMSADecodeMetadata)
248 + assert metadata.decode.decode_query_len == 1
249 + assert metadata.decode.msa_cutlass is expected_metadata
250 +
251 +
252 +def test_msa_cutlass_plan_cache_keys_query_len(
253 + monkeypatch: pytest.MonkeyPatch,
254 +) -> None:
255 + batch = 16
256 + block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda")
257 + seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda")
258 + seq_lens_cpu = torch.full((batch,), 257, dtype=torch.int32)
259 + plan_cache = MSACutlassDecodePlanCache()
260 + built_query_lens: list[int] = []
261 +
262 + def fake_build_plan(**kwargs):
263 + query_len = kwargs["decode_query_len"]
264 + built_query_lens.append(query_len)
265 + num_rows = batch * query_len
266 + return (
267 + None,
268 + None,
269 + None,
270 + {
271 + "kv_segment_lens": torch.empty(
272 + num_rows, dtype=torch.int32, device="cuda"
273 + ),
274 + "qo_offset": torch.empty(num_rows, dtype=torch.int32, device="cuda"),
275 + },
276 + )
277 +
278 + monkeypatch.setattr(plan_cache, "_build_plan", fake_build_plan)
279 + first = prepare_decode_metadata(
280 + block_table,
281 + seq_lens,
282 + seq_lens_cpu,
283 + 1,
284 + num_q_heads=64,
285 + num_kv_heads=4,
286 + page_size=BLOCK_SIZE,
287 + topk_blocks=TOPK,
288 + plan_cache=plan_cache,
289 + )
290 + repeated = prepare_decode_metadata(
291 + block_table,
292 + seq_lens,
293 + seq_lens_cpu,
294 + 1,
295 + num_q_heads=64,
296 + num_kv_heads=4,
297 + page_size=BLOCK_SIZE,
298 + topk_blocks=TOPK,
299 + plan_cache=plan_cache,
300 + )
301 + different = prepare_decode_metadata(
302 + block_table,
303 + seq_lens,
304 + seq_lens_cpu,
305 + 2,
306 + num_q_heads=64,
307 + num_kv_heads=4,
308 + page_size=BLOCK_SIZE,
309 + topk_blocks=TOPK,
310 + plan_cache=plan_cache,
311 + )
312 + current_platform.synchronize()
313 +
314 + assert first.plan is repeated.plan
315 + assert different.plan is not first.plan
316 + assert built_query_lens == [1, 2]
317 +
318 +
319 +def _make_topk(
320 + seq_lens: list[int],
321 + num_kv_heads: int,
322 + query_len: int,
323 +) -> torch.Tensor:
324 + topk = torch.full(
325 + (len(seq_lens) * query_len, num_kv_heads, TOPK),
326 + -1,
327 + dtype=torch.int32,
328 + device="cuda",
329 + )
330 + for request, seq_len in enumerate(seq_lens):
331 + for local_query in range(query_len):
332 + token = request * query_len + local_query
333 + visible_tokens = seq_len - query_len + local_query + 1
334 + visible_pages = math.ceil(visible_tokens / BLOCK_SIZE)
335 + topk[token, :, :visible_pages] = torch.arange(
336 + visible_pages, dtype=torch.int32, device="cuda"
337 + )
338 + return topk
339 +
340 +
341 +@pytest.mark.parametrize(
342 + (
343 + "num_q_heads",
344 + "num_kv_heads",
345 + "num_request_pairs",
346 + "query_len",
347 + "capture_graph",
348 + ),
349 + [
350 + pytest.param(64, 4, 8, 1, True, id="tp1-query-len-1"),
351 + pytest.param(64, 4, 8, 2, False, id="tp1-query-len-2"),
352 + pytest.param(64, 4, 8, 3, False, id="tp1-query-len-3"),
353 + pytest.param(64, 4, 8, 4, True, id="tp1-query-len-4"),
354 + pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
355 + pytest.param(16, 1, 8, 1, True, id="tp4-query-len-1"),
356 + pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
357 + ],
358 +)
359 +def test_msa_cutlass_decode_matches_triton_with_interleaved_cache(
360 + num_q_heads: int,
361 + num_kv_heads: int,
362 + num_request_pairs: int,
363 + query_len: int,
364 + capture_graph: bool,
365 +) -> None:
366 + torch.manual_seed(0)
367 + seq_lens_list = [257, 513] * num_request_pairs
368 + seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32)
369 + seq_lens = seq_lens_cpu.cuda()
370 + pages_per_request = [math.ceil(seq_len / BLOCK_SIZE) for seq_len in seq_lens_list]
371 + num_pages = sum(pages_per_request)
372 + max_pages = max(pages_per_request)
373 +
374 + block_table = torch.zeros(
375 + len(seq_lens_list), max_pages, dtype=torch.int32, device="cuda"
376 + )
377 + physical_pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda")
378 + offset = 0
379 + for request, request_pages in enumerate(pages_per_request):
380 + block_table[request, :request_pages] = physical_pages[
381 + offset : offset + request_pages
382 + ]
383 + offset += request_pages
384 +
385 + key = (
386 + torch.randn(
387 + num_pages,
388 + num_kv_heads,
389 + BLOCK_SIZE,
390 + HEAD_DIM,
391 + dtype=torch.bfloat16,
392 + device="cuda",
393 + )
394 + * 0.25
395 + ).to(torch.float8_e4m3fn)
396 + value = (torch.randn_like(key, dtype=torch.bfloat16) * 0.25).to(torch.float8_e4m3fn)
397 + kv_cache = torch.cat((key, value), dim=-1)
398 + assert kv_cache.stride() == (
399 + num_kv_heads * BLOCK_SIZE * 2 * HEAD_DIM,
400 + BLOCK_SIZE * 2 * HEAD_DIM,
401 + 2 * HEAD_DIM,
402 + 1,
403 + )
404 +
405 + num_query_tokens = len(seq_lens_list) * query_len
406 + query = torch.randn(
407 + num_query_tokens,
408 + num_q_heads,
409 + HEAD_DIM,
410 + dtype=torch.bfloat16,
411 + device="cuda",
412 + )
413 + q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
414 + query_fp8 = torch.empty_like(query, dtype=torch.float8_e4m3fn)
415 + ops.scaled_fp8_quant(
416 + query.view(num_query_tokens, -1),
417 + scale=q_scale,
418 + output=query_fp8.view(num_query_tokens, -1),
419 + )
420 + query_dequantized = query_fp8.to(torch.bfloat16) * q_scale
421 +
422 + topk_token_major = _make_topk(seq_lens_list, num_kv_heads, query_len)
423 + expected = torch.empty_like(query)
424 + minimax_m3_sparse_attn_decode(
425 + query_dequantized,
426 + kv_cache,
427 + topk_token_major.transpose(0, 1),
428 + block_table,
429 + seq_lens,
430 + num_kv_heads,
431 + SM_SCALE,
432 + expected,
433 + query_len,
434 + k_scale=None,
435 + v_scale=None,
436 + )
437 +
438 + plan_cache = MSACutlassDecodePlanCache()
439 + metadata = prepare_decode_metadata(
440 + block_table,
441 + seq_lens,
442 + seq_lens_cpu,
443 + query_len,
444 + num_q_heads=num_q_heads,
445 + num_kv_heads=num_kv_heads,
446 + page_size=BLOCK_SIZE,
447 + topk_blocks=TOPK,
448 + plan_cache=plan_cache,
449 + )
450 + assert metadata.page_table.data_ptr() == block_table.data_ptr()
451 + actual = torch.empty_like(query)
452 + msa_cutlass_sparse_decode(
453 + query_fp8,
454 + kv_cache,
455 + topk_token_major,
456 + actual,
457 + metadata,
458 + scale=SM_SCALE,
459 + q_scale_float=1.0,
460 + k_scale_float=1.0,
461 + v_scale_float=1.0,
462 + )
463 +
464 + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
465 +
466 + # The same captured plan must remain correct as ragged lengths change.
467 + updated_seq_lens_list = [129, 385] * num_request_pairs
468 + seq_lens.copy_(
469 + torch.tensor(updated_seq_lens_list, dtype=torch.int32, device="cuda")
470 + )
471 + updated_seq_lens_cpu = torch.tensor(updated_seq_lens_list, dtype=torch.int32)
472 + topk_token_major.copy_(_make_topk(updated_seq_lens_list, num_kv_heads, query_len))
473 + updated_metadata = prepare_decode_metadata(
474 + block_table,
475 + seq_lens,
476 + updated_seq_lens_cpu,
477 + query_len,
478 + num_q_heads=num_q_heads,
479 + num_kv_heads=num_kv_heads,
480 + page_size=BLOCK_SIZE,
481 + topk_blocks=TOPK,
482 + plan_cache=plan_cache,
483 + )
484 + assert updated_metadata.plan is metadata.plan
485 + assert updated_metadata.page_table.data_ptr() == metadata.page_table.data_ptr()
486 +
487 + minimax_m3_sparse_attn_decode(
488 + query_dequantized,
489 + kv_cache,
490 + topk_token_major.transpose(0, 1),
491 + block_table,
492 + seq_lens,
493 + num_kv_heads,
494 + SM_SCALE,
495 + expected,
496 + query_len,
497 + k_scale=None,
498 + v_scale=None,
499 + )
500 + msa_cutlass_sparse_decode(
501 + query_fp8,
502 + kv_cache,
503 + topk_token_major,
504 + actual,
505 + updated_metadata,
506 + scale=SM_SCALE,
507 + q_scale_float=1.0,
508 + k_scale_float=1.0,
509 + v_scale_float=1.0,
510 + )
511 + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
512 +
513 + if capture_graph:
514 + graph = torch.cuda.CUDAGraph()
515 + with torch.cuda.graph(graph):
516 + msa_cutlass_sparse_decode(
517 + query_fp8,
518 + kv_cache,
519 + topk_token_major,
520 + actual,
521 + updated_metadata,
522 + scale=SM_SCALE,
523 + q_scale_float=1.0,
524 + k_scale_float=1.0,
525 + v_scale_float=1.0,
526 + )
527 +
528 + replay_seq_lens_list = [257, 513] * num_request_pairs
529 + seq_lens.copy_(
530 + torch.tensor(replay_seq_lens_list, dtype=torch.int32, device="cuda")
531 + )
532 + topk_token_major.copy_(
533 + _make_topk(replay_seq_lens_list, num_kv_heads, query_len)
534 + )
535 + prepare_decode_metadata(
536 + block_table,
537 + seq_lens,
538 + torch.tensor(replay_seq_lens_list, dtype=torch.int32),
539 + query_len,
540 + num_q_heads=num_q_heads,
541 + num_kv_heads=num_kv_heads,
542 + page_size=BLOCK_SIZE,
543 + topk_blocks=TOPK,
544 + plan_cache=plan_cache,
545 + )
546 + minimax_m3_sparse_attn_decode(
547 + query_dequantized,
548 + kv_cache,
549 + topk_token_major.transpose(0, 1),
550 + block_table,
551 + seq_lens,
552 + num_kv_heads,
553 + SM_SCALE,
554 + expected,
555 + query_len,
556 + k_scale=None,
557 + v_scale=None,
558 + )
559 + graph.replay()
560 + current_platform.synchronize()
561 + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+15
@@ -208,6 +208,12 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
208 # index_q here (de-interleaved from the packed qkv); k/v/index_k stay in
209 # place inside qkv and are scatter-inserted into the caches.
210 q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
211 + q_fp8 = torch.empty(
212 + num_tokens,
213 + qsz,
214 + dtype=torch.float8_e4m3fn,
215 + device=device,
216 + )
217 index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device)
218
219 ops.fused_minimax_m3_qknorm_rope_kv_insert(
@@ -231,6 +237,8 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
237 q_out,
238 index_q,
239 kv_cache_dtype,
240 + q_fp8_out=q_fp8,
241 + q_fp8_scale=0.5,
242 )
243
244 # ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are
@@ -262,6 +270,13 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
270 # reference materializes bf16 after the norm (the unfused boundary), so
271 # rounding-boundary elements can differ by ~1 bf16 ulp.
272 torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2)
273 + expected_q_fp8 = torch.empty_like(q_fp8)
274 + ops.scaled_fp8_quant(
275 + q_out,
276 + scale=torch.tensor(0.5, dtype=torch.float32, device=device),
277 + output=expected_q_fp8,
278 + )
279 + torch.testing.assert_close(q_fp8, expected_q_fp8, rtol=0, atol=0)
280 torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2)
281 torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2)
282 torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
vllm/_custom_ops.py
+7
@@ -2655,6 +2655,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2655 index_q_out: torch.Tensor | None = None,
2656 kv_cache_dtype: str = "auto",
2657 skip_index_branch: bool = False,
2658 + q_fp8_out: torch.Tensor | None = None,
2659 + q_fp8_scale: float = 1.0,
2660 ) -> None:
2661 """Fused MiniMax-M3 attention pre-processing (in-place).
2662
@@ -2677,6 +2679,9 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2679 callers skip a separate ``.contiguous()`` copy before the SM100 sparse
2680 attention's flat TMA descriptor.
2681
2682 + If ``q_fp8_out`` is given, the same normalized q is also written in FP8
2683 + E4M3 using ``q_fp8_scale`` as its dequantization scale.
2684 +
2685 When ``skip_index_branch`` is true, sparse rows still keep their packed
2686 ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
2687 branches and main KV cache. This is used by MiniMax-M3 index-topk reuse
@@ -2704,6 +2709,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2709 index_q_out,
2710 kv_cache_dtype,
2711 skip_index_branch,
2712 + q_fp8_out,
2713 + q_fp8_scale,
2714 )
2715
2716
vllm/config/attention.py
+15
@@ -11,6 +11,7 @@ from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnu
11 from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
13 IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"]
14 +MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
15
16
17 @config
@@ -20,6 +21,9 @@ class AttentionConfig:
21 backend: AttentionBackendEnum | None = None
22 """Attention backend to use. Use "auto" or None for automatic selection."""
23
24 + minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton"
25 + """Sparse decode kernel used by the MiniMax M3 MSA backend."""
26 +
27 backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict)
28 """Per-KV-cache-group attention backend overrides, keyed by
29 `KVCacheSpecKind` (e.g. `{"mla_attention": "FLASHINFER_MLA",
@@ -99,6 +103,17 @@ class AttentionConfig:
103 PyTorch >= 2.9, and 128 for encoder-only attention or older PyTorch
104 versions."""
105
106 + def __post_init__(self) -> None:
107 + msa_aliases: dict[AttentionBackendEnum, MiniMaxM3MSADecodeBackend] = {
108 + AttentionBackendEnum.CUTLASS_MSA: "cutlass",
109 + AttentionBackendEnum.TRITON_MSA: "triton",
110 + }
111 + if self.backend in msa_aliases:
112 + self.minimax_m3_msa_decode_backend = msa_aliases[self.backend]
113 + # The alias selects only MiniMax's sparse decode kernel. Dense
114 + # layers still use the platform's normal automatic backend.
115 + self.backend = None
116 +
117 def compute_hash(self) -> str:
118 """
119 Provide a hash that uniquely identifies all the configs
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+2
@@ -23,6 +23,8 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
23 query: torch.Tensor,
24 kv_cache: torch.Tensor,
25 output: torch.Tensor,
26 + *,
27 + query_fp8: torch.Tensor | None = None,
28 ) -> torch.Tensor:
29 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
30 minimax_m3_sparse_attn_decode_aiter,
vllm/models/minimax_m3/common/sparse_attention.py
+30 -7
@@ -330,6 +330,7 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
330 *,
331 topk_blocks: int,
332 sparse_block_size: int,
333 + msa_decode_backend: str = "triton",
334 ) -> None:
335 self.num_heads = num_heads
336 self.head_size = head_size
@@ -355,6 +356,8 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
356 query: torch.Tensor,
357 kv_cache: torch.Tensor,
358 output: torch.Tensor,
359 + *,
360 + query_fp8: torch.Tensor | None = None,
361 ) -> torch.Tensor:
362 """Attend the queries to the indexer-selected blocks. Per kernel.
363
@@ -364,6 +367,9 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
367 """
368 raise NotImplementedError
369
370 + def should_use_msa_decode(self, layer_name: str) -> bool:
371 + return False
372 +
373
374 class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
375 """Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode."""
@@ -374,6 +380,8 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
380 query: torch.Tensor,
381 kv_cache: torch.Tensor,
382 output: torch.Tensor,
383 + *,
384 + query_fp8: torch.Tensor | None = None,
385 ) -> torch.Tensor:
386 attn_metadata = get_forward_context().attn_metadata
387 if not isinstance(attn_metadata, dict):
@@ -442,18 +450,18 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
450 return output
451
452
445 -def select_main_impl_cls(
453 +def select_main_backend_and_impl_cls(
454 *,
455 topk_blocks: int,
456 kv_cache_dtype: str,
457 num_kv_heads: int,
450 -) -> type[MiniMaxM3SparseImpl]:
451 - """Pick the main attend impl off the main KV-cache dtype.
458 +) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]:
459 + """Pick the main attention backend and implementation.
460
461 Blackwell (SM100) uses the MSA attend for supported top-k block counts
462 when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA
463 with shuffle KV cache layout; Other platforms and FP8 E5M2 fall
456 - back to Triton. The MSA modules are imported lazily avoid import errors
464 + back to Triton. The MSA modules are imported lazily to avoid import errors
465 on unsupported platforms.
466 """
467 use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(num_kv_heads)
@@ -477,11 +485,26 @@ def select_main_impl_cls(
485 MiniMaxM3SparseAiterPAImpl,
486 )
487
480 - return MiniMaxM3SparseAiterPAImpl
488 + return MiniMaxM3SparseBackend, MiniMaxM3SparseAiterPAImpl
489 if use_msa:
490 from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
491 + MiniMaxM3SparseMSABackend,
492 MiniMaxM3SparseMSAImpl,
493 )
494
486 - return MiniMaxM3SparseMSAImpl
487 - return MiniMaxM3SparseTritonImpl
495 + return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl
496 + return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl
497 +
498 +
499 +def select_main_impl_cls(
500 + *,
501 + topk_blocks: int,
502 + kv_cache_dtype: str,
503 + num_kv_heads: int,
504 +) -> type[MiniMaxM3SparseImpl]:
505 + """Backward-compatible implementation-only selector."""
506 + return select_main_backend_and_impl_cls(
507 + topk_blocks=topk_blocks,
508 + kv_cache_dtype=kv_cache_dtype,
509 + num_kv_heads=num_kv_heads,
510 + )[1]
vllm/models/minimax_m3/nvidia/model.py
+28 -6
@@ -79,7 +79,7 @@ from vllm.models.minimax_m3.common.mm_preprocess import (
79 from vllm.models.minimax_m3.common.sparse_attention import (
80 MiniMaxM3SparseBackend,
81 MiniMaxM3SparseImpl,
82 - select_main_impl_cls,
82 + select_main_backend_and_impl_cls,
83 )
84 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
85 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -505,15 +505,15 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
505 # the attend impl reads them back (so nothing crosses the eager break as a
506 # Python value, which would freeze at capture).
507 self.topk_indices_buffer = topk_indices_buffer
508 - self.attn_backend = MiniMaxM3SparseBackend
508 # Indexer (top-k selection) and main attention are separate impls, each
509 # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
510 # (broader than the AttentionImpl that AttentionLayerBase annotates).
512 - self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
511 + self.attn_backend, impl_cls = select_main_backend_and_impl_cls(
512 topk_blocks=sparse_cfg["sparse_topk_blocks"],
513 kv_cache_dtype=self.kv_cache_dtype,
514 num_kv_heads=self.num_kv_heads,
516 - )(
515 + )
516 + self.impl: MiniMaxM3SparseImpl = impl_cls( # type: ignore[assignment]
517 self.num_heads,
518 self.head_dim,
519 self.scaling,
@@ -521,6 +521,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
521 kv_cache_dtype=self.kv_cache_dtype,
522 topk_blocks=sparse_cfg["sparse_topk_blocks"],
523 sparse_block_size=sparse_cfg["sparse_block_size"],
524 + msa_decode_backend=(
525 + vllm_config.attention_config.minimax_m3_msa_decode_backend
526 + ),
527 )
528 # Self-contained nn.Module: owns its side cache, selects its impl in init.
529 self.indexer = MiniMaxM3Indexer(
@@ -594,6 +597,16 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
597 main_slot_mapping = fwd_slot_mapping[self.layer_name]
598 index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
599 q = qkv.new_empty((num_tokens, self.q_size))
600 + use_msa_decode = self.impl.should_use_msa_decode(self.layer_name)
601 + query_fp8 = (
602 + torch.empty(
603 + (num_tokens, self.q_size),
604 + dtype=torch.float8_e4m3fn,
605 + device=qkv.device,
606 + )
607 + if use_msa_decode
608 + else None
609 + )
610 # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
611 # the fused kernel emits fp8 directly when this buffer is e4m3.
612 index_q = qkv.new_empty(
@@ -621,10 +634,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
634 q,
635 index_q,
636 self.kv_cache_dtype,
637 + q_fp8_out=query_fp8,
638 + q_fp8_scale=self._q_scale_float,
639 )
640
641 output = torch.empty_like(q)
627 - attn_output = self._run_attention(q, index_q, output)
642 + attn_output = self._run_attention(q, query_fp8, index_q, output)
643 output, _ = self.o_proj(attn_output)
644 return output
645
@@ -632,6 +647,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
647 def _run_attention(
648 self,
649 query: torch.Tensor,
650 + query_fp8: torch.Tensor | None,
651 index_query: torch.Tensor,
652 output: torch.Tensor,
653 ) -> torch.Tensor:
@@ -639,7 +655,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
655 # metadata and can't be captured into a cudagraph. The indexer writes its
656 # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
657 self.indexer(index_query)
642 - return self.impl.forward(self, query, self.kv_cache, output)
658 + return self.impl.forward(
659 + self,
660 + query,
661 + self.kv_cache,
662 + output,
663 + query_fp8=query_fp8,
664 + )
665
666
667 class MiniMaxM3DecoderLayer(nn.Module):
vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py new
+316
@@ -0,0 +1,316 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MiniMax CUTLASS sparse decode using per-query-token page indices."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass, field
8 +from typing import Any
9 +
10 +import torch
11 +
12 +from vllm.config.attention import MiniMaxM3MSADecodeBackend
13 +from vllm.platforms import current_platform
14 +from vllm.triton_utils import tl, triton
15 +
16 +_MAX_NUM_Q_HEADS = 64
17 +_MAX_NUM_KV_HEADS = 4
18 +_HEAD_DIM = 128
19 +_PAGE_SIZE = 128
20 +_TOPK = 16
21 +# fmha_sm100 plans one row per query head. Keep every cached plan within the
22 +# fixed planner allocation used by the MSA decode kernel.
23 +_MAX_QUERY_HEAD_ROWS = 65536
24 +_MAX_DECODE_QUERY_LEN = 32
25 +# Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4.
26 +_MIN_CUTLASS_BATCH_SIZE = 16
27 +
28 +
29 +@dataclass
30 +class MSACutlassDecodeMetadata:
31 + plan: Any
32 + page_table: torch.Tensor
33 +
34 +
35 +@triton.jit
36 +def _update_runtime_metadata_kernel(
37 + seq_lens_ptr,
38 + kv_segment_lens_ptr,
39 + qo_offset_ptr,
40 + num_rows: tl.constexpr,
41 + decode_query_len: tl.constexpr,
42 + BLOCK_SIZE: tl.constexpr,
43 +):
44 + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
45 + mask = offsets < num_rows
46 + request = offsets // decode_query_len
47 + local_query = offsets % decode_query_len
48 + seq_len = tl.load(seq_lens_ptr + request, mask=mask)
49 + tl.store(kv_segment_lens_ptr + offsets, seq_len, mask=mask)
50 + tl.store(
51 + qo_offset_ptr + offsets,
52 + seq_len - decode_query_len + local_query,
53 + mask=mask,
54 + )
55 +
56 +
57 +@dataclass
58 +class MSACutlassDecodePlanCache:
59 + """Reusable plans whose mutable tensors retain cudagraph-stable addresses."""
60 +
61 + plans: dict[tuple[int, ...], Any] = field(init=False, default_factory=dict)
62 +
63 + def _build_plan(
64 + self,
65 + *,
66 + batch: int,
67 + decode_query_len: int,
68 + page_table_stride: int,
69 + initial_seq_lens_cpu: torch.Tensor,
70 + device: torch.device,
71 + num_q_heads: int,
72 + num_kv_heads: int,
73 + page_size: int,
74 + topk_blocks: int,
75 + ) -> Any:
76 + from vllm.third_party.fmha_sm100.api import fmha_sm100_plan
77 +
78 + qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32)
79 + kv_lens_cpu = initial_seq_lens_cpu
80 + plan = fmha_sm100_plan(
81 + qo_lens_cpu,
82 + kv_lens_cpu,
83 + num_q_heads,
84 + num_kv_heads=num_kv_heads,
85 + qo_offset=kv_lens_cpu - qo_lens_cpu,
86 + page_size=page_size,
87 + output_maxscore=False,
88 + kv_block_num=topk_blocks,
89 + causal=True,
90 + sparse_kernel_mode="decode",
91 + use_fp8_kvcache=True,
92 + split_prefill_decode=False,
93 + device=device,
94 + )
95 +
96 + plan_info = plan[3]
97 + row_starts = (
98 + torch.arange(batch, dtype=torch.int32, device=device)
99 + .mul_(page_table_stride)
100 + .repeat_interleave(decode_query_len)
101 + )
102 + page_indptr = torch.cat(
103 + (
104 + row_starts,
105 + torch.tensor(
106 + [batch * page_table_stride],
107 + dtype=torch.int32,
108 + device=device,
109 + ),
110 + )
111 + )
112 + plan_info["kv_page_indptr"].copy_(page_indptr)
113 + return plan
114 +
115 + def prepare(
116 + self,
117 + block_table: torch.Tensor,
118 + seq_lens: torch.Tensor,
119 + seq_lens_cpu: torch.Tensor,
120 + decode_query_len: int,
121 + *,
122 + num_q_heads: int,
123 + num_kv_heads: int,
124 + page_size: int,
125 + topk_blocks: int,
126 + ) -> MSACutlassDecodeMetadata:
127 + batch = int(seq_lens.shape[0])
128 + if (
129 + block_table.device.type != "cuda"
130 + or block_table.dtype != torch.int32
131 + or not block_table.is_contiguous()
132 + or block_table.shape[0] != batch
133 + ):
134 + raise ValueError(
135 + "MSA sparse decode requires a contiguous CUDA int32 block "
136 + "table with one row per request"
137 + )
138 + if (
139 + seq_lens.dtype != torch.int32
140 + or not seq_lens.is_contiguous()
141 + or seq_lens.device != block_table.device
142 + ):
143 + raise ValueError(
144 + "MSA sparse decode requires contiguous CUDA int32 sequence "
145 + "lengths on the block table device"
146 + )
147 + if (
148 + seq_lens_cpu.device.type != "cpu"
149 + or seq_lens_cpu.dtype != torch.int32
150 + or not seq_lens_cpu.is_contiguous()
151 + or seq_lens_cpu.shape != seq_lens.shape
152 + ):
153 + raise ValueError(
154 + "MSA sparse decode requires contiguous CPU int32 sequence "
155 + "lengths matching the device sequence lengths"
156 + )
157 +
158 + page_table_stride = int(block_table.stride(0))
159 + key = (
160 + batch,
161 + decode_query_len,
162 + page_table_stride,
163 + num_q_heads,
164 + num_kv_heads,
165 + page_size,
166 + topk_blocks,
167 + )
168 + plan = self.plans.get(key)
169 + if plan is None:
170 + plan = self._build_plan(
171 + batch=batch,
172 + decode_query_len=decode_query_len,
173 + page_table_stride=page_table_stride,
174 + initial_seq_lens_cpu=seq_lens_cpu,
175 + device=seq_lens.device,
176 + num_q_heads=num_q_heads,
177 + num_kv_heads=num_kv_heads,
178 + page_size=page_size,
179 + topk_blocks=topk_blocks,
180 + )
181 + self.plans[key] = plan
182 +
183 + plan_info = plan[3]
184 + num_rows = batch * decode_query_len
185 + _update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)](
186 + seq_lens,
187 + plan_info["kv_segment_lens"],
188 + plan_info["qo_offset"],
189 + num_rows=num_rows,
190 + decode_query_len=decode_query_len,
191 + BLOCK_SIZE=128,
192 + )
193 + return MSACutlassDecodeMetadata(
194 + plan=plan,
195 + page_table=block_table.view(-1),
196 + )
197 +
198 +
199 +def _supported_head_geometry(num_q_heads: int, num_kv_heads: int) -> bool:
200 + return (
201 + 0 < num_q_heads <= _MAX_NUM_Q_HEADS
202 + and 0 < num_kv_heads <= _MAX_NUM_KV_HEADS
203 + and num_q_heads % num_kv_heads == 0
204 + )
205 +
206 +
207 +def supports_cutlass_sparse_decode(
208 + *,
209 + decode_backend: MiniMaxM3MSADecodeBackend,
210 + num_q_heads: int,
211 + num_kv_heads: int,
212 + kv_cache_dtype: str,
213 + page_size: int,
214 + topk_blocks: int,
215 +) -> bool:
216 + """Return whether static model geometry supports CUTLASS sparse decode."""
217 + return (
218 + decode_backend == "cutlass"
219 + and current_platform.is_cuda()
220 + and current_platform.is_device_capability_family(100)
221 + and kv_cache_dtype in ("fp8", "fp8_e4m3")
222 + and _supported_head_geometry(num_q_heads, num_kv_heads)
223 + and page_size == _PAGE_SIZE
224 + and topk_blocks == _TOPK
225 + )
226 +
227 +
228 +def should_prepare_decode_metadata(
229 + batch_size: int,
230 + decode_query_len: int,
231 + *,
232 + decode_backend: MiniMaxM3MSADecodeBackend,
233 + num_q_heads: int,
234 + num_kv_heads: int,
235 + kv_cache_dtype: str,
236 + page_size: int,
237 + topk_blocks: int,
238 +) -> bool:
239 + """Return whether a graph shape can use the CUTLASS decode path."""
240 + total_q = batch_size * decode_query_len
241 + return (
242 + supports_cutlass_sparse_decode(
243 + decode_backend=decode_backend,
244 + num_q_heads=num_q_heads,
245 + num_kv_heads=num_kv_heads,
246 + kv_cache_dtype=kv_cache_dtype,
247 + page_size=page_size,
248 + topk_blocks=topk_blocks,
249 + )
250 + and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
251 + and batch_size >= _MIN_CUTLASS_BATCH_SIZE
252 + and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
253 + )
254 +
255 +
256 +@torch.no_grad()
257 +def prepare_decode_metadata(
258 + block_table: torch.Tensor,
259 + seq_lens: torch.Tensor,
260 + seq_lens_cpu: torch.Tensor,
261 + decode_query_len: int,
262 + *,
263 + num_q_heads: int,
264 + num_kv_heads: int,
265 + page_size: int,
266 + topk_blocks: int,
267 + plan_cache: MSACutlassDecodePlanCache | None = None,
268 +) -> MSACutlassDecodeMetadata:
269 + """Prepare graph-stable runtime metadata for one sparse decode step."""
270 + cache = plan_cache or MSACutlassDecodePlanCache()
271 + return cache.prepare(
272 + block_table,
273 + seq_lens,
274 + seq_lens_cpu,
275 + decode_query_len,
276 + num_q_heads=num_q_heads,
277 + num_kv_heads=num_kv_heads,
278 + page_size=page_size,
279 + topk_blocks=topk_blocks,
280 + )
281 +
282 +
283 +@torch.no_grad()
284 +def msa_cutlass_sparse_decode(
285 + query_fp8: torch.Tensor,
286 + kv_cache: torch.Tensor,
287 + topk: torch.Tensor,
288 + output: torch.Tensor,
289 + metadata: MSACutlassDecodeMetadata,
290 + *,
291 + scale: float,
292 + q_scale_float: float,
293 + k_scale_float: float,
294 + v_scale_float: float,
295 +) -> None:
296 + """Run CUTLASS sparse decode with metadata prepared by the MSA builder."""
297 + key, value = kv_cache.split(_HEAD_DIM, dim=-1)
298 +
299 + from vllm.third_party.fmha_sm100.api import fmha_sm100
300 +
301 + fmha_sm100(
302 + query_fp8,
303 + key,
304 + value,
305 + metadata.plan,
306 + kv_indices=metadata.page_table,
307 + kv_block_indexes=topk,
308 + out=output,
309 + output_maxscore=False,
310 + output_o=True,
311 + sm_scale=scale,
312 + q_scale=q_scale_float,
313 + k_scale=k_scale_float,
314 + v_scale=v_scale_float,
315 + o_scale=1.0,
316 + )
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+210 -19
@@ -1,28 +1,198 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""MSA (SM100/Blackwell) block-sparse attend for MiniMax M3.
3 +"""MSA (SM100/Blackwell) block-sparse attention for MiniMax M3.
4
5 -Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``);
6 -decode falls back to the Triton split-K kernel (no MSA decode yet). ``fmha_sm100``
7 -imports are function-local, so this module is import-safe on AMD/non-SM100.
5 +Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``).
6 +Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100``
7 +path for regular decode and speculative verification.
8 """
9
10 +from dataclasses import dataclass
11 +
12 import torch
13
14 +from vllm.config import VllmConfig
15 +from vllm.config.attention import MiniMaxM3MSADecodeBackend
16 from vllm.forward_context import get_forward_context
17 +from vllm.logger import init_logger
18 from vllm.models.minimax_m3.common.ops.sparse_attn import (
19 SPARSE_BLOCK_SIZE,
20 minimax_m3_sparse_attn_decode,
21 )
22 from vllm.models.minimax_m3.common.sparse_attention import (
23 + MiniMaxM3SparseBackend,
24 + MiniMaxM3SparseDecodeMetadata,
25 MiniMaxM3SparseImpl,
26 MiniMaxM3SparseMetadata,
27 + MiniMaxM3SparseMetadataBuilder,
28 +)
29 +from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
30 + MSACutlassDecodeMetadata,
31 + MSACutlassDecodePlanCache,
32 + msa_cutlass_sparse_decode,
33 + prepare_decode_metadata,
34 + should_prepare_decode_metadata,
35 + supports_cutlass_sparse_decode,
36 +)
37 +from vllm.v1.attention.backend import (
38 + AttentionLayer,
39 + CommonAttentionMetadata,
40 )
21 -from vllm.v1.attention.backend import AttentionLayer
41 +from vllm.v1.kv_cache_interface import AttentionSpec
42 +
43 +logger = init_logger(__name__)
44 +
45 +
46 +class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
47 + """MiniMax M3 backend with NVIDIA MSA-specific decode metadata."""
48 +
49 + @staticmethod
50 + def get_builder_cls() -> type["MiniMaxM3SparseMSAMetadataBuilder"]:
51 + return MiniMaxM3SparseMSAMetadataBuilder
52 +
53 +
54 +class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend):
55 + """Attention-backend alias selecting CUTLASS MSA sparse decode."""
56 +
57 + @staticmethod
58 + def get_name() -> str:
59 + return "CUTLASS_MSA"
60 +
61 +
62 +class MiniMaxM3SparseTritonBackend(MiniMaxM3SparseMSABackend):
63 + """Attention-backend alias selecting Triton MSA sparse decode."""
64 +
65 + @staticmethod
66 + def get_name() -> str:
67 + return "TRITON_MSA"
68 +
69 +
70 +@dataclass
71 +class MiniMaxM3SparseMSADecodeMetadata(MiniMaxM3SparseDecodeMetadata):
72 + msa_cutlass: MSACutlassDecodeMetadata | None = None
73 +
74 +
75 +class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
76 + """Prepare MSA plans only for decode shapes supported by ``fmha_sm100``."""
77 +
78 + def __init__(
79 + self,
80 + kv_cache_spec: AttentionSpec,
81 + layer_names: list[str],
82 + vllm_config: VllmConfig,
83 + device: torch.device,
84 + ) -> None:
85 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
86 + config = vllm_config.model_config.hf_text_config
87 + tp_size = vllm_config.parallel_config.tensor_parallel_size
88 + self.num_q_heads = config.num_attention_heads // tp_size
89 + self.num_kv_heads = kv_cache_spec.num_kv_heads
90 + self.topk_blocks = config.sparse_attention_config["sparse_topk_blocks"]
91 + # AttentionSpec stores every FP8 mode as uint8, so retain the configured
92 + # format to distinguish E4M3 (supported) from E5M2 before planning.
93 + self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
94 + self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend
95 + self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache()
96 +
97 + def build(
98 + self,
99 + common_prefix_len: int,
100 + common_attn_metadata: CommonAttentionMetadata,
101 + fast_build: bool = False,
102 + ) -> MiniMaxM3SparseMetadata:
103 + metadata = super().build(
104 + common_prefix_len,
105 + common_attn_metadata,
106 + fast_build,
107 + )
108 + decode = metadata.decode
109 + if decode is None:
110 + return metadata
111 +
112 + msa_cutlass = None
113 + if should_prepare_decode_metadata(
114 + metadata.num_decodes,
115 + decode.decode_query_len,
116 + decode_backend=self.decode_backend,
117 + num_q_heads=self.num_q_heads,
118 + num_kv_heads=self.num_kv_heads,
119 + kv_cache_dtype=self.kv_cache_dtype,
120 + page_size=SPARSE_BLOCK_SIZE,
121 + topk_blocks=self.topk_blocks,
122 + ):
123 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
124 + assert seq_lens_cpu is not None
125 + msa_cutlass = prepare_decode_metadata(
126 + decode.block_table,
127 + decode.seq_lens,
128 + seq_lens_cpu[: metadata.num_decodes],
129 + decode.decode_query_len,
130 + num_q_heads=self.num_q_heads,
131 + num_kv_heads=self.num_kv_heads,
132 + page_size=SPARSE_BLOCK_SIZE,
133 + topk_blocks=self.topk_blocks,
134 + plan_cache=self.msa_cutlass_plan_cache,
135 + )
136 + metadata.decode = MiniMaxM3SparseMSADecodeMetadata(
137 + seq_lens=decode.seq_lens,
138 + block_table=decode.block_table,
139 + decode_query_len=decode.decode_query_len,
140 + msa_cutlass=msa_cutlass,
141 + )
142 + return metadata
143
144
145 class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
25 - """MSA block-sparse attend (``fmha_sm100``); Triton split-K decode."""
146 + """MSA block-sparse attention with guarded CUTLASS sparse decode."""
147 +
148 + def __init__(
149 + self,
150 + num_heads: int,
151 + head_size: int,
152 + scale: float,
153 + num_kv_heads: int | None = None,
154 + kv_cache_dtype: str = "auto",
155 + *,
156 + topk_blocks: int,
157 + sparse_block_size: int,
158 + msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton",
159 + ) -> None:
160 + super().__init__(
161 + num_heads,
162 + head_size,
163 + scale,
164 + num_kv_heads,
165 + kv_cache_dtype,
166 + topk_blocks=topk_blocks,
167 + sparse_block_size=sparse_block_size,
168 + )
169 + self.use_cutlass_decode = supports_cutlass_sparse_decode(
170 + decode_backend=msa_decode_backend,
171 + num_q_heads=self.num_heads,
172 + num_kv_heads=self.num_kv_heads,
173 + kv_cache_dtype=self.kv_cache_dtype,
174 + page_size=self.block_size,
175 + topk_blocks=self.topk_blocks,
176 + )
177 + logger.info_once(
178 + "MiniMax M3 MSA sparse decode selected %s",
179 + "CUTLASS" if self.use_cutlass_decode else "Triton",
180 + )
181 +
182 + def should_use_msa_decode(self, layer_name: str) -> bool:
183 + if not self.use_cutlass_decode:
184 + return False
185 + attn_metadata = get_forward_context().attn_metadata
186 + if not isinstance(attn_metadata, dict):
187 + return False
188 + main_md = attn_metadata[layer_name]
189 + if not isinstance(main_md, MiniMaxM3SparseMetadata):
190 + return False
191 + decode = main_md.decode
192 + return (
193 + isinstance(decode, MiniMaxM3SparseMSADecodeMetadata)
194 + and decode.msa_cutlass is not None
195 + )
196
197 def forward(
198 self,
@@ -30,6 +200,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
200 query: torch.Tensor,
201 kv_cache: torch.Tensor,
202 output: torch.Tensor,
203 + *,
204 + query_fp8: torch.Tensor | None = None,
205 ) -> torch.Tensor:
206 attn_metadata = get_forward_context().attn_metadata
207 if not isinstance(attn_metadata, dict):
@@ -52,23 +224,42 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
224 k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
225 v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
226
55 - # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
227 + # Decode [:nd]: CUTLASS for planned shapes, otherwise Triton.
228 if main_md.num_decodes > 0:
229 d = main_md.decode
230 assert d is not None
59 - minimax_m3_sparse_attn_decode(
60 - q[:nd],
61 - kv_cache,
62 - topk[:nd].transpose(0, 1),
63 - d.block_table,
64 - d.seq_lens,
65 - self.num_kv_heads,
66 - self.scale,
67 - out[:nd],
68 - d.decode_query_len,
69 - k_scale=k_scale,
70 - v_scale=v_scale,
231 + msa_metadata = (
232 + d.msa_cutlass
233 + if isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
234 + else None
235 )
236 + if self.use_cutlass_decode and msa_metadata is not None:
237 + assert query_fp8 is not None
238 + msa_cutlass_sparse_decode(
239 + query_fp8[:nd].view(-1, self.num_heads, hd),
240 + kv_cache,
241 + topk[:nd],
242 + out[:nd],
243 + msa_metadata,
244 + scale=self.scale,
245 + q_scale_float=getattr(layer, "_q_scale_float", 1.0),
246 + k_scale_float=getattr(layer, "_k_scale_float", 1.0),
247 + v_scale_float=getattr(layer, "_v_scale_float", 1.0),
248 + )
249 + else:
250 + minimax_m3_sparse_attn_decode(
251 + q[:nd],
252 + kv_cache,
253 + topk[:nd].transpose(0, 1),
254 + d.block_table,
255 + d.seq_lens,
256 + self.num_kv_heads,
257 + self.scale,
258 + out[:nd],
259 + d.decode_query_len,
260 + k_scale=k_scale,
261 + v_scale=v_scale,
262 + )
263
264 # Prefill [nd:]: MSA sparse FMHA over the selected blocks.
265 if main_md.num_prefills > 0:
vllm/v1/attention/backends/registry.py
+8
@@ -101,6 +101,14 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
101 MINIMAX_M3_SPARSE = (
102 "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend"
103 )
104 + CUTLASS_MSA = (
105 + "vllm.models.minimax_m3.nvidia.sparse_attention_msa."
106 + "MiniMaxM3SparseCutlassBackend"
107 + )
108 + TRITON_MSA = (
109 + "vllm.models.minimax_m3.nvidia.sparse_attention_msa."
110 + "MiniMaxM3SparseTritonBackend"
111 + )
112 NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
113 FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
114 # HPC Attention Backend: