vllm.models.minimax_m3.nvidia.sparse_attention_msa

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. NVFP4 KV caches read the vLLM packed pages directly in both prefill and CUTLASS decode.

Classes:

• MiniMaxM3SparseCutlassBackend –

  Attention-backend alias selecting CUTLASS MSA sparse decode.

• MiniMaxM3SparseMSABackend –

  MiniMax M3 backend with NVIDIA MSA-specific decode metadata.

• MiniMaxM3SparseMSAImpl –

  MSA block-sparse attention with guarded CUTLASS sparse decode.

• MiniMaxM3SparseMSAMetadataBuilder –

  Prepare MSA plans only for decode shapes supported by fmha_sm100.

• MiniMaxM3SparseMSANvfp4Backend –

  MSA backend over HND NVFP4 KV cache pages in per-head K/V slots.

• MiniMaxM3SparseTritonBackend –

  Attention-backend alias selecting Triton MSA sparse decode.

MiniMaxM3SparseCutlassBackend

Bases: MiniMaxM3SparseMSABackend

Attention-backend alias selecting CUTLASS MSA sparse decode.

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend):
    """Attention-backend alias selecting CUTLASS MSA sparse decode."""

    @staticmethod
    def get_name() -> str:
        return "CUTLASS_MSA"

MiniMaxM3SparseMSABackend

Bases: MiniMaxM3SparseBackend

MiniMax M3 backend with NVIDIA MSA-specific decode metadata.

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
    """MiniMax M3 backend with NVIDIA MSA-specific decode metadata."""

    @staticmethod
    def get_builder_cls() -> type["MiniMaxM3SparseMSAMetadataBuilder"]:
        return MiniMaxM3SparseMSAMetadataBuilder

MiniMaxM3SparseMSAImpl

Bases: MiniMaxM3SparseImpl

