Skip to content

vllm.models.minimax_m3.amd.indexer_aiter

AITER (ROCm) indexer impl for MiniMax M3.

Scores index blocks and selects the top-k with AITER's fp8 MFMA kernels on both sides of the batch: pa_sparse_block_score_decode for the uniform-query-length decode rows and pa_sparse_block_score_prefill for the ragged prefill rows, then pa_sparse_block_topk for each. The two scoring passes share one tile body in AITER, so they agree block for block, and both encode the forced init/local blocks as the same sentinel scores the Triton indexer uses.

The top-k also emits the attend's page table. The winners are already in its workgroup's LDS, so resolving them through the block table there costs one wave and replaces the separate Triton pass in ops.sparse_pa; the rows it writes are one per (token, kv head) with the page ids folded head-minor, which is the layout pa_decode_gluon reads after it flattens the cache.

The score kernels are built on v_mfma_f32_16x16x32_fp8_fp8, so this impl requires an fp8 (e4m3) index cache and an fp8 index query -- the fused QK-norm/RoPE kernel emits both directly when the index cache is e4m3. See aiter_indexer_unsupported_reason for the full set of limits; select_aiter_indexer_impl_cls refuses to pick this impl unless they all hold, and the model falls back to the platform-neutral MiniMaxM3Indexer.

Classes:

Functions:

MiniMaxM3AiterIndexer

Bases: Module

MiniMaxM3Indexer's surface over the AITER impl.

