[Performance][ROCm] Integrate aiter indexer scoring and top-k kernels into MiniMax-M3 sparse attention path (#52664)

Signed-off-by: Yamini Preethi Kamisetty <YaminiPreethi.Kamisetty@amd.com>

tests/kernels/attention/test_minimax_m3.py

@@ -1694,6 +1694,58 @@ def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
     with pytest.raises(ValueError, match="num_kv_heads == 1"):
         sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(2)
 
+    # An indexer that emits the attend's page table addresses every head, so
+    # the pairing the check guards against does not arise.
+    assert sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(
+        2, emits_sparse_block_table=True
+    )
+    # The AITER path carries its own backend: its builder is what rebases the
+    # block table the indexer's top-k resolves through, so selecting the impl
+    # without it would leave that field unset.
+    backend, impl_cls = sparse_attn_mod.select_main_backend_and_impl_cls(
+        topk_blocks=TOPK,
+        kv_cache_dtype="fp8_e4m3",
+        num_kv_heads=2,
+        emits_sparse_block_table=True,
+    )
+    assert impl_cls.__name__ == "MiniMaxM3SparseAiterPAImpl"
+    assert backend.__name__ == "MiniMaxM3SparseAiterPABackend"
+    assert backend.get_builder_cls().__name__ == "MiniMaxM3SparseAiterPAMetadataBuilder"
+
+
+def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):
+    """The top-k emits the attend's page table, so it needs that attend.
+
+    Selected without it, the indexer would reach its top-k with no table to
+    resolve through; falling back instead surfaces the fp8 index cache as the
+    configuration error it is.
+    """
+    import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod
+
+    monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True)
+    kwargs = dict(
+        topk_blocks=TOPK,
+        sparse_block_size=BLOCK_SIZE,
+        num_index_heads=1,
+        index_head_dim=HEAD_DIM,
+        indexer_kv_dtype="fp8_e4m3",
+        max_model_len=8192,
+    )
+
+    monkeypatch.setattr(
+        indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False
+    )
+    reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)
+    assert reason is not None and "AITER sparse PA attend" in reason
+
+    # With the attend asked for, the gate moves on to the kernel-contract
+    # checks rather than stopping here.
+    monkeypatch.setattr(
+        indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True
+    )
+    reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)
+    assert reason is None or "AITER sparse PA attend" not in reason
+
 
 def test_indexer_cache_squeezes_to_contiguous_3d():
     """The indexer side cache is standardized 4D with H=1: under both layouts

vllm/models/minimax_m3/amd/indexer_aiter.py

@@ -0,0 +1,646 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""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``.
+"""
+
+import math
+from dataclasses import dataclass
+from functools import cache
+from typing import ClassVar
+
+import torch
+from torch import nn
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.attention import IndexerKVDType
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.models.minimax_m3.amd.ops.sparse_pa import ASM_PAGE_SIZE
+from vllm.models.minimax_m3.common.indexer import (
+    MiniMaxM3IndexerBackend,
+    MiniMaxM3IndexerCache,
+    MiniMaxM3IndexerDecodeMetadata,
+    MiniMaxM3IndexerImpl,
+    MiniMaxM3IndexerMetadata,
+    MiniMaxM3IndexerMetadataBuilder,
+    MiniMaxM3IndexerPrefillMetadata,
+)
+
+# The single source of truth for whether the AITER attend was asked for; this
+# indexer is only usable alongside it.
+from vllm.models.minimax_m3.common.sparse_attention import (
+    _minimax_m3_aiter_sparse_pa_requested,
+)
+from vllm.platforms import current_platform
+from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+logger = init_logger(__name__)
+
+# Compiled AITER MiniMax-M3 score/top-k contract
+MSA_TOPK_BLOCKS = 16
+MSA_SCORE_TYPE = "max"
+MSA_SPARSE_BLOCK_SIZE = 128
+MSA_INDEX_HEAD_DIM = 128
+# Wave width the top-k is written against: it gives one lane per output slot and
+# reads the score row in wave-wide strips.
+WAVE_SIZE = 64
+# Per-lane register slots the top-k can hold, which is what caps the context: a
+# row may span at most SLOTS_MAX * WAVE_SIZE blocks.
+SLOTS_MAX = 128
+MAX_SUPPORTED_BLOCKS = SLOTS_MAX * WAVE_SIZE
+# MFMA columns the score pass has, one per (query token, index head) pair.
+MFMA_COLS = 16
+# The fp8 MFMA the score kernels are built on exists on these targets only.
+SUPPORTED_ARCHS = ("gfx950",)
+
+
+def _pow2_ceil(n: int) -> int:
+    return 1 << (n - 1).bit_length() if n >= 1 else 1
+
+
+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
+
+
+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
+
+
+@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
+
+
+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()
+
+
+class MiniMaxM3IndexerAiterBackend(MiniMaxM3IndexerBackend):
+    """Indexer side-cache backend selecting the AITER builder."""
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3IndexerAiterMetadataBuilder"]:
+        return MiniMaxM3IndexerAiterMetadataBuilder
+
+
+@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
+
+
+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,
+        )
+
+
+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
+
+
+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
+
+
+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,
+        )

