obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-22 13:40:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jason Li
2026-08-02 13:46:55 -07:00
committed by GitHub
parent c666810676
commit 0055b8bfa3
15 changed files with 1275 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+12 -1
View File
@@ -287,8 +287,17 @@ 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
- vllm/models/minimax_m3/common/sparse_attention.py
- vllm/models/minimax_m3/nvidia/
- vllm/config/attention.py
- vllm/v1/attention/backends/registry.py
- vllm/_custom_ops.py
- csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
- csrc/libtorch_stable/ops.h
- csrc/libtorch_stable/torch_bindings.cpp
- tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
- tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
- tests/kernels/attention/test_minimax_m3.py
@@ -304,6 +313,8 @@ steps:
- pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
- pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
- pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
- pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
- pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
- pytest -v -s tests/kernels/test_top_k_per_row.py
# Quantization
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
cmake/external_projects/fmha_sm100.cmake
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+66 -23
View File
@@ -269,6 +269,22 @@ __device__ __forceinline__ void storeElemsFp8(
#endif
}
// Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the
// inverse dequantization scale and converting it to E4M3.
template <typename scalar_t>
__device__ __forceinline__ void storeScaledQElemsFp8(
uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane],
float const inv_scale) {
using Converter = vllm::_typeConvert<scalar_t>;
float scaled[kElemsPerLane];
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
auto const rounded = Converter::convert(elems[i]);
scaled[i] = static_cast<float>(rounded) * inv_scale;
}
storeElemsFp8(dst, scaled);
}
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
@@ -297,6 +313,7 @@ template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
uint8_t* __restrict__ q_fp8_out, // [N, nq*128] E4M3, or nullptr
out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte
scalar_t const* __restrict__ q_norm_w,
scalar_t const* __restrict__ k_norm_w,
@@ -308,8 +325,9 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
cache_t* __restrict__ kv_cache, // [nb,nkv,bs,2*128] or nullptr
out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
float const eps, int const rotary_dim, int const num_tokens, int const nq,
int const nkv, int const niq, int const block_size,
float const eps, float const q_fp8_inv_scale, int const rotary_dim,
int const num_tokens, int const nq, int const nkv, int const niq,
int const block_size,
// kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128].
// The content (last) dim is always innermost-contiguous (stride 1), so the
// NHD/HND layout choice is captured by the head/token strides.
@@ -445,6 +463,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
} else {
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
if (isQ && q_fp8_out != nullptr) {
storeScaledQElemsFp8<scalar_t>(
q_fp8_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
slot * kHeadDim + dim_base,
elems, q_fp8_inv_scale);
}
}
// ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -493,13 +517,14 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
void launchFusedMiniMaxM3(
scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
int64_t const* positions, int64_t const* slot_mapping,
int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
float const eps, int const rotary_dim, int const num_tokens, int const nq,
int const nkv, int const niq, int const block_size,
scalar_t* qkv, scalar_t* q_out, uint8_t* q_fp8_out, void* index_q_out,
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
scalar_t const* cos_sin_cache, int64_t const* positions,
int64_t const* slot_mapping, int64_t const* index_slot_mapping,
cache_t* kv_cache, void* index_cache, float const eps,
float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens,
int const nq, int const nkv, int const niq, int const block_size,
int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
bool const process_index, bool const fp8_idx, cudaStream_t stream) {
@@ -540,10 +565,11 @@ void launchFusedMiniMaxM3(
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
HAS_INDEX, INSERT, \
PROCESS_INDEX, FP8>, \
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
kv_s_head, kv_s_token, kv_s_dim)
#else
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
@@ -552,12 +578,12 @@ void launchFusedMiniMaxM3(
fusedMiniMaxM3QNormRopeKVInsertKernel< \
scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \
FP8><<<grid, kBlockSize, 0, stream>>>( \
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, num_tokens, \
nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token, \
kv_s_dim)
reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
kv_s_head, kv_s_token, kv_s_dim)
// clang-format on
#endif
@@ -599,6 +625,9 @@ void launchFusedMiniMaxM3(
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
reinterpret_cast<st*>(qkv.data_ptr()), \
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
q_fp8_out.has_value() \
? reinterpret_cast<uint8_t*>(q_fp8_out->data_ptr()) \
: nullptr, \
index_q_out.has_value() \
? reinterpret_cast<void*>(index_q_out->data_ptr()) \
: nullptr, \
@@ -622,10 +651,10 @@ void launchFusedMiniMaxM3(
(insert_kv && process_index) \
? reinterpret_cast<void*>(index_cache->data_ptr()) \
: nullptr, \
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_head, \
kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx, \
stream)
static_cast<float>(eps), 1.0f / static_cast<float>(q_fp8_scale), \
static_cast<int>(rotary_dim), num_tokens, nq, nkv, niq, \
static_cast<int>(block_size), kv_s_block, kv_s_head, kv_s_token, \
kv_s_dim, has_index, insert_kv, process_index, fp8_idx, stream)
// clang-format on
// ────────────────────────────────────────────────────────────────────────────
@@ -649,7 +678,10 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
std::optional<torch::stable::Tensor>
index_q_out, // [N, niq*128] contiguous
const std::string& kv_cache_dtype, bool skip_index_branch) {
const std::string& kv_cache_dtype, bool skip_index_branch,
std::optional<torch::stable::Tensor>
q_fp8_out, // [N, nq*128] contiguous E4M3
double q_fp8_scale) {
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
"qkv must be contiguous CUDA");
STD_TORCH_CHECK(
@@ -780,6 +812,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
"q_out must have num_tokens * num_heads * 128 elements");
}
if (q_fp8_out.has_value()) {
STD_TORCH_CHECK(q_fp8_out->is_cuda() && q_fp8_out->is_contiguous() &&
q_fp8_out->scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn,
"q_fp8_out must be a contiguous CUDA fp8 e4m3 tensor");
STD_TORCH_CHECK(
q_fp8_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
"q_fp8_out must have num_tokens * num_heads * 128 elements");
STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0,
"q_fp8_scale must be finite and positive");
}
if (index_q_out.has_value()) {
STD_TORCH_CHECK(process_index,
"index_q_out requires index branch processing");
csrc/libtorch_stable/ops.h
+2 -1
View File
@@ -376,7 +376,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
std::optional<torch::stable::Tensor> q_out,
std::optional<torch::stable::Tensor> index_q_out,
const std::string& kv_cache_dtype, bool skip_index_branch);
const std::string& kv_cache_dtype, bool skip_index_branch,
std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale);
#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
void fused_kda_decode(
csrc/libtorch_stable/torch_bindings.cpp
+2 -1
View File
@@ -514,7 +514,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor? slot_mapping, Tensor? index_slot_mapping, "
"Tensor!? kv_cache, Tensor!? index_cache, "
"int block_size, Tensor!? q_out, Tensor!? index_q_out, "
"str kv_cache_dtype, bool skip_index_branch=False) -> ()");
"str kv_cache_dtype, bool skip_index_branch=False, "
"Tensor!? q_fp8_out=None, float q_fp8_scale=1.0) -> ()");
#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
ops.def(
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+561
View File
@@ -0,0 +1,561 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Correctness tests for MiniMax M3 CUTLASS sparse decode."""
import math
from types import SimpleNamespace
import pytest
import torch
from vllm import _custom_ops as ops
from vllm.config import AttentionConfig
from vllm.models.minimax_m3.common.ops.sparse_attn import (
minimax_m3_sparse_attn_decode,
)
from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseBackend,
MiniMaxM3SparseMetadataBuilder,
MiniMaxM3SparseTritonImpl,
select_main_backend_and_impl_cls,
)
from vllm.models.minimax_m3.nvidia import (
sparse_attention_msa as sparse_attention_msa_module,
)
from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
MSACutlassDecodePlanCache,
msa_cutlass_sparse_decode,
prepare_decode_metadata,
should_prepare_decode_metadata,
)
from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
MiniMaxM3SparseMSABackend,
MiniMaxM3SparseMSADecodeMetadata,
MiniMaxM3SparseMSAImpl,
MiniMaxM3SparseMSAMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backends.registry import AttentionBackendEnum
if not current_platform.is_device_capability_family(100):
pytest.skip(
"fmha_sm100 sparse decode requires SM100 (Blackwell).",
allow_module_level=True,
)
HEAD_DIM = 128
BLOCK_SIZE = 128
TOPK = 16
DEFAULT_QUERY_LEN = 4
SM_SCALE = HEAD_DIM**-0.5
@pytest.mark.parametrize(
(
"batch_size",
"decode_query_len",
"num_q_heads",
"num_kv_heads",
"expected",
),
[
pytest.param(8, 4, 64, 4, False, id="tp1-below-min-batch"),
pytest.param(16, 4, 64, 4, True, id="tp1-supported"),
pytest.param(16, 4, 16, 1, True, id="tp4-min-batch"),
pytest.param(24, 4, 16, 1, True, id="tp4-intermediate-batch"),
pytest.param(32, 4, 16, 1, True, id="tp4-supported"),
pytest.param(16, 1, 64, 4, True, id="tp1-query-len-1"),
pytest.param(16, 1, 16, 1, True, id="tp4-query-len-1"),
pytest.param(16, 2, 64, 4, True, id="tp1-query-len-2"),
pytest.param(16, 2, 16, 1, True, id="tp4-query-len-2"),
pytest.param(16, 32, 64, 4, True, id="query-len-upper-bound"),
pytest.param(16, 0, 64, 4, False, id="query-len-zero"),
pytest.param(16, 33, 64, 4, False, id="query-len-above-bound"),
],
)
def test_msa_cutlass_decode_static_dispatch(
batch_size: int,
decode_query_len: int,
expected: bool,
num_q_heads: int,
num_kv_heads: int,
) -> None:
assert (
should_prepare_decode_metadata(
batch_size,
decode_query_len,
decode_backend="cutlass",
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
kv_cache_dtype="fp8_e4m3",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
is expected
)
def test_msa_cutlass_decode_static_dispatch_requires_opt_in() -> None:
assert not should_prepare_decode_metadata(
32,
DEFAULT_QUERY_LEN,
decode_backend="triton",
num_q_heads=16,
num_kv_heads=1,
kv_cache_dtype="fp8_e4m3",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
def test_msa_cutlass_decode_static_dispatch_accepts_fp8_alias() -> None:
assert should_prepare_decode_metadata(
32,
DEFAULT_QUERY_LEN,
decode_backend="cutlass",
num_q_heads=16,
num_kv_heads=1,
kv_cache_dtype="fp8",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
@pytest.mark.parametrize(
"kv_cache_dtype",
["auto", "bfloat16", "float16", "fp8_e5m2"],
)
def test_msa_cutlass_decode_static_dispatch_requires_fp8_e4m3(
kv_cache_dtype: str,
) -> None:
assert not should_prepare_decode_metadata(
32,
DEFAULT_QUERY_LEN,
decode_backend="cutlass",
num_q_heads=16,
num_kv_heads=1,
kv_cache_dtype=kv_cache_dtype,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
def test_msa_cutlass_decode_static_dispatch_requires_sm100(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(
current_platform,
"is_device_capability_family",
lambda _: False,
)
assert not should_prepare_decode_metadata(
32,
DEFAULT_QUERY_LEN,
decode_backend="cutlass",
num_q_heads=16,
num_kv_heads=1,
kv_cache_dtype="fp8",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
@pytest.mark.parametrize(
("backend", "expected"),
[
(AttentionBackendEnum.CUTLASS_MSA, "cutlass"),
(AttentionBackendEnum.TRITON_MSA, "triton"),
],
)
def test_msa_attention_backend_alias(
backend: AttentionBackendEnum,
expected: str,
) -> None:
config = AttentionConfig(backend=backend)
assert config.backend is None
assert config.minimax_m3_msa_decode_backend == expected
def test_msa_backend_owns_msa_metadata_builder(
monkeypatch: pytest.MonkeyPatch,
) -> None:
backend_cls, impl_cls = select_main_backend_and_impl_cls(
topk_blocks=TOPK,
kv_cache_dtype="fp8_e4m3",
num_kv_heads=1,
)
assert backend_cls is MiniMaxM3SparseMSABackend
assert impl_cls is MiniMaxM3SparseMSAImpl
monkeypatch.setattr(
current_platform,
"is_device_capability_family",
lambda _: False,
)
backend_cls, impl_cls = select_main_backend_and_impl_cls(
topk_blocks=TOPK,
kv_cache_dtype="fp8_e4m3",
num_kv_heads=1,
)
assert backend_cls is MiniMaxM3SparseBackend
assert impl_cls is MiniMaxM3SparseTritonImpl
def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
monkeypatch: pytest.MonkeyPatch,
) -> None:
batch = 16
block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda")
seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda")
base_metadata = SimpleNamespace(
num_decodes=batch,
decode=SimpleNamespace(
block_table=block_table,
seq_lens=seq_lens,
decode_query_len=1,
),
)
monkeypatch.setattr(
MiniMaxM3SparseMetadataBuilder,
"build",
lambda *args, **kwargs: base_metadata,
)
expected_metadata = object()
monkeypatch.setattr(
sparse_attention_msa_module,
"prepare_decode_metadata",
lambda *args, **kwargs: expected_metadata,
)
builder = object.__new__(MiniMaxM3SparseMSAMetadataBuilder)
builder.num_q_heads = 64
builder.num_kv_heads = 4
builder.topk_blocks = TOPK
builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4)
builder.kv_cache_dtype = "fp8_e4m3"
builder.decode_backend = "cutlass"
builder.msa_cutlass_plan_cache = object()
metadata = builder.build(
0,
SimpleNamespace(
seq_lens_cpu_upper_bound=torch.full((batch,), 257, dtype=torch.int32)
),
)
assert isinstance(metadata.decode, MiniMaxM3SparseMSADecodeMetadata)
assert metadata.decode.decode_query_len == 1
assert metadata.decode.msa_cutlass is expected_metadata
def test_msa_cutlass_plan_cache_keys_query_len(
monkeypatch: pytest.MonkeyPatch,
) -> None:
batch = 16
block_table = torch.zeros(batch, 3, dtype=torch.int32, device="cuda")
seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda")
seq_lens_cpu = torch.full((batch,), 257, dtype=torch.int32)
plan_cache = MSACutlassDecodePlanCache()
built_query_lens: list[int] = []
def fake_build_plan(**kwargs):
query_len = kwargs["decode_query_len"]
built_query_lens.append(query_len)
num_rows = batch * query_len
return (
None,
None,
None,
{
"kv_segment_lens": torch.empty(
num_rows, dtype=torch.int32, device="cuda"
),
"qo_offset": torch.empty(num_rows, dtype=torch.int32, device="cuda"),
},
)
monkeypatch.setattr(plan_cache, "_build_plan", fake_build_plan)
first = prepare_decode_metadata(
block_table,
seq_lens,
seq_lens_cpu,
1,
num_q_heads=64,
num_kv_heads=4,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
repeated = prepare_decode_metadata(
block_table,
seq_lens,
seq_lens_cpu,
1,
num_q_heads=64,
num_kv_heads=4,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
different = prepare_decode_metadata(
block_table,
seq_lens,
seq_lens_cpu,
2,
num_q_heads=64,
num_kv_heads=4,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
current_platform.synchronize()
assert first.plan is repeated.plan
assert different.plan is not first.plan
assert built_query_lens == [1, 2]
def _make_topk(
seq_lens: list[int],
num_kv_heads: int,
query_len: int,
) -> torch.Tensor:
topk = torch.full(
(len(seq_lens) * query_len, num_kv_heads, TOPK),
-1,
dtype=torch.int32,
device="cuda",
)
for request, seq_len in enumerate(seq_lens):
for local_query in range(query_len):
token = request * query_len + local_query
visible_tokens = seq_len - query_len + local_query + 1
visible_pages = math.ceil(visible_tokens / BLOCK_SIZE)
topk[token, :, :visible_pages] = torch.arange(
visible_pages, dtype=torch.int32, device="cuda"
)
return topk
@pytest.mark.parametrize(
(
"num_q_heads",
"num_kv_heads",
"num_request_pairs",
"query_len",
"capture_graph",
),
[
pytest.param(64, 4, 8, 1, True, id="tp1-query-len-1"),
pytest.param(64, 4, 8, 2, False, id="tp1-query-len-2"),
pytest.param(64, 4, 8, 3, False, id="tp1-query-len-3"),
pytest.param(64, 4, 8, 4, True, id="tp1-query-len-4"),
pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
pytest.param(16, 1, 8, 1, True, id="tp4-query-len-1"),
pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
],
)
def test_msa_cutlass_decode_matches_triton_with_interleaved_cache(
num_q_heads: int,
num_kv_heads: int,
num_request_pairs: int,
query_len: int,
capture_graph: bool,
) -> None:
torch.manual_seed(0)
seq_lens_list = [257, 513] * num_request_pairs
seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32)
seq_lens = seq_lens_cpu.cuda()
pages_per_request = [math.ceil(seq_len / BLOCK_SIZE) for seq_len in seq_lens_list]
num_pages = sum(pages_per_request)
max_pages = max(pages_per_request)
block_table = torch.zeros(
len(seq_lens_list), max_pages, dtype=torch.int32, device="cuda"
)
physical_pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda")
offset = 0
for request, request_pages in enumerate(pages_per_request):
block_table[request, :request_pages] = physical_pages[
offset : offset + request_pages
]
offset += request_pages
key = (
torch.randn(
num_pages,
num_kv_heads,
BLOCK_SIZE,
HEAD_DIM,
dtype=torch.bfloat16,
device="cuda",
)
* 0.25
).to(torch.float8_e4m3fn)
value = (torch.randn_like(key, dtype=torch.bfloat16) * 0.25).to(torch.float8_e4m3fn)
kv_cache = torch.cat((key, value), dim=-1)
assert kv_cache.stride() == (
num_kv_heads * BLOCK_SIZE * 2 * HEAD_DIM,
BLOCK_SIZE * 2 * HEAD_DIM,
2 * HEAD_DIM,
1,
)
num_query_tokens = len(seq_lens_list) * query_len
query = torch.randn(
num_query_tokens,
num_q_heads,
HEAD_DIM,
dtype=torch.bfloat16,
device="cuda",
)
q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
query_fp8 = torch.empty_like(query, dtype=torch.float8_e4m3fn)
ops.scaled_fp8_quant(
query.view(num_query_tokens, -1),
scale=q_scale,
output=query_fp8.view(num_query_tokens, -1),
)
query_dequantized = query_fp8.to(torch.bfloat16) * q_scale
topk_token_major = _make_topk(seq_lens_list, num_kv_heads, query_len)
expected = torch.empty_like(query)
minimax_m3_sparse_attn_decode(
query_dequantized,
kv_cache,
topk_token_major.transpose(0, 1),
block_table,
seq_lens,
num_kv_heads,
SM_SCALE,
expected,
query_len,
k_scale=None,
v_scale=None,
)
plan_cache = MSACutlassDecodePlanCache()
metadata = prepare_decode_metadata(
block_table,
seq_lens,
seq_lens_cpu,
query_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
assert metadata.page_table.data_ptr() == block_table.data_ptr()
actual = torch.empty_like(query)
msa_cutlass_sparse_decode(
query_fp8,
kv_cache,
topk_token_major,
actual,
metadata,
scale=SM_SCALE,
q_scale_float=1.0,
k_scale_float=1.0,
v_scale_float=1.0,
)
torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
# The same captured plan must remain correct as ragged lengths change.
updated_seq_lens_list = [129, 385] * num_request_pairs
seq_lens.copy_(
torch.tensor(updated_seq_lens_list, dtype=torch.int32, device="cuda")
)
updated_seq_lens_cpu = torch.tensor(updated_seq_lens_list, dtype=torch.int32)
topk_token_major.copy_(_make_topk(updated_seq_lens_list, num_kv_heads, query_len))
updated_metadata = prepare_decode_metadata(
block_table,
seq_lens,
updated_seq_lens_cpu,
query_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
assert updated_metadata.plan is metadata.plan
assert updated_metadata.page_table.data_ptr() == metadata.page_table.data_ptr()
minimax_m3_sparse_attn_decode(
query_dequantized,
kv_cache,
topk_token_major.transpose(0, 1),
block_table,
seq_lens,
num_kv_heads,
SM_SCALE,
expected,
query_len,
k_scale=None,
v_scale=None,
)
msa_cutlass_sparse_decode(
query_fp8,
kv_cache,
topk_token_major,
actual,
updated_metadata,
scale=SM_SCALE,
q_scale_float=1.0,
k_scale_float=1.0,
v_scale_float=1.0,
)
torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
if capture_graph:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
msa_cutlass_sparse_decode(
query_fp8,
kv_cache,
topk_token_major,
actual,
updated_metadata,
scale=SM_SCALE,
q_scale_float=1.0,
k_scale_float=1.0,
v_scale_float=1.0,
)
replay_seq_lens_list = [257, 513] * num_request_pairs
seq_lens.copy_(
torch.tensor(replay_seq_lens_list, dtype=torch.int32, device="cuda")
)
topk_token_major.copy_(
_make_topk(replay_seq_lens_list, num_kv_heads, query_len)
)
prepare_decode_metadata(
block_table,
seq_lens,
torch.tensor(replay_seq_lens_list, dtype=torch.int32),
query_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
plan_cache=plan_cache,
)
minimax_m3_sparse_attn_decode(
query_dequantized,
kv_cache,
topk_token_major.transpose(0, 1),
block_table,
seq_lens,
num_kv_heads,
SM_SCALE,
expected,
query_len,
k_scale=None,
v_scale=None,
)
graph.replay()
current_platform.synchronize()
torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+15
View File
@@ -208,6 +208,12 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
# index_q here (de-interleaved from the packed qkv); k/v/index_k stay in
# place inside qkv and are scatter-inserted into the caches.
q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
q_fp8 = torch.empty(
num_tokens,
qsz,
dtype=torch.float8_e4m3fn,
device=device,
)
index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device)
ops.fused_minimax_m3_qknorm_rope_kv_insert(
@@ -231,6 +237,8 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
q_out,
index_q,
kv_cache_dtype,
q_fp8_out=q_fp8,
q_fp8_scale=0.5,
)
# ── 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):
# reference materializes bf16 after the norm (the unfused boundary), so
# rounding-boundary elements can differ by ~1 bf16 ulp.
torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2)
expected_q_fp8 = torch.empty_like(q_fp8)
ops.scaled_fp8_quant(
q_out,
scale=torch.tensor(0.5, dtype=torch.float32, device=device),
output=expected_q_fp8,
)
torch.testing.assert_close(q_fp8, expected_q_fp8, rtol=0, atol=0)
torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2)
torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2)
torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
vllm/_custom_ops.py
+7
View File
@@ -2655,6 +2655,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
index_q_out: torch.Tensor | None = None,
kv_cache_dtype: str = "auto",
skip_index_branch: bool = False,
q_fp8_out: torch.Tensor | None = None,
q_fp8_scale: float = 1.0,
) -> None:
"""Fused MiniMax-M3 attention pre-processing (in-place).
@@ -2677,6 +2679,9 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
callers skip a separate ``.contiguous()`` copy before the SM100 sparse
attention's flat TMA descriptor.
If ``q_fp8_out`` is given, the same normalized q is also written in FP8
E4M3 using ``q_fp8_scale`` as its dequantization scale.
When ``skip_index_branch`` is true, sparse rows still keep their packed
``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
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(
index_q_out,
kv_cache_dtype,
skip_index_branch,
q_fp8_out,
q_fp8_scale,
)
vllm/config/attention.py
+15
View File
@@ -11,6 +11,7 @@ from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnu
from vllm.v1.attention.backends.registry import AttentionBackendEnum
IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"]
MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
@config
@@ -20,6 +21,9 @@ class AttentionConfig:
backend: AttentionBackendEnum | None = None
"""Attention backend to use. Use "auto" or None for automatic selection."""
minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton"
"""Sparse decode kernel used by the MiniMax M3 MSA backend."""
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",
@@ -99,6 +103,17 @@ class AttentionConfig:
PyTorch >= 2.9, and 128 for encoder-only attention or older PyTorch
versions."""
def __post_init__(self) -> None:
msa_aliases: dict[AttentionBackendEnum, MiniMaxM3MSADecodeBackend] = {
AttentionBackendEnum.CUTLASS_MSA: "cutlass",
AttentionBackendEnum.TRITON_MSA: "triton",
}
if self.backend in msa_aliases:
self.minimax_m3_msa_decode_backend = msa_aliases[self.backend]
# The alias selects only MiniMax's sparse decode kernel. Dense
# layers still use the platform's normal automatic backend.
self.backend = None
def compute_hash(self) -> str:
"""
Provide a hash that uniquely identifies all the configs
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+2
View File
@@ -23,6 +23,8 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
query: torch.Tensor,
kv_cache: torch.Tensor,
output: torch.Tensor,
*,
query_fp8: torch.Tensor | None = None,
) -> torch.Tensor:
from vllm.models.minimax_m3.amd.ops.sparse_pa import (
minimax_m3_sparse_attn_decode_aiter,
vllm/models/minimax_m3/common/sparse_attention.py
+30 -7
View File
@@ -330,6 +330,7 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
*,
topk_blocks: int,
sparse_block_size: int,
msa_decode_backend: str = "triton",
) -> None:
self.num_heads = num_heads
self.head_size = head_size
@@ -355,6 +356,8 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
query: torch.Tensor,
kv_cache: torch.Tensor,
output: torch.Tensor,
*,
query_fp8: torch.Tensor | None = None,
) -> torch.Tensor:
"""Attend the queries to the indexer-selected blocks. Per kernel.
@@ -364,6 +367,9 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
"""
raise NotImplementedError
def should_use_msa_decode(self, layer_name: str) -> bool:
return False
class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
"""Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode."""
@@ -374,6 +380,8 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
query: torch.Tensor,
kv_cache: torch.Tensor,
output: torch.Tensor,
*,
query_fp8: torch.Tensor | None = None,
) -> torch.Tensor:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
@@ -442,18 +450,18 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
return output
def select_main_impl_cls(
def select_main_backend_and_impl_cls(
*,
topk_blocks: int,
kv_cache_dtype: str,
num_kv_heads: int,
) -> type[MiniMaxM3SparseImpl]:
"""Pick the main attend impl off the main KV-cache dtype.
) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]:
"""Pick the main attention backend and implementation.
Blackwell (SM100) uses the MSA attend for supported top-k block counts
when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA
with shuffle KV cache layout; Other platforms and FP8 E5M2 fall
back to Triton. The MSA modules are imported lazily avoid import errors
back to Triton. The MSA modules are imported lazily to avoid import errors
on unsupported platforms.
"""
use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(num_kv_heads)
@@ -477,11 +485,26 @@ def select_main_impl_cls(
MiniMaxM3SparseAiterPAImpl,
)
return MiniMaxM3SparseAiterPAImpl
return MiniMaxM3SparseBackend, MiniMaxM3SparseAiterPAImpl
if use_msa:
from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
MiniMaxM3SparseMSABackend,
MiniMaxM3SparseMSAImpl,
)
return MiniMaxM3SparseMSAImpl
return MiniMaxM3SparseTritonImpl
return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl
return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl
def select_main_impl_cls(
*,
topk_blocks: int,
kv_cache_dtype: str,
num_kv_heads: int,
) -> type[MiniMaxM3SparseImpl]:
"""Backward-compatible implementation-only selector."""
return select_main_backend_and_impl_cls(
topk_blocks=topk_blocks,
kv_cache_dtype=kv_cache_dtype,
num_kv_heads=num_kv_heads,
)[1]
vllm/models/minimax_m3/nvidia/model.py
+28 -6
View File
@@ -79,7 +79,7 @@ from vllm.models.minimax_m3.common.mm_preprocess import (
from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseBackend,
MiniMaxM3SparseImpl,
select_main_impl_cls,
select_main_backend_and_impl_cls,
)
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -505,15 +505,15 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
# the attend impl reads them back (so nothing crosses the eager break as a
# Python value, which would freeze at capture).
self.topk_indices_buffer = topk_indices_buffer
self.attn_backend = MiniMaxM3SparseBackend
# Indexer (top-k selection) and main attention are separate impls, each
# picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
# (broader than the AttentionImpl that AttentionLayerBase annotates).
self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
self.attn_backend, impl_cls = select_main_backend_and_impl_cls(
topk_blocks=sparse_cfg["sparse_topk_blocks"],
kv_cache_dtype=self.kv_cache_dtype,
num_kv_heads=self.num_kv_heads,
)(
)
self.impl: MiniMaxM3SparseImpl = impl_cls( # type: ignore[assignment]
self.num_heads,
self.head_dim,
self.scaling,
@@ -521,6 +521,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
kv_cache_dtype=self.kv_cache_dtype,
topk_blocks=sparse_cfg["sparse_topk_blocks"],
sparse_block_size=sparse_cfg["sparse_block_size"],
msa_decode_backend=(
vllm_config.attention_config.minimax_m3_msa_decode_backend
),
)
# Self-contained nn.Module: owns its side cache, selects its impl in init.
self.indexer = MiniMaxM3Indexer(
@@ -594,6 +597,16 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
main_slot_mapping = fwd_slot_mapping[self.layer_name]
index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
q = qkv.new_empty((num_tokens, self.q_size))
use_msa_decode = self.impl.should_use_msa_decode(self.layer_name)
query_fp8 = (
torch.empty(
(num_tokens, self.q_size),
dtype=torch.float8_e4m3fn,
device=qkv.device,
)
if use_msa_decode
else None
)
# index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
# the fused kernel emits fp8 directly when this buffer is e4m3.
index_q = qkv.new_empty(
@@ -621,10 +634,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
q,
index_q,
self.kv_cache_dtype,
q_fp8_out=query_fp8,
q_fp8_scale=self._q_scale_float,
)
output = torch.empty_like(q)
attn_output = self._run_attention(q, index_q, output)
attn_output = self._run_attention(q, query_fp8, index_q, output)
output, _ = self.o_proj(attn_output)
return output
@@ -632,6 +647,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
def _run_attention(
self,
query: torch.Tensor,
query_fp8: torch.Tensor | None,
index_query: torch.Tensor,
output: torch.Tensor,
) -> torch.Tensor:
@@ -639,7 +655,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
# metadata and can't be captured into a cudagraph. The indexer writes its
# top-k into the shared ``topk_indices_buffer``; the attend reads it back.
self.indexer(index_query)
return self.impl.forward(self, query, self.kv_cache, output)
return self.impl.forward(
self,
query,
self.kv_cache,
output,
query_fp8=query_fp8,
)
class MiniMaxM3DecoderLayer(nn.Module):
vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
+316
View File
@@ -0,0 +1,316 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MiniMax CUTLASS sparse decode using per-query-token page indices."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import torch
from vllm.config.attention import MiniMaxM3MSADecodeBackend
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
_MAX_NUM_Q_HEADS = 64
_MAX_NUM_KV_HEADS = 4
_HEAD_DIM = 128
_PAGE_SIZE = 128
_TOPK = 16
# fmha_sm100 plans one row per query head. Keep every cached plan within the
# fixed planner allocation used by the MSA decode kernel.
_MAX_QUERY_HEAD_ROWS = 65536
_MAX_DECODE_QUERY_LEN = 32
# Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4.
_MIN_CUTLASS_BATCH_SIZE = 16
@dataclass
class MSACutlassDecodeMetadata:
plan: Any
page_table: torch.Tensor
@triton.jit
def _update_runtime_metadata_kernel(
seq_lens_ptr,
kv_segment_lens_ptr,
qo_offset_ptr,
num_rows: tl.constexpr,
decode_query_len: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < num_rows
request = offsets // decode_query_len
local_query = offsets % decode_query_len
seq_len = tl.load(seq_lens_ptr + request, mask=mask)
tl.store(kv_segment_lens_ptr + offsets, seq_len, mask=mask)
tl.store(
qo_offset_ptr + offsets,
seq_len - decode_query_len + local_query,
mask=mask,
)
@dataclass
class MSACutlassDecodePlanCache:
"""Reusable plans whose mutable tensors retain cudagraph-stable addresses."""
plans: dict[tuple[int, ...], Any] = field(init=False, default_factory=dict)
def _build_plan(
self,
*,
batch: int,
decode_query_len: int,
page_table_stride: int,
initial_seq_lens_cpu: torch.Tensor,
device: torch.device,
num_q_heads: int,
num_kv_heads: int,
page_size: int,
topk_blocks: int,
) -> Any:
from vllm.third_party.fmha_sm100.api import fmha_sm100_plan
qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32)
kv_lens_cpu = initial_seq_lens_cpu
plan = fmha_sm100_plan(
qo_lens_cpu,
kv_lens_cpu,
num_q_heads,
num_kv_heads=num_kv_heads,
qo_offset=kv_lens_cpu - qo_lens_cpu,
page_size=page_size,
output_maxscore=False,
kv_block_num=topk_blocks,
causal=True,
sparse_kernel_mode="decode",
use_fp8_kvcache=True,
split_prefill_decode=False,
device=device,
)
plan_info = plan[3]
row_starts = (
torch.arange(batch, dtype=torch.int32, device=device)
.mul_(page_table_stride)
.repeat_interleave(decode_query_len)
)
page_indptr = torch.cat(
(
row_starts,
torch.tensor(
[batch * page_table_stride],
dtype=torch.int32,
device=device,
),
)
)
plan_info["kv_page_indptr"].copy_(page_indptr)
return plan
def prepare(
self,
block_table: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
decode_query_len: int,
*,
num_q_heads: int,
num_kv_heads: int,
page_size: int,
topk_blocks: int,
) -> MSACutlassDecodeMetadata:
batch = int(seq_lens.shape[0])
if (
block_table.device.type != "cuda"
or block_table.dtype != torch.int32
or not block_table.is_contiguous()
or block_table.shape[0] != batch
):
raise ValueError(
"MSA sparse decode requires a contiguous CUDA int32 block "
"table with one row per request"
)
if (
seq_lens.dtype != torch.int32
or not seq_lens.is_contiguous()
or seq_lens.device != block_table.device
):
raise ValueError(
"MSA sparse decode requires contiguous CUDA int32 sequence "
"lengths on the block table device"
)
if (
seq_lens_cpu.device.type != "cpu"
or seq_lens_cpu.dtype != torch.int32
or not seq_lens_cpu.is_contiguous()
or seq_lens_cpu.shape != seq_lens.shape
):
raise ValueError(
"MSA sparse decode requires contiguous CPU int32 sequence "
"lengths matching the device sequence lengths"
)
page_table_stride = int(block_table.stride(0))
key = (
batch,
decode_query_len,
page_table_stride,
num_q_heads,
num_kv_heads,
page_size,
topk_blocks,
)
plan = self.plans.get(key)
if plan is None:
plan = self._build_plan(
batch=batch,
decode_query_len=decode_query_len,
page_table_stride=page_table_stride,
initial_seq_lens_cpu=seq_lens_cpu,
device=seq_lens.device,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=page_size,
topk_blocks=topk_blocks,
)
self.plans[key] = plan
plan_info = plan[3]
num_rows = batch * decode_query_len
_update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)](
seq_lens,
plan_info["kv_segment_lens"],
plan_info["qo_offset"],
num_rows=num_rows,
decode_query_len=decode_query_len,
BLOCK_SIZE=128,
)
return MSACutlassDecodeMetadata(
plan=plan,
page_table=block_table.view(-1),
)
def _supported_head_geometry(num_q_heads: int, num_kv_heads: int) -> bool:
return (
0 < num_q_heads <= _MAX_NUM_Q_HEADS
and 0 < num_kv_heads <= _MAX_NUM_KV_HEADS
and num_q_heads % num_kv_heads == 0
)
def supports_cutlass_sparse_decode(
*,
decode_backend: MiniMaxM3MSADecodeBackend,
num_q_heads: int,
num_kv_heads: int,
kv_cache_dtype: str,
page_size: int,
topk_blocks: int,
) -> bool:
"""Return whether static model geometry supports CUTLASS sparse decode."""
return (
decode_backend == "cutlass"
and current_platform.is_cuda()
and current_platform.is_device_capability_family(100)
and kv_cache_dtype in ("fp8", "fp8_e4m3")
and _supported_head_geometry(num_q_heads, num_kv_heads)
and page_size == _PAGE_SIZE
and topk_blocks == _TOPK
)
def should_prepare_decode_metadata(
batch_size: int,
decode_query_len: int,
*,
decode_backend: MiniMaxM3MSADecodeBackend,
num_q_heads: int,
num_kv_heads: int,
kv_cache_dtype: str,
page_size: int,
topk_blocks: int,
) -> bool:
"""Return whether a graph shape can use the CUTLASS decode path."""
total_q = batch_size * decode_query_len
return (
supports_cutlass_sparse_decode(
decode_backend=decode_backend,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
kv_cache_dtype=kv_cache_dtype,
page_size=page_size,
topk_blocks=topk_blocks,
)
and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
and batch_size >= _MIN_CUTLASS_BATCH_SIZE
and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
)
@torch.no_grad()
def prepare_decode_metadata(
block_table: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
decode_query_len: int,
*,
num_q_heads: int,
num_kv_heads: int,
page_size: int,
topk_blocks: int,
plan_cache: MSACutlassDecodePlanCache | None = None,
) -> MSACutlassDecodeMetadata:
"""Prepare graph-stable runtime metadata for one sparse decode step."""
cache = plan_cache or MSACutlassDecodePlanCache()
return cache.prepare(
block_table,
seq_lens,
seq_lens_cpu,
decode_query_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=page_size,
topk_blocks=topk_blocks,
)
@torch.no_grad()
def msa_cutlass_sparse_decode(
query_fp8: torch.Tensor,
kv_cache: torch.Tensor,
topk: torch.Tensor,
output: torch.Tensor,
metadata: MSACutlassDecodeMetadata,
*,
scale: float,
q_scale_float: float,
k_scale_float: float,
v_scale_float: float,
) -> None:
"""Run CUTLASS sparse decode with metadata prepared by the MSA builder."""
key, value = kv_cache.split(_HEAD_DIM, dim=-1)
from vllm.third_party.fmha_sm100.api import fmha_sm100
fmha_sm100(
query_fp8,
key,
value,
metadata.plan,
kv_indices=metadata.page_table,
kv_block_indexes=topk,
out=output,
output_maxscore=False,
output_o=True,
sm_scale=scale,
q_scale=q_scale_float,
k_scale=k_scale_float,
v_scale=v_scale_float,
o_scale=1.0,
)
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+210 -19
View File
@@ -1,28 +1,198 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MSA (SM100/Blackwell) block-sparse attend for MiniMax M3.
"""MSA (SM100/Blackwell) block-sparse attention for MiniMax M3.
Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``);
decode falls back to the Triton split-K kernel (no MSA decode yet). ``fmha_sm100``
imports are function-local, so this module is import-safe on AMD/non-SM100.
Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``).
Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100``
path for regular decode and speculative verification.
"""
from dataclasses import dataclass
import torch
from vllm.config import VllmConfig
from vllm.config.attention import MiniMaxM3MSADecodeBackend
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.models.minimax_m3.common.ops.sparse_attn import (
SPARSE_BLOCK_SIZE,
minimax_m3_sparse_attn_decode,
)
from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseBackend,
MiniMaxM3SparseDecodeMetadata,
MiniMaxM3SparseImpl,
MiniMaxM3SparseMetadata,
MiniMaxM3SparseMetadataBuilder,
)
from vllm.v1.attention.backend import AttentionLayer
from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
MSACutlassDecodeMetadata,
MSACutlassDecodePlanCache,
msa_cutlass_sparse_decode,
prepare_decode_metadata,
should_prepare_decode_metadata,
supports_cutlass_sparse_decode,
)
from vllm.v1.attention.backend import (
AttentionLayer,
CommonAttentionMetadata,
)
from vllm.v1.kv_cache_interface import AttentionSpec
logger = init_logger(__name__)
class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
"""MiniMax M3 backend with NVIDIA MSA-specific decode metadata."""
@staticmethod
def get_builder_cls() -> type["MiniMaxM3SparseMSAMetadataBuilder"]:
return MiniMaxM3SparseMSAMetadataBuilder
class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend):
"""Attention-backend alias selecting CUTLASS MSA sparse decode."""
@staticmethod
def get_name() -> str:
return "CUTLASS_MSA"
class MiniMaxM3SparseTritonBackend(MiniMaxM3SparseMSABackend):
"""Attention-backend alias selecting Triton MSA sparse decode."""
@staticmethod
def get_name() -> str:
return "TRITON_MSA"
@dataclass
class MiniMaxM3SparseMSADecodeMetadata(MiniMaxM3SparseDecodeMetadata):
msa_cutlass: MSACutlassDecodeMetadata | None = None
class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
"""Prepare MSA plans only for decode shapes supported by ``fmha_sm100``."""
def __init__(
self,
kv_cache_spec: AttentionSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
) -> None:
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
config = vllm_config.model_config.hf_text_config
tp_size = vllm_config.parallel_config.tensor_parallel_size
self.num_q_heads = config.num_attention_heads // tp_size
self.num_kv_heads = kv_cache_spec.num_kv_heads
self.topk_blocks = config.sparse_attention_config["sparse_topk_blocks"]
# AttentionSpec stores every FP8 mode as uint8, so retain the configured
# format to distinguish E4M3 (supported) from E5M2 before planning.
self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend
self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache()
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> MiniMaxM3SparseMetadata:
metadata = super().build(
common_prefix_len,
common_attn_metadata,
fast_build,
)
decode = metadata.decode
if decode is None:
return metadata
msa_cutlass = None
if should_prepare_decode_metadata(
metadata.num_decodes,
decode.decode_query_len,
decode_backend=self.decode_backend,
num_q_heads=self.num_q_heads,
num_kv_heads=self.num_kv_heads,
kv_cache_dtype=self.kv_cache_dtype,
page_size=SPARSE_BLOCK_SIZE,
topk_blocks=self.topk_blocks,
):
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
assert seq_lens_cpu is not None
msa_cutlass = prepare_decode_metadata(
decode.block_table,
decode.seq_lens,
seq_lens_cpu[: metadata.num_decodes],
decode.decode_query_len,
num_q_heads=self.num_q_heads,
num_kv_heads=self.num_kv_heads,
page_size=SPARSE_BLOCK_SIZE,
topk_blocks=self.topk_blocks,
plan_cache=self.msa_cutlass_plan_cache,
)
metadata.decode = MiniMaxM3SparseMSADecodeMetadata(
seq_lens=decode.seq_lens,
block_table=decode.block_table,
decode_query_len=decode.decode_query_len,
msa_cutlass=msa_cutlass,
)
return metadata
class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
"""MSA block-sparse attend (``fmha_sm100``); Triton split-K decode."""
"""MSA block-sparse attention with guarded CUTLASS sparse decode."""
def __init__(
self,
num_heads: int,
head_size: int,
scale: float,
num_kv_heads: int | None = None,
kv_cache_dtype: str = "auto",
*,
topk_blocks: int,
sparse_block_size: int,
msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton",
) -> None:
super().__init__(
num_heads,
head_size,
scale,
num_kv_heads,
kv_cache_dtype,
topk_blocks=topk_blocks,
sparse_block_size=sparse_block_size,
)
self.use_cutlass_decode = supports_cutlass_sparse_decode(
decode_backend=msa_decode_backend,
num_q_heads=self.num_heads,
num_kv_heads=self.num_kv_heads,
kv_cache_dtype=self.kv_cache_dtype,
page_size=self.block_size,
topk_blocks=self.topk_blocks,
)
logger.info_once(
"MiniMax M3 MSA sparse decode selected %s",
"CUTLASS" if self.use_cutlass_decode else "Triton",
)
def should_use_msa_decode(self, layer_name: str) -> bool:
if not self.use_cutlass_decode:
return False
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
return False
main_md = attn_metadata[layer_name]
if not isinstance(main_md, MiniMaxM3SparseMetadata):
return False
decode = main_md.decode
return (
isinstance(decode, MiniMaxM3SparseMSADecodeMetadata)
and decode.msa_cutlass is not None
)
def forward(
self,
@@ -30,6 +200,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
query: torch.Tensor,
kv_cache: torch.Tensor,
output: torch.Tensor,
*,
query_fp8: torch.Tensor | None = None,
) -> torch.Tensor:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
@@ -52,23 +224,42 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
# Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
# Decode [:nd]: CUTLASS for planned shapes, otherwise Triton.
if main_md.num_decodes > 0:
d = main_md.decode
assert d is not None
minimax_m3_sparse_attn_decode(
q[:nd],
kv_cache,
topk[:nd].transpose(0, 1),
d.block_table,
d.seq_lens,
self.num_kv_heads,
self.scale,
out[:nd],
d.decode_query_len,
k_scale=k_scale,
v_scale=v_scale,
msa_metadata = (
d.msa_cutlass
if isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
else None
)
if self.use_cutlass_decode and msa_metadata is not None:
assert query_fp8 is not None
msa_cutlass_sparse_decode(
query_fp8[:nd].view(-1, self.num_heads, hd),
kv_cache,
topk[:nd],
out[:nd],
msa_metadata,
scale=self.scale,
q_scale_float=getattr(layer, "_q_scale_float", 1.0),
k_scale_float=getattr(layer, "_k_scale_float", 1.0),
v_scale_float=getattr(layer, "_v_scale_float", 1.0),
)
else:
minimax_m3_sparse_attn_decode(
q[:nd],
kv_cache,
topk[:nd].transpose(0, 1),
d.block_table,
d.seq_lens,
self.num_kv_heads,
self.scale,
out[:nd],
d.decode_query_len,
k_scale=k_scale,
v_scale=v_scale,
)
# Prefill [nd:]: MSA sparse FMHA over the selected blocks.
if main_md.num_prefills > 0:
vllm/v1/attention/backends/registry.py
+8
View File
@@ -101,6 +101,14 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
MINIMAX_M3_SPARSE = (
"vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend"
)
CUTLASS_MSA = (
"vllm.models.minimax_m3.nvidia.sparse_attention_msa."
"MiniMaxM3SparseCutlassBackend"
)
TRITON_MSA = (
"vllm.models.minimax_m3.nvidia.sparse_attention_msa."
"MiniMaxM3SparseTritonBackend"
)
NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
# HPC Attention Backend: