vllm.models.minimax_m3.common.sparse_attention ¶

class MiniMaxM3SparseBackend(AttentionBackend):
    """Block-sparse GQA backend for MiniMax M3 sparse attention layers."""

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
    # bf16 or fp8 (e4m3/e5m2): the Triton kernels dequant fp8 before the dots.
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "bfloat16",
        "fp8",
        "fp8_e4m3",
        "fp8_e5m2",
    ]

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

    @staticmethod
    def get_impl_cls() -> type["MiniMaxM3SparseImpl"]:
        # Concrete impl chosen by select_main_impl_cls; base for introspection.
        return MiniMaxM3SparseImpl

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

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        return [128]

    @staticmethod
    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
        # Page size == sparse block size (one sparse block per KV page).
        return [128]

    @classmethod
    def is_sparse(cls) -> bool:
        return True

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
        # AITER sparse PA reads K and V as separate page-16 tensors. LHBNC
        # gives each side a dense plane; LBHNC interleaves the sides inside a
        # block, which the page-id rebasing in ops/sparse_pa.py absorbs. LBHNC
        # is listed first because these layers publish two head slots while the
        # indexer publishes one, and mixed HNC shapes need a block-compact
        # layout to share the block-id pool.
        if _minimax_m3_aiter_sparse_pa_requested():
            return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
        return super().supported_kv_cache_layouts()

@dataclass
class MiniMaxM3SparseDecodeMetadata:
    """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform
    per-request query length (1, or 1 + num_speculative_tokens)."""

    seq_lens: torch.Tensor  # [num_decodes] int32
    block_table: torch.Tensor
    decode_query_len: int

class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
    """Abstract base for block-sparse GQA over the indexer-selected blocks.

    Inherits ``AttentionImplBase`` for a custom forward signature (the layer
    pre-inserts K/V and runs the indexer, which writes the selected blocks into
    the shared ``layer.topk_indices_buffer``; the attend reads them back from
    there). The Triton and MSA subclasses each own a full ``forward`` -- no
    shared forward code.
    """

    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: str = "triton",
    ) -> None:
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = scale
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.kv_cache_dtype = kv_cache_dtype
        self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
        if "e5m2" in kv_cache_dtype:
            self.kv_cache_fp8_dtype = (
                torch.float8_e5m2fnuz
                if current_platform.is_fp8_fnuz()
                else torch.float8_e5m2
            )
        else:
            self.kv_cache_fp8_dtype = current_platform.fp8_dtype()
        # Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE).
        self.topk_blocks = topk_blocks
        self.block_size = sparse_block_size

    def forward(
        self,
        layer: AttentionLayer,
        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.

        The indexer has already written the top-k block ids into
        ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at
        ``[:, nd:num_tokens]``); the attend reads them from there.
        """
        raise NotImplementedError

    def should_use_msa_decode(self, layer_name: str) -> bool:
        return False

@dataclass
class MiniMaxM3SparseMetadata(AttentionMetadata):
    """Sparse-attention metadata, split into prefill and decode sub-metadata."""

    seq_lens: torch.Tensor
    max_seq_len: int
    slot_mapping: torch.Tensor

    num_actual_tokens: int  # total query tokens (decode-first batch)

    # Split counts (batch reordered decode-first).
    num_decodes: int
    num_decode_tokens: int
    num_prefills: int
    num_prefill_tokens: int

    prefill: MiniMaxM3SparsePrefillMetadata | None = None
    decode: MiniMaxM3SparseDecodeMetadata | None = None

    # ``slot_mapping`` rebased onto the page-16 numbering AITER's KV writer
    # uses when both K/V sides share a block. Only set for the AITER sparse PA
    # path; see ``MiniMaxM3SparseMetadataBuilder.build``.
    page16_slot_mapping: torch.Tensor | None = None

@dataclass
class MiniMaxM3SparsePrefillMetadata:
    """Per-prefill state; ``cu_seqlens_k``/``total_kv_blocks`` feed the MSA CSR."""

    cu_seqlens_q: torch.Tensor  # [num_prefills + 1] int32, rebased to 0
    cu_seqlens_k: torch.Tensor  # [num_prefills + 1] int32, cumulative KV lengths
    seq_lens: torch.Tensor  # [num_prefills] int32, total KV lengths
    context_lens: torch.Tensor  # [num_prefills] int32 (cached/context tokens)
    block_table: torch.Tensor
    max_query_len: int
    max_seq_len: int
    total_kv_blocks: int

    # Per-query-token request id and absolute KV position, built once per step
    # for the AITER sparse PA block-table kernel. See
    # ``minimax_m3_query_token_positions``.
    query_req_id: torch.Tensor | None = None
    query_abs_pos: torch.Tensor | None = None

class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
    """Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode."""

    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 buffer: decode [:, :nd], prefill [:, nd:].
        topk_buffer = layer.topk_indices_buffer  # type: ignore[attr-defined]
        assert topk_buffer is not None

        topk = (
            topk_buffer
            if current_platform.is_rocm()
            else topk_buffer[:num_tokens].transpose(0, 1)
        )
        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)
        kv_cache = (
            kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
        )
        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]: split-K over the selected blocks (request-major chunks).
        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, :],
                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:]: cu_seqlens_q already rebased to 0.
        if main_md.num_prefills > 0:
            p = main_md.prefill
            assert p is not None
            minimax_m3_sparse_attn(
                q[nd:],
                kv_cache,
                topk[:, nd:num_tokens, :],
                p.block_table,
                p.cu_seqlens_q,
                p.seq_lens,
                p.context_lens,
                p.max_query_len,
                self.num_kv_heads,
                self.scale,
                out[nd:],
                k_scale=k_scale,
                v_scale=v_scale,
            )
        return output