vllm/models/minimax_m3/amd/model.py

@@ -87,17 +87,24 @@ from vllm.model_executor.models.utils import (
     maybe_prefix,
 )
 from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
+from vllm.models.minimax_m3.amd.indexer_aiter import (
+    MiniMaxM3AiterIndexer,
+    select_aiter_indexer_impl_cls,
+)
 from vllm.models.minimax_m3.amd.ops import (
     gemma_fused_add_rmsnorm,
     gemma_rmsnorm,
     swiglu_oai_split,
 )
 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
+    ASM_PAGE_SIZE,
     minimax_m3_alloc_sparse_block_table,
     minimax_m3_sparse_block_page_stride,
 )
 from vllm.models.minimax_m3.amd.sparse_attention_msa import (
+    MiniMaxM3SparseAiterPADecodeMetadata,
     MiniMaxM3SparseAiterPAImpl,
+    MiniMaxM3SparseAiterPAPrefillMetadata,
 )
 from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
 from vllm.models.minimax_m3.common.mm_preprocess import (
@@ -111,7 +118,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
     MiniMaxM3SparseMetadata,
     minimax_m3_rebase_slots_to_page16,
     minimax_m3_use_aiter_sparse_pa,
-    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
@@ -596,6 +603,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         prefix: str = "",
         cache_config: CacheConfig | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -684,18 +692,45 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
         set_default_quant_scales(self, register_buffer=True)
 
+        # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main cache
+        # (--attention-config '{"indexer_kv_dtype": ...}'). fp8 e4m3 is what the
+        # AITER indexer needs; bf16 keeps the Triton indexer.
+        self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
+            "bf16"
+        )
+
         # Shared top-k buffer: the indexer writes the selected blocks into it and
         # the attend impl reads them back (no Python value crosses the break).
         self.topk_indices_buffer = topk_indices_buffer
-        self.attn_backend = MiniMaxM3SparseBackend
-        # Indexer and main attention are separate impls. On ROCm the SM100 gate
-        # is always False, so both pick Triton and the index cache stays bf16.
-        # impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
-        self.impl: MiniMaxM3SparseImpl = select_main_impl_cls(  # type: ignore[assignment]
+        # Indexer and main attention are separate impls. The AITER indexer is
+        # ROCm-only, so it is selected here rather than inside the neutral
+        # MiniMaxM3Indexer, which knows nothing about it and would pick Triton.
+        indexer_impl_cls = select_aiter_indexer_impl_cls(
+            topk_blocks=sparse_cfg["sparse_topk_blocks"],
+            sparse_block_size=sparse_cfg["sparse_block_size"],
+            num_index_heads=self.num_idx_heads,
+            index_head_dim=self.idx_head_dim,
+            indexer_kv_dtype=self.indexer_kv_dtype,
+            score_type=sparse_cfg.get("sparse_score_type", "max"),
+        )
+        # The attend gate below needs this: an emitted table is what lets it
+        # serve more than one KV head per rank. Buffers to write the table into
+        # only arrive when the model already resolved the attend to the AITER
+        # path the table addresses, so that is not rechecked here.
+        self.indexer_emits_table = (
+            indexer_impl_cls is not None and sparse_table_buffers is not None
+        )
+        # The backend names the metadata builder, which for the AITER path is
+        # the one that rebases the block table the indexer's top-k resolves
+        # through, so both have to come from the same selection.
+        self.attn_backend, main_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,
-        )(
+            emits_sparse_block_table=self.indexer_emits_table,
+        )
+        # impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
+        self.impl: MiniMaxM3SparseImpl = main_impl_cls(  # type: ignore[assignment]
             self.num_heads,
             self.head_dim,
             self.scaling,
@@ -704,14 +739,14 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             topk_blocks=sparse_cfg["sparse_topk_blocks"],
             sparse_block_size=sparse_cfg["sparse_block_size"],
         )