The platform-neutral wrapper picks its impl through common's selector and forwards a Triton-only set of fused-table arguments, neither of which can reach this impl without editing common. This holds the same three members the attention layer uses -- impl, index_cache, num_index_heads -- and forwards the one argument AITER needs instead.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
class MiniMaxM3AiterIndexer(nn.Module):
    """``MiniMaxM3Indexer``'s surface over the AITER impl.

    The platform-neutral wrapper picks its impl through ``common``'s selector
    and forwards a Triton-only set of fused-table arguments, neither of which
    can reach this impl without editing ``common``. This holds the same three
    members the attention layer uses -- ``impl``, ``index_cache``,
    ``num_index_heads`` -- and forwards the one argument AITER needs instead.
    """

    def __init__(
        self,
        *,
        impl_cls: type[MiniMaxM3IndexerAiterImpl],
        sparse_bt_buffer: torch.Tensor | None = None,
        sparse_ctx_buffer: torch.Tensor | None = None,
        **impl_kwargs,
    ) -> None:
        super().__init__()
        self.impl = impl_cls(**impl_kwargs)
        # Assigned rather than passed: the impl's base ``__init__`` lives in
        # common and takes no table buffers.
        self.impl.sparse_bt_buffer = sparse_bt_buffer
        self.impl.sparse_ctx_buffer = sparse_ctx_buffer

    @property
    def index_cache(self) -> MiniMaxM3IndexerCache:
        return self.impl.index_cache

    @property
    def num_index_heads(self) -> int:
        return self.impl.num_index_heads

    def forward(
        self,
        index_query: torch.Tensor,
        *,
        decode_page16_block_table: torch.Tensor | None = None,
        prefill_page16_block_table: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        return self.impl(
            index_query,
            decode_page16_block_table=decode_page16_block_table,
            prefill_page16_block_table=prefill_page16_block_table,
        )

MiniMaxM3IndexerAiterBackend

Bases: MiniMaxM3IndexerBackend

Indexer side-cache backend selecting the AITER builder.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
class MiniMaxM3IndexerAiterBackend(MiniMaxM3IndexerBackend):
    """Indexer side-cache backend selecting the AITER builder."""

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

MiniMaxM3IndexerAiterImpl

Bases: MiniMaxM3IndexerImpl

AITER fp8 score + top-k for both prefill and decode.

Attributes:

  • pages_per_block (int) –

    Physical pages one selected block expands into for the attend.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
class MiniMaxM3IndexerAiterImpl(MiniMaxM3IndexerImpl):
    """AITER fp8 score + top-k for both prefill and decode."""

    indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerAiterBackend

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        # Both passes are v_mfma_f32_16x16x32_fp8_fp8 with no bf16 instantiation,
        # so an index cache of any other dtype would be read as e4m3 bytes.
        # The selector will not get here, but nothing else may either.
        if self.indexer_kv_dtype not in ("fp8", "fp8_e4m3"):
            raise ValueError(
                "The AITER indexer requires an fp8 e4m3 index cache, got "
                f"indexer_kv_dtype={self.indexer_kv_dtype!r}"
            )
        # Shared, stable-address page table + per-row context bound the top-k
        # emits for the attend. Owned by the model so one allocation serves
        # every layer; left None when it reserved none, in which case the
        # attend rebuilds the table itself.
        self.sparse_bt_buffer: torch.Tensor | None = None
        self.sparse_ctx_buffer: torch.Tensor | None = None

    @property
    def pages_per_block(self) -> int:
        """Physical pages one selected block expands into for the attend."""
        return self.block_size // ASM_PAGE_SIZE

    def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]:
        """The page table and context rows covering score rows ``[lo, hi)``.

        One table row per (token, kv head), head minor, which is the order
        ``pa_decode_gluon`` reads once it flattens the cache. Slices of the
        shared buffers rather than copies, so the attend sees the writes; a
        model that reserved no buffers gets throwaway ones, and the attend
        rebuilds the table itself in that case.
        """
        kv_heads = self.num_kv_heads
        if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None:
            rows = (hi - lo) * kv_heads
            device = self.index_cache.kv_cache.device
            return (
                torch.empty(
                    (rows, self.topk_blocks * self.pages_per_block),
                    dtype=torch.int32,
                    device=device,
                ),
                torch.empty(rows, dtype=torch.int32, device=device),
            )
        return (
            self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads],
            self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads],
        )

    def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor:
        """Score buffer for ``rows`` query rows.

        Left uninitialized on purpose: the score pass writes every block up to
        the longest row it covers, and the top-k reads only the blocks its own
        row can see, so nothing downstream observes the padded tail. Filling it
        would cost a write over the whole [H, rows, width] extent, which at long
        context is the largest tensor in the indexer.
        """
        return torch.empty(
            (
                self.num_index_heads,
                rows,
                score_block_width(max_seq_len, self.block_size),
            ),
            dtype=torch.float32,
            device=self.index_cache.kv_cache.device,
        )

    def forward(
        self,
        index_query: torch.Tensor,
        *,
        decode_page16_block_table: torch.Tensor | None = None,
        prefill_page16_block_table: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        from aiter.ops.msa_attention import (
            pa_sparse_block_score_decode,
            pa_sparse_block_score_prefill,
            pa_sparse_block_topk,
        )

        attn_metadata = get_forward_context().attn_metadata
        if not isinstance(attn_metadata, dict):
            return None, None  # profiling run; caches unbound
        md = attn_metadata[self.index_cache.prefix]
        assert isinstance(md, MiniMaxM3IndexerAiterMetadata)
        # The emitted page table addresses the main cache, whose blocks are a
        # different group's than the index cache's, so the top-k resolves the
        # selection through the attend's block table -- only the score pass reads
        # the index cache and takes the indexer's. It has to be the page-16
        # rebase of that table, which the attend's own metadata builder does
        # once per step; the indexer's metadata cannot reach another group's
        # blocks, so the layer reads it off the attend and hands it in.
        num_tokens = md.num_actual_tokens
        nd = md.num_decode_tokens
        iq = index_query[:num_tokens].view(
            -1, self.num_index_heads, self.index_head_dim
        )
        kv = self.index_cache.kv_cache

        # Both sides write into the single shared persistent buffer (decode at
        # [:, :nd], prefill at [:, nd:]) and return views into it. The top-k
        # takes the head/row strides, so these slices need no copy back.
        buf = self.topk_indices_buffer
        if buf is None:
            buf = torch.empty(
                (self.num_index_heads, num_tokens, self.topk_blocks),
                dtype=torch.int32,
                device=iq.device,
            )

        decode_topk: torch.Tensor | None = None
        prefill_topk: torch.Tensor | None = None

        if md.num_decodes > 0:
            d = md.decode
            assert d is not None
            assert decode_page16_block_table is not None, (
                "the AITER indexer's top-k emits the attend's page table and "
                "needs the page-16 rebase of the attend's decode block table"
            )
            score = self._new_score(nd, d.max_seq_len)
            pa_sparse_block_score_decode(
                iq[:nd],
                kv,
                score,
                d.block_table,
                d.seq_lens,
                init_blocks=self.init_blocks,
                local_blocks=self.local_blocks,
                query_len=d.decode_query_len,
                max_seq_len=d.max_seq_len,
            )
            decode_topk = buf[:, :nd, :]
            sparse_bt, sparse_ctx = self._table_rows(0, nd)
            # A decode row's causal length is seq_len - query_len + token + 1,
            # which the kernel derives itself, so the emitted table covers
            # speculative rows without any extra per-row shape.
            pa_sparse_block_topk(
                score,
                decode_topk,
                decode_page16_block_table,
                d.seq_lens,
                sparse_bt,
                sparse_ctx,
                max_seq_len=d.max_seq_len,
                block_size=self.block_size,
                query_len=d.decode_query_len,
                num_kv_heads=self.num_kv_heads,
                pages_per_block=self.pages_per_block,
            )

        if md.num_prefills > 0:
            p = md.prefill
            assert p is not None
            assert prefill_page16_block_table is not None, (
                "the AITER indexer's top-k emits the attend's page table and "
                "needs the page-16 rebase of the attend's prefill block table"
            )
            assert md.prefill_num_valid_pages is not None
            assert md.prefill_row_req_id is not None
            assert md.prefill_kv_lens is not None
            score = self._new_score(num_tokens - nd, p.max_seq_len)
            pa_sparse_block_score_prefill(
                iq[nd:],
                kv,
                score,
                p.block_table,
                p.cu_seqlens_q,
                p.seq_lens,
                init_blocks=self.init_blocks,
                local_blocks=self.local_blocks,
                max_query_len=p.max_query_len,
                max_seq_len=p.max_seq_len,
            )
            prefill_topk = buf[:, nd:num_tokens, :]
            sparse_bt, sparse_ctx = self._table_rows(nd, num_tokens)
            # Ragged rows carry their request and causal length explicitly: the
            # block count alone cannot place the tail block the table ends on.
            pa_sparse_block_topk(
                score,
                prefill_topk,
                prefill_page16_block_table,
                p.seq_lens,
                sparse_bt,
                sparse_ctx,
                max_seq_len=p.max_seq_len,
                block_size=self.block_size,
                num_valid_pages=md.prefill_num_valid_pages,
                row_req_id=md.prefill_row_req_id,
                kv_lens=md.prefill_kv_lens,
                num_kv_heads=self.num_kv_heads,
                pages_per_block=self.pages_per_block,
            )

        return decode_topk, prefill_topk

pages_per_block property

Physical pages one selected block expands into for the attend.

_new_score(rows, max_seq_len)

Score buffer for rows query rows.

Left uninitialized on purpose: the score pass writes every block up to the longest row it covers, and the top-k reads only the blocks its own row can see, so nothing downstream observes the padded tail. Filling it would cost a write over the whole [H, rows, width] extent, which at long context is the largest tensor in the indexer.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor:
    """Score buffer for ``rows`` query rows.

    Left uninitialized on purpose: the score pass writes every block up to
    the longest row it covers, and the top-k reads only the blocks its own
    row can see, so nothing downstream observes the padded tail. Filling it
    would cost a write over the whole [H, rows, width] extent, which at long
    context is the largest tensor in the indexer.
    """
    return torch.empty(
        (
            self.num_index_heads,
            rows,
            score_block_width(max_seq_len, self.block_size),
        ),
        dtype=torch.float32,
        device=self.index_cache.kv_cache.device,
    )

_table_rows(lo, hi)

The page table and context rows covering score rows [lo, hi).

One table row per (token, kv head), head minor, which is the order pa_decode_gluon reads once it flattens the cache. Slices of the shared buffers rather than copies, so the attend sees the writes; a model that reserved no buffers gets throwaway ones, and the attend rebuilds the table itself in that case.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]:
    """The page table and context rows covering score rows ``[lo, hi)``.

    One table row per (token, kv head), head minor, which is the order
    ``pa_decode_gluon`` reads once it flattens the cache. Slices of the
    shared buffers rather than copies, so the attend sees the writes; a
    model that reserved no buffers gets throwaway ones, and the attend
    rebuilds the table itself in that case.
    """
    kv_heads = self.num_kv_heads
    if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None:
        rows = (hi - lo) * kv_heads
        device = self.index_cache.kv_cache.device
        return (
            torch.empty(
                (rows, self.topk_blocks * self.pages_per_block),
                dtype=torch.int32,
                device=device,
            ),
            torch.empty(rows, dtype=torch.int32, device=device),
        )
    return (
        self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads],
        self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads],
    )

