From 83252ea899c6538eaa0c1fb31f28a92c661bbffc Mon Sep 17 00:00:00 2001 From: ykamiset Date: Wed, 9 Sep 2026 17:40:33 -0500 Subject: [PATCH] [Performance][ROCm] Integrate aiter indexer scoring and top-k kernels into MiniMax-M3 sparse attention path (#52664) Signed-off-by: Yamini Preethi Kamisetty --- tests/kernels/attention/test_minimax_m3.py | 52 ++ vllm/models/minimax_m3/amd/indexer_aiter.py | 646 ++++++++++++++++++ vllm/models/minimax_m3/amd/model.py | 194 +++++- vllm/models/minimax_m3/amd/ops/sparse_pa.py | 53 +- .../minimax_m3/amd/sparse_attention_msa.py | 157 ++++- .../minimax_m3/common/sparse_attention.py | 33 +- 6 files changed, 1077 insertions(+), 58 deletions(-) create mode 100644 vllm/models/minimax_m3/amd/indexer_aiter.py diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index d88540a3b5..c45b87d04b 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/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 diff --git a/vllm/models/minimax_m3/amd/indexer_aiter.py b/vllm/models/minimax_m3/amd/indexer_aiter.py new file mode 100644 index 0000000000..adf505aa7d --- /dev/null +++ b/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, + ) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 71babc80ec..50cdb55183 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/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", ) diff --git a/vllm/models/minimax_m3/amd/ops/sparse_pa.py b/vllm/models/minimax_m3/amd/ops/sparse_pa.py index 1431bc579c..6d0ea00b1b 100644 --- a/vllm/models/minimax_m3/amd/ops/sparse_pa.py +++ b/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, diff --git a/vllm/models/minimax_m3/amd/sparse_attention_msa.py b/vllm/models/minimax_m3/amd/sparse_attention_msa.py index c5e0d31681..a2c49d13ab 100644 --- a/vllm/models/minimax_m3/amd/sparse_attention_msa.py +++ b/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 diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 148a20b9f3..50c03dd9bf 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/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]