-        self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(self.num_kv_heads)
+        self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
+            self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table
+        )
         self.kv_cache_k = torch.tensor([])
         self.kv_cache_v = torch.tensor([])
         self._aiter_sparse_pa_cache_data_ptr = 0
         self._aiter_sparse_pa_block_page_stride = 0
-        # Self-contained nn.Module: owns its side cache, selects its impl in init
-        # (Triton on ROCm, where the SM100 gate is always False).
-        self.indexer = MiniMaxM3Indexer(
+        indexer_kwargs = dict(
             num_kv_heads=self.num_kv_heads,
             scale=self.scaling,
             topk_blocks=sparse_cfg["sparse_topk_blocks"],
@@ -723,8 +758,24 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             local_blocks=sparse_cfg.get("sparse_local_block", 0),
             score_type=sparse_cfg.get("sparse_score_type", "max"),
             cache_config=cache_config,
+            indexer_kv_dtype=self.indexer_kv_dtype,
             topk_indices_buffer=topk_indices_buffer,
         )
+        self.sparse_bt_buffer: torch.Tensor | None = None
+        self.sparse_ctx_buffer: torch.Tensor | None = None
+        if indexer_impl_cls is None:
+            # Self-contained nn.Module: owns its side cache, selects its impl.
+            self.indexer = MiniMaxM3Indexer(**indexer_kwargs)
+        else:
+            if self.indexer_emits_table:
+                assert sparse_table_buffers is not None
+                self.sparse_bt_buffer, self.sparse_ctx_buffer = sparse_table_buffers
+            self.indexer = MiniMaxM3AiterIndexer(
+                impl_cls=indexer_impl_cls,
+                sparse_bt_buffer=self.sparse_bt_buffer,
+                sparse_ctx_buffer=self.sparse_ctx_buffer,
+                **indexer_kwargs,
+            )
 
         # Register the main K/V cache so the KV-cache manager allocates it.
         compilation_config = vllm_config.compilation_config
@@ -952,7 +1003,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
                 )
         else:
             index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