MiniMaxM3IndexerAiterMetadata dataclass

Bases: MiniMaxM3IndexerMetadata

Adds the per-row shape the ragged top-k needs.

The uniform decode rows are recovered inside the kernel from seq_lens and the shared query length, which also clamps cudagraph padding rows to nothing. Prefill rows have no such shape, so it is materialized here once per forward and shared by every layer -- which is also where the emitted page table's tail block comes from, since a block count alone does not say how many tokens the last block holds.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
@dataclass
class MiniMaxM3IndexerAiterMetadata(MiniMaxM3IndexerMetadata):
    """Adds the per-row shape the ragged top-k needs.

    The uniform decode rows are recovered inside the kernel from ``seq_lens`` and
    the shared query length, which also clamps cudagraph padding rows to nothing.
    Prefill rows have no such shape, so it is materialized here once per forward
    and shared by every layer -- which is also where the emitted page table's
    tail block comes from, since a block count alone does not say how many tokens
    the last block holds.
    """

    # [num_prefill_tokens] int32, cdiv(position + 1, sparse_block_size) per prefill row.
    prefill_num_valid_pages: torch.Tensor | None = None
    # [num_prefill_tokens] int32, the request each prefill row belongs to.
    prefill_row_req_id: torch.Tensor | None = None
    # [num_prefill_tokens] int32, causal token count (position + 1) per row.
    prefill_kv_lens: torch.Tensor | None = None