MSA block-sparse attention with guarded CUTLASS sparse decode.

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
    """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_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype)
        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,
        )
        if self.use_nvfp4_kv and not self.use_cutlass_decode:
            raise ValueError(
                "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode "
                f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, "
                f"block_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,
        layer: AttentionLayer,
        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):
            return output  # profiling run; caches unbound
        main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
        assert isinstance(main_md, MiniMaxM3SparseMetadata)

        nd = main_md.num_decode_tokens
        num_tokens = main_md.num_actual_tokens
        # Indexer top-k from the shared token-major buffer [total_q, H, MK]; the
        # kernels want [H, tokens, MK], so slice tokens on dim 0 then transpose.
        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
        assert topk is not None
        hd = self.head_size
        q = query[:num_tokens].view(-1, self.num_heads, hd)
        out = output[:num_tokens].view(-1, self.num_heads, hd)
        # Given query_fp8, it is the only q the fused insert wrote: CUTLASS
        # decode reads it directly, and the bf16-query kernels read it
        # dequantized into ``query``.
        q_fp8 = (
            None
            if query_fp8 is None
            else query_fp8[:num_tokens].view(-1, self.num_heads, hd)
        )
        q_scale_float = getattr(layer, "_q_scale_float", 1.0)
        if self.use_fp8_kv and not self.use_nvfp4_kv:
            kv_cache = kv_cache.view(self.kv_cache_fp8_dtype)
        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]: CUTLASS for planned shapes, otherwise Triton.
        if main_md.num_decodes > 0:
            d = main_md.decode
            assert d is not None
            msa_metadata = (
                d.msa_cutlass
                if isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
                else None
            )
            if self.use_cutlass_decode and msa_metadata is not None:
                assert q_fp8 is not None
                msa_cutlass_sparse_decode(
                    q_fp8[:nd],
                    kv_cache,
                    topk[:nd],
                    out[:nd],
                    msa_metadata,
                    scale=self.scale,
                    q_scale_float=q_scale_float,
                    k_scale_float=getattr(layer, "_k_scale_float", 1.0),
                    v_scale_float=getattr(layer, "_v_scale_float", 1.0),
                    k_scale=k_scale if self.use_nvfp4_kv else None,
                    v_scale=v_scale if self.use_nvfp4_kv else None,
                )
            elif self.use_nvfp4_kv:
                raise RuntimeError(
                    "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; "
                    f"CUTLASS decode was not planned for {main_md.num_decodes} "
                    f"requests with query length {d.decode_query_len}"
                )
            else:
                minimax_m3_sparse_attn_decode(
                    q[:nd]
                    if q_fp8 is None
                    else _dequantize_query(q_fp8[:nd], q_scale_float, 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:
            from vllm.third_party.fmha_sm100.sparse import (
                build_k2q_csr,
                sparse_atten_func,
                sparse_atten_nvfp4_kv_func,
            )

            p = main_md.prefill
            assert p is not None
            # [H, prefill, MK] transposed view; build_k2q_csr consumes the
            # strided view directly (topK stays innermost-contiguous).
            prefill_topk = topk[nd:num_tokens].transpose(0, 1)
            qp = (
                q[nd:]
                if q_fp8 is None
                else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:])
            )
            k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
                prefill_topk,
                p.cu_seqlens_q,
                p.cu_seqlens_k,
                SPARSE_BLOCK_SIZE,
                total_k=0,
                max_seqlen_k=p.max_seq_len,
                max_seqlen_q=p.max_query_len,
                total_rows=p.total_kv_blocks,
                qhead_per_kv=qp.shape[1] // self.num_kv_heads,
                return_schedule=True,
            )
            if self.use_nvfp4_kv:
                k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
                sparse_atten_nvfp4_kv_func(
                    qp,
                    k_data,
                    v_data,
                    k_sf,
                    v_sf,
                    k_scale,
                    v_scale,
                    k2q_row_ptr,
                    k2q_q_indices,
                    self.topk_blocks,
                    cu_seqlens_q=p.cu_seqlens_q,
                    cu_seqlens_k=p.cu_seqlens_k,
                    max_seqlen_q=p.max_query_len,
                    max_seqlen_k=p.max_seq_len,
                    blk_kv=SPARSE_BLOCK_SIZE,
                    causal=True,
                    softmax_scale=self.scale,
                    page_table=p.block_table,
                    seqused_k=p.seq_lens,
                    schedule=schedule,
                    out=out[nd:],
                    kv_layout="vllm",
                )
                return output
            k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
            sparse_atten_func(
                qp,
                k_cache,
                v_cache,
                k2q_row_ptr,
                k2q_q_indices,
                topK=self.topk_blocks,
                blk_kv=SPARSE_BLOCK_SIZE,
                causal=True,
                softmax_scale=self.scale,
                cu_seqlens_q=p.cu_seqlens_q,
                cu_seqlens_k=p.cu_seqlens_k,
                max_seqlen_q=p.max_query_len,
                max_seqlen_k=p.max_seq_len,
                page_table=p.block_table,
                seqused_k=p.seq_lens,
                schedule=schedule,
                out=out[nd:],
            )
        return output

MiniMaxM3SparseMSAMetadataBuilder

Bases: MiniMaxM3SparseMetadataBuilder

Prepare MSA plans only for decode shapes supported by fmha_sm100.

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

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

MiniMaxM3SparseMSANvfp4Backend

Bases: MiniMaxM3SparseMSABackend

MSA backend over HND NVFP4 KV cache pages in per-head K/V slots.

Methods:

• customize_spec –

  Store each head's K and V as two slots of packed data + block scales

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

class MiniMaxM3SparseMSANvfp4Backend(MiniMaxM3SparseMSABackend):
    """MSA backend over HND NVFP4 KV cache pages in per-head K/V slots."""

    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["nvfp4"]

    @classmethod
    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
        """Store each head's K and V as two slots of packed data + block scales
        (see ``nvfp4_kv_cache_views``)."""
        if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
            return spec
        return replace(
            spec,
            num_head_slots=2 * spec.num_kv_heads,
            state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size)
            * get_dtype_size(spec.dtype),
        )

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
        # Every head's K/V slots are one contiguous run of the page.
        return (KVCacheLayout.LBHNC,)

customize_spec(spec) classmethod

Store each head's K and V as two slots of packed data + block scales (see nvfp4_kv_cache_views).

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@classmethod
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
    """Store each head's K and V as two slots of packed data + block scales
    (see ``nvfp4_kv_cache_views``)."""
    if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
        return spec
    return replace(
        spec,
        num_head_slots=2 * spec.num_kv_heads,
        state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size)
        * get_dtype_size(spec.dtype),
    )

MiniMaxM3SparseTritonBackend

Bases: MiniMaxM3SparseMSABackend

Attention-backend alias selecting Triton MSA sparse decode.

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

class MiniMaxM3SparseTritonBackend(MiniMaxM3SparseMSABackend):
    """Attention-backend alias selecting Triton MSA sparse decode."""

    @staticmethod
    def get_name() -> str:
        return "TRITON_MSA"

_dequantize_query(query_fp8, q_scale, out)

Write query_fp8 * q_scale into out (E4M3 is exact in BF16).

Source code in vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

def _dequantize_query(
    query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor
) -> torch.Tensor:
    """Write ``query_fp8 * q_scale`` into ``out`` (E4M3 is exact in BF16)."""
    out.copy_(query_fp8)
    if q_scale != 1.0:
        out.mul_(q_scale)
    return out