-            index_q = qkv.new_empty((num_tokens, self.index_q_size))
+            # index_q matches the index-K cache dtype (e4m3 for the AITER fp8
+            # score path); the fused kernel emits fp8 directly when this buffer
+            # is e4m3.
+            index_q = qkv.new_empty(
+                (num_tokens, self.index_q_size),
+                dtype=self.indexer.index_cache.dtype,
+            )
             if self.use_aiter_sparse_pa:
                 ops.fused_minimax_m3_qknorm_rope_kv_insert(
                     qkv,
@@ -1034,31 +1091,56 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None
         if not self.skip_index_topk:
             assert index_query is not None
-            attn_metadata = get_forward_context().attn_metadata
-            if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
-                main_md = attn_metadata[self.layer_name]
-                assert isinstance(main_md, MiniMaxM3SparseMetadata)
-                if main_md.num_decodes > 0:
-                    d = main_md.decode
-                    assert d is not None
-                    topk = self.topk_indices_buffer
-                    assert topk is not None
-                    decode_sparse_table = minimax_m3_alloc_sparse_block_table(
-                        topk[:, : main_md.num_decode_tokens, :]
-                    )
-                    block_page_stride = self._aiter_sparse_pa_block_page_stride
-                    assert block_page_stride > 0
-                    self.indexer(
-                        index_query,
-                        attention_block_table=d.block_table,
-                        sparse_block_table_out=decode_sparse_table[0],
-                        sparse_context_lens_out=decode_sparse_table[1],
-                        block_page_stride=block_page_stride,
-                    )
+            # The AITER indexer emits the attend table into its persistent
+            # buffers, resolving its selection through the page-16 rebase of the
+            # attend's block table. The Triton indexer can instead fuse decode
+            # top-k with main's per-forward sparse-table allocation.
+            if self.indexer_emits_table:
+                attn_metadata = get_forward_context().attn_metadata
+                decode_page16 = prefill_page16 = None
+                if isinstance(attn_metadata, dict):
+                    main_md = attn_metadata[self.layer_name]
+                    assert isinstance(main_md, MiniMaxM3SparseMetadata)
+                    # Rebased once per step by the AITER attend's builder, which
+                    # this path selected along with the impl.
+                    d, p = main_md.decode, main_md.prefill
+                    if d is not None:
+                        assert isinstance(d, MiniMaxM3SparseAiterPADecodeMetadata)
+                        decode_page16 = d.page16_block_table
+                    if p is not None:
+                        assert isinstance(p, MiniMaxM3SparseAiterPAPrefillMetadata)
+                        prefill_page16 = p.page16_block_table
+                self.indexer(
+                    index_query,
+                    decode_page16_block_table=decode_page16,
+                    prefill_page16_block_table=prefill_page16,
+                )
+            else:
+                attn_metadata = get_forward_context().attn_metadata
+                if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
+                    main_md = attn_metadata[self.layer_name]
+                    assert isinstance(main_md, MiniMaxM3SparseMetadata)
+                    if main_md.num_decodes > 0:
+                        d = main_md.decode
+                        assert d is not None
+                        topk = self.topk_indices_buffer
+                        assert topk is not None
+                        decode_sparse_table = minimax_m3_alloc_sparse_block_table(
+                            topk[:, : main_md.num_decode_tokens, :]
+                        )
+                        block_page_stride = self._aiter_sparse_pa_block_page_stride
+                        assert block_page_stride > 0
+                        self.indexer(
+                            index_query,
+                            attention_block_table=d.block_table,
+                            sparse_block_table_out=decode_sparse_table[0],
+                            sparse_context_lens_out=decode_sparse_table[1],
+                            block_page_stride=block_page_stride,
+                        )
+                    else:
+                        self.indexer(index_query)
                 else:
                     self.indexer(index_query)
-            else:
-                self.indexer(index_query)
         if self.use_aiter_sparse_pa:
             assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
             return self.impl.forward(
@@ -1081,6 +1163,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
         force_sparse_attn: bool = False,
         force_moe: bool = False,
         topk_indices_buffer: torch.Tensor | None = None,
+        sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -1101,6 +1184,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
                 prefix=f"{prefix}.self_attn",
                 cache_config=cache_config,
                 topk_indices_buffer=topk_indices_buffer,
+                sparse_table_buffers=sparse_table_buffers,
             )
         else:
             self.self_attn = MiniMaxM3Attention(
@@ -1189,15 +1273,52 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
         # layers (mirrors DeepseekV4); the indexer writes its per-head decode/
         # prefill block selection into it, the attend reads it back.
         sparse_cfg = getattr(config, "sparse_attention_config", None)
+        self.sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None
         if sparse_cfg is not None:
             tp_size = get_tensor_model_parallel_world_size()
             num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
+            max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
             self.topk_indices_buffer = torch.empty(
                 num_index_heads,
-                vllm_config.scheduler_config.max_num_batched_tokens,
+                max_tokens,
                 sparse_cfg["sparse_topk_blocks"],
                 dtype=torch.int32,
             )
+
+            num_kv_heads = max(1, config.num_key_value_heads // tp_size)
+            indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
+                "bf16"
+            )
+            # Probed with the head count the layer will pass (it derives
+            # num_idx_heads from num_kv_heads, not from sparse_num_index_heads),
+            # so this cannot land on the other side of the MFMA column limit
+            # from the selection the layer makes.
+            emits_table = (
+                select_aiter_indexer_impl_cls(
+                    topk_blocks=sparse_cfg["sparse_topk_blocks"],
+                    sparse_block_size=sparse_cfg["sparse_block_size"],
+                    num_index_heads=num_kv_heads,
+                    index_head_dim=sparse_cfg["sparse_index_dim"],
+                    indexer_kv_dtype=indexer_kv_dtype,
+                    score_type=sparse_cfg.get("sparse_score_type", "max"),
+                )
+                is not None
+            )
+            if emits_table and minimax_m3_use_aiter_sparse_pa(
+                num_kv_heads, emits_sparse_block_table=True
+            ):
+                # One row per (token, kv head): the table addresses the attend's
+                # cache, not the index cache.
+                rows = max_tokens * num_kv_heads
+                self.sparse_table_buffers = (
+                    torch.empty(
+                        rows,
+                        sparse_cfg["sparse_topk_blocks"]
+                        * (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE),
+                        dtype=torch.int32,
+                    ),
+                    torch.empty(rows, dtype=torch.int32),
+                )
         else:
             self.topk_indices_buffer = None
 
@@ -1209,6 +1330,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                 cache_config=cache_config,
                 quant_config=quant_config,
                 topk_indices_buffer=self.topk_indices_buffer,
+                sparse_table_buffers=self.sparse_table_buffers,
             ),
             prefix=f"{prefix}.layers",
         )