MiniMaxM3IndexerAiterMetadataBuilder

Bases: MiniMaxM3IndexerMetadataBuilder

The Triton indexer's metadata plus the prefill rows' causal shape.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
class MiniMaxM3IndexerAiterMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
    """The Triton indexer's metadata plus the prefill rows' causal shape."""

    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)
        hf_config = vllm_config.model_config.hf_config
        text_config = getattr(hf_config, "text_config", hf_config)
        self.sparse_block_size = int(
            text_config.sparse_attention_config["sparse_block_size"]
        )
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        # Companions to the base's num_valid_pages_buffer, for the two vectors
        # only the emitted table needs.
        self.row_req_id_buffer = torch.empty(
            max_tokens, dtype=torch.int32, device=device
        )
        self.kv_lens_buffer = torch.empty(max_tokens, dtype=torch.int32, device=device)

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> MiniMaxM3IndexerAiterMetadata:
        num_reqs = common_attn_metadata.num_reqs
        num_tokens = common_attn_metadata.num_actual_tokens
        query_start_loc = common_attn_metadata.query_start_loc
        seq_lens = common_attn_metadata.seq_lens
        block_table = common_attn_metadata.block_table_tensor

        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
            split_decodes_and_prefills(
                common_attn_metadata,
                decode_threshold=self.reorder_batch_threshold,
                require_uniform=True,
            )
        )
        assert num_decodes + num_prefills == num_reqs
        assert num_decode_tokens + num_prefill_tokens == num_tokens

        # Decode-first batch: context lengths into the stable cudagraph buffer.
        context_lens = self.context_len_buffer[:num_reqs]
        context_lens.copy_(
            common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
        )

        prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None
        prefill_num_valid_pages: torch.Tensor | None = None
        prefill_row_req_id: torch.Tensor | None = None
        prefill_kv_lens: torch.Tensor | None = None
        if num_prefills > 0:
            cu_seqlens_q = (query_start_loc[num_decodes:] - num_decode_tokens).to(
                torch.int32
            )
            prefill_metadata = MiniMaxM3IndexerPrefillMetadata(
                cu_seqlens_q=cu_seqlens_q,
                seq_lens=seq_lens[num_decodes:],
                context_lens=context_lens[num_decodes:],
                block_table=block_table[num_decodes:],
                max_query_len=common_attn_metadata.max_query_len,
                max_seq_len=common_attn_metadata.max_seq_len,
            )
            # A prefill row sees its own position, so its causal length and block
            # count both follow from that alone; the request it belongs to comes
            # from the query offsets. Prefill batches are never captured, so the
            # stable buffers are only being reused here, not required.
            positions = common_attn_metadata.positions
            assert positions is not None
            row_positions = positions[num_decode_tokens:num_tokens]
            prefill_num_valid_pages = self.num_valid_pages_buffer[
                num_decode_tokens:num_tokens
            ]
            prefill_num_valid_pages.copy_(
                row_positions // self.sparse_block_size + 1, non_blocking=True
            )
            prefill_kv_lens = self.kv_lens_buffer[num_decode_tokens:num_tokens]
            prefill_kv_lens.copy_(row_positions + 1, non_blocking=True)
            prefill_row_req_id = self.row_req_id_buffer[num_decode_tokens:num_tokens]
            prefill_row_req_id.copy_(
                torch.searchsorted(
                    cu_seqlens_q[1:].contiguous(),
                    torch.arange(
                        num_prefill_tokens,
                        dtype=torch.int32,
                        device=cu_seqlens_q.device,
                    ),
                    right=True,
                ),
                non_blocking=True,
            )

        decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None
        if num_decodes > 0:
            qsl_cpu = common_attn_metadata.query_start_loc_cpu
            query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
            decode_query_len = int(query_lens_cpu[0].item())
            assert decode_query_len > 0
            assert torch.all(
                (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
            )
            assert num_decode_tokens == num_decodes * decode_query_len
            decode_metadata = MiniMaxM3IndexerDecodeMetadata(
                seq_lens=seq_lens[:num_decodes],
                block_table=block_table[:num_decodes],
                max_seq_len=common_attn_metadata.max_seq_len,
                decode_query_len=decode_query_len,
                max_decode_query_len=self.max_decode_query_len,
            )

        return MiniMaxM3IndexerAiterMetadata(
            seq_lens=seq_lens,
            max_seq_len=common_attn_metadata.max_seq_len,
            slot_mapping=common_attn_metadata.slot_mapping,
            num_actual_tokens=num_tokens,
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
            num_prefill_tokens=num_prefill_tokens,
            prefill=prefill_metadata,
            decode=decode_metadata,
            prefill_num_valid_pages=prefill_num_valid_pages,
            prefill_row_req_id=prefill_row_req_id,
            prefill_kv_lens=prefill_kv_lens,
        )

aiter_indexer_max_decode_query_len(vllm_config)

Longest query a decode row can carry, which spec decode is what sets.

Mirrors _init_reorder_batch_threshold(1, supports_spec_as_decode=True), since that is what the builder splits the batch on and therefore what the decode kernel will actually be handed.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def aiter_indexer_max_decode_query_len(vllm_config: VllmConfig) -> int:
    """Longest query a decode row can carry, which spec decode is what sets.

    Mirrors ``_init_reorder_batch_threshold(1, supports_spec_as_decode=True)``,
    since that is what the builder splits the batch on and therefore what the
    decode kernel will actually be handed.
    """
    spec = vllm_config.speculative_config
    if spec is None or spec.num_speculative_tokens is None:
        return 1
    return 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens

aiter_indexer_unsupported_reason(*, topk_blocks, sparse_block_size, num_index_heads, index_head_dim, indexer_kv_dtype, max_model_len, max_decode_query_len=1, score_type='max')

Return why this config cannot use the AITER indexer, or None if it can.

Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache dtype, the compiled score/top-k contract, MFMA column limits, max context in blocks, and whether AITER exposes the MSA kernels. select_aiter_indexer_impl_cls logs the string and falls back when it is not None.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def aiter_indexer_unsupported_reason(
    *,
    topk_blocks: int,
    sparse_block_size: int,
    num_index_heads: int,
    index_head_dim: int,
    indexer_kv_dtype: IndexerKVDType,
    max_model_len: int,
    max_decode_query_len: int = 1,
    score_type: str = "max",
) -> str | None:
    """Return why this config cannot use the AITER indexer, or None if it can.

    Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache
    dtype, the compiled score/top-k contract, MFMA column limits, max context
    in blocks, and whether AITER exposes the MSA kernels.
    ``select_aiter_indexer_impl_cls`` logs the string and falls back when it is
    not None.
    """
    if not current_platform.is_rocm():
        return (
            "needs ROCm for the AITER fp8 MFMA score/top-k, "
            f"got platform={current_platform.device_type!r}"
        )
    if not _minimax_m3_aiter_sparse_pa_requested():
        # The top-k emits the attend's page table in page-16 numbering, which
        # only addresses the interleaved cache the AITER attend reads. Paired
        # with any other attend there is nowhere valid to write it, so this
        # indexer is not usable on its own.
        return (
            "needs the AITER sparse PA attend, whose page table its top-k "
            "emits (rocm_aiter_ops + shuffle KV cache layout)"
        )
    if indexer_kv_dtype not in ("fp8", "fp8_e4m3"):
        # The score kernels are fp8 MFMA; there is no bf16 instantiation.
        return (
            f"needs an fp8 e4m3 index cache, got indexer_kv_dtype={indexer_kv_dtype!r}"
        )
    from vllm.platforms.rocm import on_gfx950

    if not on_gfx950():
        return f"needs {' or '.join(SUPPORTED_ARCHS)} for the fp8 MFMA"
    if score_type != MSA_SCORE_TYPE:
        return f"needs score_type={MSA_SCORE_TYPE!r}, got score_type={score_type!r}"
    if topk_blocks != MSA_TOPK_BLOCKS:
        return f"needs topk_blocks={MSA_TOPK_BLOCKS}, got topk_blocks={topk_blocks}"
    if sparse_block_size != MSA_SPARSE_BLOCK_SIZE:
        return (
            f"needs sparse_block_size={MSA_SPARSE_BLOCK_SIZE}, "
            f"got sparse_block_size={sparse_block_size}"
        )
    if index_head_dim != MSA_INDEX_HEAD_DIM:
        return (
            f"needs index_head_dim={MSA_INDEX_HEAD_DIM}, "
            f"got index_head_dim={index_head_dim}"
        )
    if num_index_heads > MFMA_COLS:
        return f"num_index_heads={num_index_heads} exceeds the {MFMA_COLS} MFMA columns"
    # A decode row's whole query shares one MFMA tile, one column per (token,
    # head) pair, so spec decode trades columns against the head count.
    if num_index_heads * max_decode_query_len > MFMA_COLS:
        return (
            f"num_index_heads={num_index_heads} x max_decode_query_len="
            f"{max_decode_query_len} exceeds the {MFMA_COLS} MFMA columns"
        )
    max_blocks = math.ceil(max_model_len / sparse_block_size)
    if max_blocks > MAX_SUPPORTED_BLOCKS:
        return (
            f"max_model_len={max_model_len} needs {max_blocks} blocks per row, "
            f"more than the top-k's {MAX_SUPPORTED_BLOCKS}"
        )
    return aiter_msa_kernels_unavailable_reason()

aiter_msa_kernels_unavailable_reason() cached

Return why the AITER MSA score/top-k ops cannot be imported, or None.

They are a recent addition, so an AITER that predates them imports fine while these three names do not exist, and the failure would otherwise surface as an ImportError from the middle of a forward. compile_ops binds them lazily, so this costs the module import only -- the kernel build itself still happens on the first call.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
@cache
def aiter_msa_kernels_unavailable_reason() -> str | None:
    """Return why the AITER MSA score/top-k ops cannot be imported, or None.

    They are a recent addition, so an AITER that predates them imports fine
    while these three names do not exist, and the failure would otherwise
    surface as an ImportError from the middle of a forward. ``compile_ops``
    binds them lazily, so this costs the module import only -- the kernel build
    itself still happens on the first call.
    """
    try:
        from aiter.ops.msa_attention import (  # noqa: F401
            pa_sparse_block_score_decode,
            pa_sparse_block_score_prefill,
            pa_sparse_block_topk,
        )
    except ImportError as exc:
        return f"AITER cannot supply the MSA score/top-k kernels ({exc})"
    return None

score_block_width(max_seq_len, block_size)

Block-axis width the top-k requires of the score buffer.

Lanes read whole wave-wide strips with no tail guard and each holds a power-of-two count of them, so the axis is padded past the block count.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def score_block_width(max_seq_len: int, block_size: int) -> int:
    """Block-axis width the top-k requires of the score buffer.

    Lanes read whole wave-wide strips with no tail guard and each holds a
    power-of-two count of them, so the axis is padded past the block count.
    """
    max_blk = math.ceil(max(max_seq_len, 1) / block_size)
    return _pow2_ceil(math.ceil(max_blk / WAVE_SIZE)) * WAVE_SIZE

select_aiter_indexer_impl_cls(*, topk_blocks, sparse_block_size, num_index_heads, index_head_dim, indexer_kv_dtype, score_type='max')

The AITER indexer impl if this config can use it, else None.

None sends the caller to the platform-neutral MiniMaxM3Indexer, which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8 index cache that lands here has nowhere to go.

Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
def select_aiter_indexer_impl_cls(
    *,
    topk_blocks: int,
    sparse_block_size: int,
    num_index_heads: int,
    index_head_dim: int,
    indexer_kv_dtype: IndexerKVDType,
    score_type: str = "max",
) -> type[MiniMaxM3IndexerAiterImpl] | None:
    """The AITER indexer impl if this config can use it, else None.

    ``None`` sends the caller to the platform-neutral ``MiniMaxM3Indexer``,
    which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8
    index cache that lands here has nowhere to go.
    """
    reason = aiter_indexer_unsupported_reason(
        topk_blocks=topk_blocks,
        sparse_block_size=sparse_block_size,
        num_index_heads=num_index_heads,
        index_head_dim=index_head_dim,
        indexer_kv_dtype=indexer_kv_dtype,
        max_model_len=get_current_vllm_config().model_config.max_model_len,
        max_decode_query_len=aiter_indexer_max_decode_query_len(
            get_current_vllm_config()
        ),
        score_type=score_type,
    )
    if reason is not None:
        logger.info_once("MiniMax M3 indexer: AITER unavailable (%s)", reason)
        return None
    logger.info_once(
        "MiniMax M3 indexer: selected AITER (fp8 MFMA score + top-k) "
        "[topk_blocks=%d, indexer_kv_dtype=%s]",
        topk_blocks,
        indexer_kv_dtype,
    )
    return MiniMaxM3IndexerAiterImpl