From 60cbbace9de2fb140fd475edc301c2f08a08c673 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 11 Sep 2026 08:18:37 -0700 Subject: [PATCH] feat(minimax-m3): add FlashInfer MSA decode backend Assisted-by: OpenAI Codex Signed-off-by: Shreyas Misra --- ...st_minimax_m3_msa_cutlass_sparse_decode.py | 215 ++++++++++++++++++ vllm/config/attention.py | 3 +- vllm/models/minimax_m3/nvidia/model.py | 5 +- .../nvidia/msa_flashinfer_sparse_decode.py | 156 +++++++++++++ .../minimax_m3/nvidia/sparse_attention_msa.py | 61 ++++- vllm/v1/attention/backends/registry.py | 4 + 6 files changed, 437 insertions(+), 7 deletions(-) create mode 100644 vllm/models/minimax_m3/nvidia/msa_flashinfer_sparse_decode.py diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index f6ea80cca020..112d205eed1d 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -31,6 +31,10 @@ prepare_decode_metadata, should_prepare_decode_metadata, ) +from vllm.models.minimax_m3.nvidia.msa_flashinfer_sparse_decode import ( + msa_flashinfer_sparse_decode, + supports_flashinfer_sparse_decode, +) from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSADecodeMetadata, @@ -164,10 +168,42 @@ def test_msa_cutlass_decode_static_dispatch_requires_sm100( ) +@pytest.mark.parametrize( + ("decode_backend", "num_q_heads", "num_kv_heads", "kv_cache_dtype", "expected"), + [ + ("flashinfer", 16, 1, "fp8_e4m3", True), + ("flashinfer", 16, 1, "fp8", True), + ("triton", 16, 1, "fp8_e4m3", False), + ("flashinfer", 16, 1, "bfloat16", False), + ("flashinfer", 17, 1, "fp8_e4m3", False), + ("flashinfer", 16, 3, "fp8_e4m3", False), + ], +) +def test_msa_flashinfer_decode_static_dispatch( + decode_backend: str, + num_q_heads: int, + num_kv_heads: int, + kv_cache_dtype: str, + expected: bool, +) -> None: + assert ( + supports_flashinfer_sparse_decode( + decode_backend=decode_backend, # type: ignore[arg-type] + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + kv_cache_dtype=kv_cache_dtype, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + is expected + ) + + @pytest.mark.parametrize( ("backend", "expected"), [ (AttentionBackendEnum.CUTLASS_MSA, "cutlass"), + (AttentionBackendEnum.FLASHINFER_MSA, "flashinfer"), (AttentionBackendEnum.TRITON_MSA, "triton"), ], ) @@ -418,6 +454,185 @@ def _make_topk( return topk +def _reference_sparse_attention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + q2k_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens: list[int], + query_len: int, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(query) + lse = torch.empty(query.shape[:2], dtype=torch.float32, device=query.device) + group_size = query.shape[1] // key_cache.shape[1] + for request, seq_len in enumerate(seq_lens): + for local_query in range(query_len): + token = request * query_len + local_query + query_position = seq_len - query_len + local_query + for query_head in range(query.shape[1]): + kv_head = query_head // group_size + keys = [] + values = [] + for logical_page in q2k_indices[kv_head, token].tolist(): + if logical_page < 0: + break + physical_page = int(block_table[request, logical_page]) + page_start = logical_page * BLOCK_SIZE + page_end = min( + page_start + BLOCK_SIZE, + seq_len, + query_position + 1, + ) + if page_end <= page_start: + continue + page_tokens = page_end - page_start + keys.append(key_cache[physical_page, kv_head, :page_tokens]) + values.append(value_cache[physical_page, kv_head, :page_tokens]) + selected_keys = torch.cat(keys).to(torch.float32) + selected_values = torch.cat(values).to(torch.float32) + scores = ( + selected_keys @ query[token, query_head].to(torch.float32) + ) * SM_SCALE + probabilities = torch.softmax(scores, dim=0) + output[token, query_head] = probabilities @ selected_values + lse[token, query_head] = torch.logsumexp(scores, dim=0) + return output, lse + + +@pytest.mark.parametrize("query_len", [1, 4]) +def test_msa_flashinfer_decode_matches_triton_and_lse_reference( + query_len: int, +) -> None: + torch.manual_seed(1) + num_q_heads = 16 + num_kv_heads = 4 + seq_lens_list = [257, 513] + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") + pages_per_request = [math.ceil(length / BLOCK_SIZE) for length 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_cache = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.25 + ).to(torch.float8_e4m3fn) + value_cache = (torch.randn_like(key_cache, dtype=torch.bfloat16) * 0.25).to( + torch.float8_e4m3fn + ) + assert key_cache.is_contiguous() and value_cache.is_contiguous() + + total_q = len(seq_lens_list) * query_len + query = torch.randn( + total_q, + num_q_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + query_scale_float = 0.75 + query_scale = torch.tensor(query_scale_float, dtype=torch.float32, device="cuda") + query_fp8 = torch.empty_like(query, dtype=torch.float8_e4m3fn) + ops.scaled_fp8_quant( + query.view(total_q, -1), + scale=query_scale, + output=query_fp8.view(total_q, -1), + ) + query_dequantized = query_fp8.to(torch.bfloat16) * query_scale + topk_token_major = _make_topk(seq_lens_list, num_kv_heads, query_len) + q2k_indices = topk_token_major.transpose(0, 1) + assert not q2k_indices.is_contiguous() + + packed_cache = torch.cat((key_cache, value_cache), dim=-1) + triton_output = torch.empty_like(query) + minimax_m3_sparse_attn_decode( + query_dequantized, + packed_cache, + q2k_indices, + block_table, + seq_lens, + num_kv_heads, + SM_SCALE, + triton_output, + query_len, + k_scale=None, + v_scale=None, + ) + reference_output, reference_lse = _reference_sparse_attention( + query_dequantized, + key_cache, + value_cache, + q2k_indices, + block_table, + seq_lens_list, + query_len, + ) + + provided_output = torch.empty_like(query) + provided_lse = torch.empty( + query.shape[0], query.shape[1], dtype=torch.float32, device=query.device + ) + capture_stream = torch.cuda.Stream() + with torch.cuda.stream(capture_stream): + actual, actual_lse = msa_flashinfer_sparse_decode( + query_fp8, + packed_cache, + q2k_indices, + block_table, + seq_lens, + query_len, + scale=SM_SCALE, + q_scale_float=query_scale_float, + out=provided_output, + lse_out=provided_lse, + ) + capture_stream.synchronize() + + assert actual.data_ptr() == provided_output.data_ptr() + assert actual_lse.data_ptr() == provided_lse.data_ptr() + torch.testing.assert_close(actual, triton_output, atol=0.02, rtol=0.02) + torch.testing.assert_close(actual, reference_output, atol=0.02, rtol=0.02) + torch.testing.assert_close(actual_lse, reference_lse, atol=0.02, rtol=0.02) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + captured_output, captured_lse = msa_flashinfer_sparse_decode( + query_fp8, + packed_cache, + q2k_indices, + block_table, + seq_lens, + query_len, + scale=SM_SCALE, + q_scale_float=query_scale_float, + out=provided_output, + lse_out=provided_lse, + ) + assert captured_output.data_ptr() == actual.data_ptr() + assert captured_lse.data_ptr() == actual_lse.data_ptr() + graph.replay() + current_platform.synchronize() + torch.testing.assert_close(captured_output, reference_output, atol=0.02, rtol=0.02) + torch.testing.assert_close(captured_lse, reference_lse, atol=0.02, rtol=0.02) + + @pytest.mark.parametrize( ( "num_q_heads", diff --git a/vllm/config/attention.py b/vllm/config/attention.py index d5a150c92a31..884da6736f49 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -14,7 +14,7 @@ logger = init_logger(__name__) IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] -MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] +MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass", "flashinfer"] @config @@ -108,6 +108,7 @@ class AttentionConfig: def __post_init__(self) -> None: msa_aliases: dict[AttentionBackendEnum, MiniMaxM3MSADecodeBackend] = { AttentionBackendEnum.CUTLASS_MSA: "cutlass", + AttentionBackendEnum.FLASHINFER_MSA: "flashinfer", AttentionBackendEnum.TRITON_MSA: "triton", } if self.backend in msa_aliases: diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 970b06d407ea..c60fe07b9a63 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -565,7 +565,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: ) def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None: - if not getattr(self.impl, "use_cutlass_decode", False): + if not ( + getattr(self.impl, "use_cutlass_decode", False) + or getattr(self.impl, "use_flashinfer_decode", False) + ): return None return torch.empty( (qkv.shape[0], self.q_size), diff --git a/vllm/models/minimax_m3/nvidia/msa_flashinfer_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_flashinfer_sparse_decode.py new file mode 100644 index 000000000000..169521978b0b --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/msa_flashinfer_sparse_decode.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer sparse decode adapter for MiniMax M3 on SM100.""" + +from __future__ import annotations + +import torch + +from vllm.config.attention import MiniMaxM3MSADecodeBackend +from vllm.platforms import current_platform + +_HEAD_DIM = 128 +_PAGE_SIZE = 128 +_TOPK = 16 + + +def supports_flashinfer_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 geometry supports packed FlashInfer MSA decode.""" + + return ( + decode_backend == "flashinfer" + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and kv_cache_dtype in ("fp8", "fp8_e4m3") + and 0 < num_kv_heads <= num_q_heads + and num_q_heads % num_kv_heads == 0 + and num_q_heads // num_kv_heads <= 16 + and page_size == _PAGE_SIZE + and topk_blocks == _TOPK + ) + + +def _validate_inputs( + query: torch.Tensor, + kv_cache: torch.Tensor, + q2k_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, +) -> None: + if query.ndim != 3 or query.shape[-1] != _HEAD_DIM: + raise ValueError("query must have shape (total_q, num_q_heads, 128)") + if query.dtype != torch.float8_e4m3fn: + raise ValueError("FlashInfer packed MSA decode requires an FP8 E4M3 query") + if query.shape[0] % decode_query_len != 0: + raise ValueError("query rows must be divisible by decode_query_len") + if kv_cache.ndim != 4 or kv_cache.shape[2:] != ( + _PAGE_SIZE, + 2 * _HEAD_DIM, + ): + raise ValueError("kv_cache must have shape (num_pages, num_kv_heads, 128, 256)") + if kv_cache.dtype != torch.float8_e4m3fn or not kv_cache.is_contiguous(): + raise ValueError("kv_cache must be a contiguous FP8 E4M3 packed cache") + + total_q, num_q_heads, _ = query.shape + num_kv_heads = kv_cache.shape[1] + if num_q_heads % num_kv_heads != 0 or num_q_heads // num_kv_heads > 16: + raise ValueError("FlashInfer MSA requires a GQA group size of at most 16") + if q2k_indices.shape != (num_kv_heads, total_q, _TOPK): + raise ValueError("q2k_indices must have shape (num_kv_heads, total_q, 16)") + if q2k_indices.dtype != torch.int32 or ( + not q2k_indices.is_contiguous() + and q2k_indices.stride() != (16, num_kv_heads * 16, 1) + ): + raise ValueError( + "q2k_indices must be compact head-major int32 or an exact compact " + "token-major transpose" + ) + + batch_size = total_q // decode_query_len + if ( + block_table.ndim != 2 + or block_table.shape[0] != batch_size + or block_table.dtype != torch.int32 + or not block_table.is_contiguous() + ): + raise ValueError( + "block_table must be contiguous int32 with one row per request" + ) + if ( + seq_lens.shape != (batch_size,) + or seq_lens.dtype != torch.int32 + or not seq_lens.is_contiguous() + ): + raise ValueError("seq_lens must be contiguous int32 with one entry per request") + + tensors = ( + kv_cache, + q2k_indices, + block_table, + seq_lens, + ) + if any(tensor.device != query.device for tensor in tensors): + raise ValueError("all FlashInfer MSA inputs must be on the query device") + + +@torch.no_grad() +def msa_flashinfer_sparse_decode( + query: torch.Tensor, + kv_cache: torch.Tensor, + q2k_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + *, + scale: float, + q_scale_float: float = 1.0, + k_scale_float: float = 1.0, + v_scale_float: float = 1.0, + out: torch.Tensor, + lse_out: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run paged MSA decode directly on vLLM's packed K/V cache.""" + if decode_query_len <= 0: + raise ValueError("decode_query_len must be positive") + _validate_inputs( + query, + kv_cache, + q2k_indices, + block_table, + seq_lens, + decode_query_len, + ) + + from flashinfer.msa_ops import msa_sparse_decode_attention + + key_cache, value_cache = kv_cache.split(_HEAD_DIM, dim=-1) + + result = msa_sparse_decode_attention( + query, + key_cache, + value_cache, + q2k_indices, + page_table=block_table, + seqused_k=seq_lens, + seqlen_q=decode_query_len, + causal=True, + softmax_scale=scale * q_scale_float, + return_softmax_lse=True, + k_global_scale=k_scale_float, + v_global_scale=v_scale_float, + force_fused=True, + out=out, + lse_out=lse_out, + ) + assert isinstance(result, tuple) + output, lse = result + return output, lse diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index d29e1346bd27..9377d4705674 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -3,8 +3,8 @@ """MSA (SM100/Blackwell) block-sparse attention for MiniMax M3. 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. +Decode uses Triton split-K by default, with opt-in CUTLASS ``fmha_sm100`` or +FlashInfer MSA paths for regular decode and speculative verification. """ from dataclasses import dataclass @@ -34,6 +34,10 @@ should_prepare_decode_metadata, supports_cutlass_sparse_decode, ) +from vllm.models.minimax_m3.nvidia.msa_flashinfer_sparse_decode import ( + msa_flashinfer_sparse_decode, + supports_flashinfer_sparse_decode, +) from vllm.v1.attention.backend import ( AttentionLayer, CommonAttentionMetadata, @@ -59,6 +63,14 @@ def get_name() -> str: return "CUTLASS_MSA" +class MiniMaxM3SparseFlashInferBackend(MiniMaxM3SparseMSABackend): + """Attention-backend alias selecting FlashInfer MSA sparse decode.""" + + @staticmethod + def get_name() -> str: + return "FLASHINFER_MSA" + + class MiniMaxM3SparseTritonBackend(MiniMaxM3SparseMSABackend): """Attention-backend alias selecting Triton MSA sparse decode.""" @@ -174,13 +186,30 @@ def __init__( page_size=self.block_size, topk_blocks=self.topk_blocks, ) + self.use_flashinfer_decode = supports_flashinfer_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, + ) + selected_backend = ( + "FlashInfer" + if self.use_flashinfer_decode + else "CUTLASS" + if self.use_cutlass_decode + else "Triton" + ) logger.info_once( "MiniMax M3 MSA sparse decode selected %s", - "CUTLASS" if self.use_cutlass_decode else "Triton", + selected_backend, ) def should_use_msa_decode(self, layer_name: str) -> bool: - if not self.use_cutlass_decode: + use_cutlass_decode = getattr(self, "use_cutlass_decode", False) + use_flashinfer_decode = getattr(self, "use_flashinfer_decode", False) + if not (use_cutlass_decode or use_flashinfer_decode): return False attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): @@ -189,6 +218,8 @@ def should_use_msa_decode(self, layer_name: str) -> bool: if not isinstance(main_md, MiniMaxM3SparseMetadata): return False decode = main_md.decode + if use_flashinfer_decode: + return isinstance(decode, MiniMaxM3SparseMSADecodeMetadata) return ( isinstance(decode, MiniMaxM3SparseMSADecodeMetadata) and decode.msa_cutlass is not None @@ -233,7 +264,27 @@ def forward( if isinstance(d, MiniMaxM3SparseMSADecodeMetadata) else None ) - if self.use_cutlass_decode and msa_metadata is not None: + if getattr(self, "use_flashinfer_decode", False): + assert query_fp8 is not None + flashinfer_lse = torch.empty( + (nd, self.num_heads), dtype=torch.float32, device=query.device + ) + flashinfer_output, _ = msa_flashinfer_sparse_decode( + query_fp8[:nd].view(-1, self.num_heads, hd), + kv_cache, + topk[:nd].transpose(0, 1), + d.block_table, + d.seq_lens, + d.decode_query_len, + 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), + out=out[:nd], + lse_out=flashinfer_lse, + ) + assert flashinfer_output.data_ptr() == out[:nd].data_ptr() + elif 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), diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 0c9b39c9886f..adb552202391 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -110,6 +110,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.minimax_m3.nvidia.sparse_attention_msa." "MiniMaxM3SparseCutlassBackend" ) + FLASHINFER_MSA = ( + "vllm.models.minimax_m3.nvidia.sparse_attention_msa." + "MiniMaxM3SparseFlashInferBackend" + ) TRITON_MSA = ( "vllm.models.minimax_m3.nvidia.sparse_attention_msa." "MiniMaxM3SparseTritonBackend"