vllm/models/minimax_m3/amd/ops/sparse_pa.py

@@ -312,7 +312,9 @@ def _insert_index_cache_kernel(
     )
     mask = (slot >= 0) & (offs_d < HEAD_DIM)
     value = tl.load(src, mask=offs_d < HEAD_DIM, other=0.0)
-    tl.store(dst, value, mask=mask)
+    # The side cache carries its own dtype (bf16, or e4m3 for the fp8 indexer),
+    # so convert rather than leaning on the store's implicit cast.
+    tl.store(dst, value.to(dst.dtype.element_ty), mask=mask)
 
 
 @torch.no_grad()
@@ -368,6 +370,32 @@ def minimax_m3_sparse_block_page_stride(
     return PAGES_PER_SPARSE_BLOCK * (2 if _sides_are_packed(k_cache, v_cache) else 1)
 
 
+# K/V sides sharing one block, which is how many sides' pages a block spans.
+# The attend publishes two head slots and the indexer's side cache one, and
+# specs that mix HNC shapes narrow the resolved layout to a block-compact one
+# (LHBNC is not), so this path always lands on an interleaved block.
+PAGE16_SIDES_PER_BLOCK = 2
+
+
+def minimax_m3_rebase_block_table_to_page16(
+    block_table: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Rebase a logical page table onto AITER's page-16 page numbering.
+
+    The page ids AITER's top-k emits for the attend are
+    ``block_table[blk] * pages_per_block + j``, and ``pages_per_block`` is fixed
+    at AITER build time to one side's pages. An interleaved block spans both
+    sides, so the only way to reach its pages through that kernel is to hand it
+    a table already scaled to the wider stride. The caller does this once per
+    step, since every sparse layer resolves its selection through the same
+    table.
+    """
+    if out is None:
+        out = torch.empty_like(block_table)
+    return torch.mul(block_table, PAGE16_SIDES_PER_BLOCK, out=out)
+
+
 def _gluon_scale_arg(
     scale: torch.Tensor | None,
     *,
@@ -461,14 +489,23 @@ def minimax_m3_sparse_attn_prefill_aiter(
     output: torch.Tensor,
     k_scale: torch.Tensor | None = None,
     v_scale: torch.Tensor | None = None,
+    sparse_bt: torch.Tensor | None = None,
+    sparse_ctx: torch.Tensor | None = None,
 ) -> None:
-    sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
-        topk_idx,
-        block_table,
-        query_req_id,
-        query_abs_pos,
-        minimax_m3_sparse_block_page_stride(k_cache, v_cache),
-    )
+    if (sparse_bt is None) != (sparse_ctx is None):
+        raise ValueError(
+            "MiniMax-M3 prepared sparse block table and context lengths "
+            "must be provided together"
+        )
+    if sparse_bt is None:
+        sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
+            topk_idx,
+            block_table,
+            query_req_id,
+            query_abs_pos,
+            minimax_m3_sparse_block_page_stride(k_cache, v_cache),
+        )
+    assert sparse_ctx is not None
     _run_gluon_decode(
         q,
         k_cache,

vllm/models/minimax_m3/amd/sparse_attention_msa.py

@@ -2,16 +2,118 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """MSA AITER block-sparse attend for MiniMax M3."""
 
+from dataclasses import dataclass
+
 import torch
 
+from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context
 from vllm.models.minimax_m3.common.sparse_attention import (
+    MiniMaxM3SparseBackend,
+    MiniMaxM3SparseDecodeMetadata,
     MiniMaxM3SparseImpl,
     MiniMaxM3SparseMetadata,
+    MiniMaxM3SparseMetadataBuilder,
+    MiniMaxM3SparsePrefillMetadata,
 )
+from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import (
     AttentionLayer,
 )
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+
+class MiniMaxM3SparseAiterPABackend(MiniMaxM3SparseBackend):
+    """MiniMax M3 backend carrying the AITER page-16 block-table rebase."""
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3SparseAiterPAMetadataBuilder"]:
+        return MiniMaxM3SparseAiterPAMetadataBuilder
+
+
+@dataclass
+class MiniMaxM3SparseAiterPAPrefillMetadata(MiniMaxM3SparsePrefillMetadata):
+    # ``block_table`` rebased onto AITER's page-16 numbering; see
+    # ``MiniMaxM3SparseAiterPAMetadataBuilder.build``.
+    page16_block_table: torch.Tensor | None = None
+
+
+@dataclass
+class MiniMaxM3SparseAiterPADecodeMetadata(MiniMaxM3SparseDecodeMetadata):
+    page16_block_table: torch.Tensor | None = None
+
+
+class MiniMaxM3SparseAiterPAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
+    """Adds the page-16 rebase of the block table the indexer's top-k needs.
+
+    The AITER indexer's top-k emits the attend's page table, expanding a
+    selected block into a compile-time number of pages -- one side's worth,
+    while an interleaved block holds both -- so it has to resolve the selection
+    through a table pre-scaled to the wider stride. Every sparse layer selects
+    through the same table, so it is rebased here once per step rather than
+    once per layer, alongside the slot mapping the base rebases for the writer.
+    """
+
+    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)
+        # Stable address for the same reason as the base's own buffers: a
+        # captured decode replay writes through the address it captured.
+        self.page16_block_table_buffer = torch.empty(
+            (
+                vllm_config.scheduler_config.max_num_seqs,
+                cdiv(vllm_config.model_config.max_model_len, kv_cache_spec.block_size),
+            ),
+            dtype=torch.int32,
+            device=device,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> MiniMaxM3SparseMetadata:
+        from vllm.models.minimax_m3.amd.ops.sparse_pa import (
+            minimax_m3_rebase_block_table_to_page16,
+        )
+
+        metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
+        block_table = common_attn_metadata.block_table_tensor
+        num_reqs, cols = block_table.shape
+        assert (
+            num_reqs <= self.page16_block_table_buffer.shape[0]
+            and cols <= self.page16_block_table_buffer.shape[1]
+        ), (
+            f"block table {(num_reqs, cols)} exceeds the page-16 rebase buffer "
+            f"{tuple(self.page16_block_table_buffer.shape)}"
+        )
+        page16 = minimax_m3_rebase_block_table_to_page16(
+            block_table,
+            out=self.page16_block_table_buffer[:num_reqs, :cols],
+        )
+
+        # Decode-first batch, so each side takes the rows its own block table
+        # was sliced from. Reconstructed from the base's fields rather than
+        # listed out, so a field added there does not silently drop here.
+        nd = metadata.num_decodes
+        if metadata.decode is not None:
+            metadata.decode = MiniMaxM3SparseAiterPADecodeMetadata(
+                **vars(metadata.decode),
+                page16_block_table=page16[:nd],
+            )
+        if metadata.prefill is not None:
+            metadata.prefill = MiniMaxM3SparseAiterPAPrefillMetadata(
+                **vars(metadata.prefill),
+                page16_block_table=page16[nd:],
+            )
+        return metadata
 
 
 class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
@@ -28,8 +130,11 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
         decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         from vllm.models.minimax_m3.amd.ops.sparse_pa import (
+            PAGE16_SIDES_PER_BLOCK,
+            PAGES_PER_SPARSE_BLOCK,
             minimax_m3_sparse_attn_decode_aiter,
             minimax_m3_sparse_attn_prefill_aiter,
+            minimax_m3_sparse_block_page_stride,
         )
 
         attn_metadata = get_forward_context().attn_metadata
@@ -42,10 +147,16 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
         num_tokens = main_md.num_actual_tokens
         topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
         assert topk is not None
-        if self.num_kv_heads != 1:
+        # Set only when the layer's indexer emitted the page table alongside its
+        # selection; otherwise these stay None and the table is built below.
+        sparse_bt_buf = getattr(layer, "sparse_bt_buffer", None)
+        sparse_ctx_buf = getattr(layer, "sparse_ctx_buffer", None)
+        kvh = self.num_kv_heads
+        if kvh != 1 and sparse_bt_buf is None:
             raise NotImplementedError(
-                "MiniMax-M3 AITER sparse PA currently requires per-rank "
-                f"num_kv_heads == 1, got {self.num_kv_heads}"
+                "MiniMax-M3 AITER sparse PA needs the page table the indexer's "
+                f"top-k emits to serve per-rank num_kv_heads == {kvh}; the "
+                "Triton builders it falls back to address one head's cache."
             )
 
         hd = self.head_size
@@ -55,9 +166,31 @@ class MiniMaxM3SparseAiterPAImpl(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
 
+        if sparse_bt_buf is not None:
+            # An emitted table was scaled by the metadata builder, which reads
+            # the packing off the layout rather than off these tensors. The two
+            # only disagree if the resolved layout gives each side its own
+            # plane, and a silent disagreement reads unrelated pages.
+            stride = minimax_m3_sparse_block_page_stride(k_cache, v_cache)
+            expected = PAGES_PER_SPARSE_BLOCK * PAGE16_SIDES_PER_BLOCK
+            if stride != expected:
+                raise RuntimeError(
+                    "MiniMax-M3 AITER sparse PA: the indexer emitted a page "
+                    f"table for a block spanning {expected} pages, but this "
+                    f"KV cache lays a block out over {stride}. The page-16 "
+                    "rebase assumes both K/V sides share a block."
+                )
+
         if main_md.num_decodes > 0:
             d = main_md.decode
             assert d is not None
+            prepared_decode_table = decode_sparse_table
+            if sparse_bt_buf is not None:
+                assert sparse_ctx_buf is not None
+                prepared_decode_table = (
+                    sparse_bt_buf[: nd * kvh],
+                    sparse_ctx_buf[: nd * kvh],
+                )
             minimax_m3_sparse_attn_decode_aiter(
                 q[:nd],
                 k_cache,
@@ -72,10 +205,14 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
                 v_scale=v_scale,
                 decode_query_len=d.decode_query_len,
                 sparse_block_table=(
-                    decode_sparse_table[0] if decode_sparse_table is not None else None
+                    prepared_decode_table[0]
+                    if prepared_decode_table is not None
+                    else None
                 ),
                 sparse_context_lens=(
-                    decode_sparse_table[1] if decode_sparse_table is not None else None
+                    prepared_decode_table[1]
+                    if prepared_decode_table is not None
+                    else None
                 ),
             )
 
@@ -96,5 +233,15 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
                 out[nd:],
                 k_scale=k_scale,
                 v_scale=v_scale,
+                sparse_bt=(
+                    None
+                    if sparse_bt_buf is None
+                    else sparse_bt_buf[nd * kvh : num_tokens * kvh]
+                ),
+                sparse_ctx=(
+                    None
+                    if sparse_ctx_buf is None
+                    else sparse_ctx_buf[nd * kvh : num_tokens * kvh]
+                ),
             )
         return output

vllm/models/minimax_m3/common/sparse_attention.py

@@ -109,13 +109,22 @@ def minimax_m3_query_token_positions(
     return query_req_id, query_abs_pos
 
 
-def minimax_m3_use_aiter_sparse_pa(num_kv_heads: int) -> bool:
-    """Whether to use the ROCm AITER page-16 sparse PA prototype."""
+def minimax_m3_use_aiter_sparse_pa(
+    num_kv_heads: int, *, emits_sparse_block_table: bool = False
+) -> bool:
+    """Whether to use the ROCm AITER page-16 sparse PA prototype.
+
+    More than one KV head per rank is only served when the layer's indexer
+    emits the attend's page table itself, since the Triton builders this path
+    otherwise falls back to address one head's cache. Whether an indexer does
+    that is a ROCm-side fact, so its caller passes it in.
+    """
     requested = _minimax_m3_aiter_sparse_pa_requested()
-    if requested and num_kv_heads != 1:
+    if requested and num_kv_heads != 1 and not emits_sparse_block_table:
         raise ValueError(
             "MiniMax M3 AITER sparse paged attention requires "
-            f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}."
+            f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}, "
+            "unless the indexer emits the attend's page table."
         )
     return requested
 
@@ -252,9 +261,9 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
         # Every sparse layer shares one slot mapping, so the AITER page-16
         # rebase is done here once per step instead of once per layer. Stable
         # buffer for the same reason as the context lengths above.
-        self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
-            kv_cache_spec.num_kv_heads
-        )
+        # The request itself, not the gate: the buffer is the same whatever the
+        # head count, and a builder cannot see which indexer its layers picked.
+        self.use_aiter_sparse_pa = _minimax_m3_aiter_sparse_pa_requested()
         self.page16_slot_mapping_buffer: torch.Tensor | None = None
         if self.use_aiter_sparse_pa:
             self.page16_slot_mapping_buffer = torch.empty(
@@ -514,6 +523,7 @@ def select_main_backend_and_impl_cls(
     topk_blocks: int,
     kv_cache_dtype: str,
     num_kv_heads: int,
+    emits_sparse_block_table: bool = False,
 ) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]:
     """Pick the main attention backend and implementation.
 
@@ -523,7 +533,9 @@ def select_main_backend_and_impl_cls(
     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)
+    use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
+        num_kv_heads, emits_sparse_block_table=emits_sparse_block_table
+    )
     use_msa = (
         current_platform.is_cuda()
         and current_platform.is_device_capability_family(100)
@@ -541,10 +553,11 @@ def select_main_backend_and_impl_cls(
     )
     if use_aiter_sparse_pa:
         from vllm.models.minimax_m3.amd.sparse_attention_msa import (
+            MiniMaxM3SparseAiterPABackend,
             MiniMaxM3SparseAiterPAImpl,
         )
 
-        return MiniMaxM3SparseBackend, MiniMaxM3SparseAiterPAImpl
+        return MiniMaxM3SparseAiterPABackend, MiniMaxM3SparseAiterPAImpl
     if use_msa:
         from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
             MiniMaxM3SparseMSABackend,
@@ -560,10 +573,12 @@ def select_main_impl_cls(
     topk_blocks: int,
     kv_cache_dtype: str,
     num_kv_heads: int,
+    emits_sparse_block_table: bool = False,
 ) -> 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,
+        emits_sparse_block_table=emits_sparse_block_table,
     )[1]