@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 83252ea899

[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>

ykamiset committed Sep 9, 2026 at 17:40 UTC 83252ea899c6538eaa0c1fb31f28a92c661bbffc
6 files changed +1077 -58
tests/kernels/attention/test_minimax_m3.py
+52
@@ -1694,6 +1694,58 @@ def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
1694 with pytest.raises(ValueError, match="num_kv_heads == 1"):
1695 sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(2)
1696
1697 + # An indexer that emits the attend's page table addresses every head, so
1698 + # the pairing the check guards against does not arise.
1699 + assert sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(
1700 + 2, emits_sparse_block_table=True
1701 + )
1702 + # The AITER path carries its own backend: its builder is what rebases the
1703 + # block table the indexer's top-k resolves through, so selecting the impl
1704 + # without it would leave that field unset.
1705 + backend, impl_cls = sparse_attn_mod.select_main_backend_and_impl_cls(
1706 + topk_blocks=TOPK,
1707 + kv_cache_dtype="fp8_e4m3",
1708 + num_kv_heads=2,
1709 + emits_sparse_block_table=True,
1710 + )
1711 + assert impl_cls.__name__ == "MiniMaxM3SparseAiterPAImpl"
1712 + assert backend.__name__ == "MiniMaxM3SparseAiterPABackend"
1713 + assert backend.get_builder_cls().__name__ == "MiniMaxM3SparseAiterPAMetadataBuilder"
1714 +
1715 +
1716 +def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):
1717 + """The top-k emits the attend's page table, so it needs that attend.
1718 +
1719 + Selected without it, the indexer would reach its top-k with no table to
1720 + resolve through; falling back instead surfaces the fp8 index cache as the
1721 + configuration error it is.
1722 + """
1723 + import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod
1724 +
1725 + monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True)
1726 + kwargs = dict(
1727 + topk_blocks=TOPK,
1728 + sparse_block_size=BLOCK_SIZE,
1729 + num_index_heads=1,
1730 + index_head_dim=HEAD_DIM,
1731 + indexer_kv_dtype="fp8_e4m3",
1732 + max_model_len=8192,
1733 + )
1734 +
1735 + monkeypatch.setattr(
1736 + indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False
1737 + )
1738 + reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)
1739 + assert reason is not None and "AITER sparse PA attend" in reason
1740 +
1741 + # With the attend asked for, the gate moves on to the kernel-contract
1742 + # checks rather than stopping here.
1743 + monkeypatch.setattr(
1744 + indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True
1745 + )
1746 + reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)
1747 + assert reason is None or "AITER sparse PA attend" not in reason
1748 +
1749
1750 def test_indexer_cache_squeezes_to_contiguous_3d():
1751 """The indexer side cache is standardized 4D with H=1: under both layouts
vllm/models/minimax_m3/amd/indexer_aiter.py new
+646
@@ -0,0 +1,646 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AITER (ROCm) indexer impl for MiniMax M3.
4 +
5 +Scores index blocks and selects the top-k with AITER's fp8 MFMA kernels on both
6 +sides of the batch: ``pa_sparse_block_score_decode`` for the uniform-query-length
7 +decode rows and ``pa_sparse_block_score_prefill`` for the ragged prefill rows,
8 +then ``pa_sparse_block_topk`` for each. The two scoring passes share one tile
9 +body in AITER, so they agree block for block, and both encode the forced
10 +init/local blocks as the same sentinel scores the Triton indexer uses.
11 +
12 +The top-k also emits the attend's page table. The winners are already in its
13 +workgroup's LDS, so resolving them through the block table there costs one wave
14 +and replaces the separate Triton pass in ``ops.sparse_pa``; the rows it writes
15 +are one per (token, kv head) with the page ids folded head-minor, which is the
16 +layout ``pa_decode_gluon`` reads after it flattens the cache.
17 +
18 +The score kernels are built on ``v_mfma_f32_16x16x32_fp8_fp8``, so this impl
19 +requires an fp8 (e4m3) index cache and an fp8 index query -- the fused
20 +QK-norm/RoPE kernel emits both directly when the index cache is e4m3. See
21 +``aiter_indexer_unsupported_reason`` for the full set of limits;
22 +``select_aiter_indexer_impl_cls`` refuses to pick this impl unless they all
23 +hold, and the model falls back to the platform-neutral ``MiniMaxM3Indexer``.
24 +"""
25 +
26 +import math
27 +from dataclasses import dataclass
28 +from functools import cache
29 +from typing import ClassVar
30 +
31 +import torch
32 +from torch import nn
33 +
34 +from vllm.config import VllmConfig, get_current_vllm_config
35 +from vllm.config.attention import IndexerKVDType
36 +from vllm.forward_context import get_forward_context
37 +from vllm.logger import init_logger
38 +from vllm.models.minimax_m3.amd.ops.sparse_pa import ASM_PAGE_SIZE
39 +from vllm.models.minimax_m3.common.indexer import (
40 + MiniMaxM3IndexerBackend,
41 + MiniMaxM3IndexerCache,
42 + MiniMaxM3IndexerDecodeMetadata,
43 + MiniMaxM3IndexerImpl,
44 + MiniMaxM3IndexerMetadata,
45 + MiniMaxM3IndexerMetadataBuilder,
46 + MiniMaxM3IndexerPrefillMetadata,
47 +)
48 +
49 +# The single source of truth for whether the AITER attend was asked for; this
50 +# indexer is only usable alongside it.
51 +from vllm.models.minimax_m3.common.sparse_attention import (
52 + _minimax_m3_aiter_sparse_pa_requested,
53 +)
54 +from vllm.platforms import current_platform
55 +from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
56 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills
57 +from vllm.v1.kv_cache_interface import AttentionSpec
58 +
59 +logger = init_logger(__name__)
60 +
61 +# Compiled AITER MiniMax-M3 score/top-k contract
62 +MSA_TOPK_BLOCKS = 16
63 +MSA_SCORE_TYPE = "max"
64 +MSA_SPARSE_BLOCK_SIZE = 128
65 +MSA_INDEX_HEAD_DIM = 128
66 +# Wave width the top-k is written against: it gives one lane per output slot and
67 +# reads the score row in wave-wide strips.
68 +WAVE_SIZE = 64
69 +# Per-lane register slots the top-k can hold, which is what caps the context: a
70 +# row may span at most SLOTS_MAX * WAVE_SIZE blocks.
71 +SLOTS_MAX = 128
72 +MAX_SUPPORTED_BLOCKS = SLOTS_MAX * WAVE_SIZE
73 +# MFMA columns the score pass has, one per (query token, index head) pair.
74 +MFMA_COLS = 16
75 +# The fp8 MFMA the score kernels are built on exists on these targets only.
76 +SUPPORTED_ARCHS = ("gfx950",)
77 +
78 +
79 +def _pow2_ceil(n: int) -> int:
80 + return 1 << (n - 1).bit_length() if n >= 1 else 1
81 +
82 +
83 +def score_block_width(max_seq_len: int, block_size: int) -> int:
84 + """Block-axis width the top-k requires of the score buffer.
85 +
86 + Lanes read whole wave-wide strips with no tail guard and each holds a
87 + power-of-two count of them, so the axis is padded past the block count.
88 + """
89 + max_blk = math.ceil(max(max_seq_len, 1) / block_size)
90 + return _pow2_ceil(math.ceil(max_blk / WAVE_SIZE)) * WAVE_SIZE
91 +
92 +
93 +def aiter_indexer_max_decode_query_len(vllm_config: VllmConfig) -> int:
94 + """Longest query a decode row can carry, which spec decode is what sets.
95 +
96 + Mirrors ``_init_reorder_batch_threshold(1, supports_spec_as_decode=True)``,
97 + since that is what the builder splits the batch on and therefore what the
98 + decode kernel will actually be handed.
99 + """
100 + spec = vllm_config.speculative_config
101 + if spec is None or spec.num_speculative_tokens is None:
102 + return 1
103 + return 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens
104 +
105 +
106 +@cache
107 +def aiter_msa_kernels_unavailable_reason() -> str | None:
108 + """Return why the AITER MSA score/top-k ops cannot be imported, or None.
109 +
110 + They are a recent addition, so an AITER that predates them imports fine
111 + while these three names do not exist, and the failure would otherwise
112 + surface as an ImportError from the middle of a forward. ``compile_ops``
113 + binds them lazily, so this costs the module import only -- the kernel build
114 + itself still happens on the first call.
115 + """
116 + try:
117 + from aiter.ops.msa_attention import ( # noqa: F401
118 + pa_sparse_block_score_decode,
119 + pa_sparse_block_score_prefill,
120 + pa_sparse_block_topk,
121 + )
122 + except ImportError as exc:
123 + return f"AITER cannot supply the MSA score/top-k kernels ({exc})"
124 + return None
125 +
126 +
127 +def aiter_indexer_unsupported_reason(
128 + *,
129 + topk_blocks: int,
130 + sparse_block_size: int,
131 + num_index_heads: int,
132 + index_head_dim: int,
133 + indexer_kv_dtype: IndexerKVDType,
134 + max_model_len: int,
135 + max_decode_query_len: int = 1,
136 + score_type: str = "max",
137 +) -> str | None:
138 + """Return why this config cannot use the AITER indexer, or None if it can.
139 +
140 + Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache
141 + dtype, the compiled score/top-k contract, MFMA column limits, max context
142 + in blocks, and whether AITER exposes the MSA kernels.
143 + ``select_aiter_indexer_impl_cls`` logs the string and falls back when it is
144 + not None.
145 + """
146 + if not current_platform.is_rocm():
147 + return (
148 + "needs ROCm for the AITER fp8 MFMA score/top-k, "
149 + f"got platform={current_platform.device_type!r}"
150 + )
151 + if not _minimax_m3_aiter_sparse_pa_requested():
152 + # The top-k emits the attend's page table in page-16 numbering, which
153 + # only addresses the interleaved cache the AITER attend reads. Paired
154 + # with any other attend there is nowhere valid to write it, so this
155 + # indexer is not usable on its own.
156 + return (
157 + "needs the AITER sparse PA attend, whose page table its top-k "
158 + "emits (rocm_aiter_ops + shuffle KV cache layout)"
159 + )
160 + if indexer_kv_dtype not in ("fp8", "fp8_e4m3"):
161 + # The score kernels are fp8 MFMA; there is no bf16 instantiation.
162 + return (
163 + f"needs an fp8 e4m3 index cache, got indexer_kv_dtype={indexer_kv_dtype!r}"
164 + )
165 + from vllm.platforms.rocm import on_gfx950
166 +
167 + if not on_gfx950():
168 + return f"needs {' or '.join(SUPPORTED_ARCHS)} for the fp8 MFMA"
169 + if score_type != MSA_SCORE_TYPE:
170 + return f"needs score_type={MSA_SCORE_TYPE!r}, got score_type={score_type!r}"
171 + if topk_blocks != MSA_TOPK_BLOCKS:
172 + return f"needs topk_blocks={MSA_TOPK_BLOCKS}, got topk_blocks={topk_blocks}"
173 + if sparse_block_size != MSA_SPARSE_BLOCK_SIZE:
174 + return (
175 + f"needs sparse_block_size={MSA_SPARSE_BLOCK_SIZE}, "
176 + f"got sparse_block_size={sparse_block_size}"
177 + )
178 + if index_head_dim != MSA_INDEX_HEAD_DIM:
179 + return (
180 + f"needs index_head_dim={MSA_INDEX_HEAD_DIM}, "
181 + f"got index_head_dim={index_head_dim}"
182 + )
183 + if num_index_heads > MFMA_COLS:
184 + return f"num_index_heads={num_index_heads} exceeds the {MFMA_COLS} MFMA columns"
185 + # A decode row's whole query shares one MFMA tile, one column per (token,
186 + # head) pair, so spec decode trades columns against the head count.
187 + if num_index_heads * max_decode_query_len > MFMA_COLS:
188 + return (
189 + f"num_index_heads={num_index_heads} x max_decode_query_len="
190 + f"{max_decode_query_len} exceeds the {MFMA_COLS} MFMA columns"
191 + )
192 + max_blocks = math.ceil(max_model_len / sparse_block_size)
193 + if max_blocks > MAX_SUPPORTED_BLOCKS:
194 + return (
195 + f"max_model_len={max_model_len} needs {max_blocks} blocks per row, "
196 + f"more than the top-k's {MAX_SUPPORTED_BLOCKS}"
197 + )
198 + return aiter_msa_kernels_unavailable_reason()
199 +
200 +
201 +class MiniMaxM3IndexerAiterBackend(MiniMaxM3IndexerBackend):
202 + """Indexer side-cache backend selecting the AITER builder."""
203 +
204 + @staticmethod
205 + def get_builder_cls() -> type["MiniMaxM3IndexerAiterMetadataBuilder"]:
206 + return MiniMaxM3IndexerAiterMetadataBuilder
207 +
208 +
209 +@dataclass
210 +class MiniMaxM3IndexerAiterMetadata(MiniMaxM3IndexerMetadata):
211 + """Adds the per-row shape the ragged top-k needs.
212 +
213 + The uniform decode rows are recovered inside the kernel from ``seq_lens`` and
214 + the shared query length, which also clamps cudagraph padding rows to nothing.
215 + Prefill rows have no such shape, so it is materialized here once per forward
216 + and shared by every layer -- which is also where the emitted page table's
217 + tail block comes from, since a block count alone does not say how many tokens
218 + the last block holds.
219 + """
220 +
221 + # [num_prefill_tokens] int32, cdiv(position + 1, sparse_block_size) per prefill row.
222 + prefill_num_valid_pages: torch.Tensor | None = None
223 + # [num_prefill_tokens] int32, the request each prefill row belongs to.
224 + prefill_row_req_id: torch.Tensor | None = None
225 + # [num_prefill_tokens] int32, causal token count (position + 1) per row.
226 + prefill_kv_lens: torch.Tensor | None = None
227 +
228 +
229 +class MiniMaxM3IndexerAiterMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
230 + """The Triton indexer's metadata plus the prefill rows' causal shape."""
231 +
232 + def __init__(
233 + self,
234 + kv_cache_spec: AttentionSpec,
235 + layer_names: list[str],
236 + vllm_config: VllmConfig,
237 + device: torch.device,
238 + ) -> None:
239 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
240 + hf_config = vllm_config.model_config.hf_config
241 + text_config = getattr(hf_config, "text_config", hf_config)
242 + self.sparse_block_size = int(
243 + text_config.sparse_attention_config["sparse_block_size"]
244 + )
245 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
246 + # Companions to the base's num_valid_pages_buffer, for the two vectors
247 + # only the emitted table needs.
248 + self.row_req_id_buffer = torch.empty(
249 + max_tokens, dtype=torch.int32, device=device
250 + )
251 + self.kv_lens_buffer = torch.empty(max_tokens, dtype=torch.int32, device=device)
252 +
253 + def build(
254 + self,
255 + common_prefix_len: int,
256 + common_attn_metadata: CommonAttentionMetadata,
257 + fast_build: bool = False,
258 + ) -> MiniMaxM3IndexerAiterMetadata:
259 + num_reqs = common_attn_metadata.num_reqs
260 + num_tokens = common_attn_metadata.num_actual_tokens
261 + query_start_loc = common_attn_metadata.query_start_loc
262 + seq_lens = common_attn_metadata.seq_lens
263 + block_table = common_attn_metadata.block_table_tensor
264 +
265 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
266 + split_decodes_and_prefills(
267 + common_attn_metadata,
268 + decode_threshold=self.reorder_batch_threshold,
269 + require_uniform=True,
270 + )
271 + )
272 + assert num_decodes + num_prefills == num_reqs
273 + assert num_decode_tokens + num_prefill_tokens == num_tokens
274 +
275 + # Decode-first batch: context lengths into the stable cudagraph buffer.
276 + context_lens = self.context_len_buffer[:num_reqs]
277 + context_lens.copy_(
278 + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
279 + )
280 +
281 + prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None
282 + prefill_num_valid_pages: torch.Tensor | None = None
283 + prefill_row_req_id: torch.Tensor | None = None
284 + prefill_kv_lens: torch.Tensor | None = None
285 + if num_prefills > 0:
286 + cu_seqlens_q = (query_start_loc[num_decodes:] - num_decode_tokens).to(
287 + torch.int32
288 + )
289 + prefill_metadata = MiniMaxM3IndexerPrefillMetadata(
290 + cu_seqlens_q=cu_seqlens_q,
291 + seq_lens=seq_lens[num_decodes:],
292 + context_lens=context_lens[num_decodes:],
293 + block_table=block_table[num_decodes:],
294 + max_query_len=common_attn_metadata.max_query_len,
295 + max_seq_len=common_attn_metadata.max_seq_len,
296 + )
297 + # A prefill row sees its own position, so its causal length and block
298 + # count both follow from that alone; the request it belongs to comes
299 + # from the query offsets. Prefill batches are never captured, so the
300 + # stable buffers are only being reused here, not required.
301 + positions = common_attn_metadata.positions
302 + assert positions is not None
303 + row_positions = positions[num_decode_tokens:num_tokens]
304 + prefill_num_valid_pages = self.num_valid_pages_buffer[
305 + num_decode_tokens:num_tokens
306 + ]
307 + prefill_num_valid_pages.copy_(
308 + row_positions // self.sparse_block_size + 1, non_blocking=True
309 + )
310 + prefill_kv_lens = self.kv_lens_buffer[num_decode_tokens:num_tokens]
311 + prefill_kv_lens.copy_(row_positions + 1, non_blocking=True)
312 + prefill_row_req_id = self.row_req_id_buffer[num_decode_tokens:num_tokens]
313 + prefill_row_req_id.copy_(
314 + torch.searchsorted(
315 + cu_seqlens_q[1:].contiguous(),
316 + torch.arange(
317 + num_prefill_tokens,
318 + dtype=torch.int32,
319 + device=cu_seqlens_q.device,
320 + ),
321 + right=True,
322 + ),
323 + non_blocking=True,
324 + )
325 +
326 + decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None
327 + if num_decodes > 0:
328 + qsl_cpu = common_attn_metadata.query_start_loc_cpu
329 + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
330 + decode_query_len = int(query_lens_cpu[0].item())
331 + assert decode_query_len > 0
332 + assert torch.all(
333 + (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
334 + )
335 + assert num_decode_tokens == num_decodes * decode_query_len
336 + decode_metadata = MiniMaxM3IndexerDecodeMetadata(
337 + seq_lens=seq_lens[:num_decodes],
338 + block_table=block_table[:num_decodes],
339 + max_seq_len=common_attn_metadata.max_seq_len,
340 + decode_query_len=decode_query_len,
341 + max_decode_query_len=self.max_decode_query_len,
342 + )
343 +
344 + return MiniMaxM3IndexerAiterMetadata(
345 + seq_lens=seq_lens,
346 + max_seq_len=common_attn_metadata.max_seq_len,
347 + slot_mapping=common_attn_metadata.slot_mapping,
348 + num_actual_tokens=num_tokens,
349 + num_decodes=num_decodes,
350 + num_decode_tokens=num_decode_tokens,
351 + num_prefills=num_prefills,
352 + num_prefill_tokens=num_prefill_tokens,
353 + prefill=prefill_metadata,
354 + decode=decode_metadata,
355 + prefill_num_valid_pages=prefill_num_valid_pages,
356 + prefill_row_req_id=prefill_row_req_id,
357 + prefill_kv_lens=prefill_kv_lens,
358 + )
359 +
360 +
361 +class MiniMaxM3IndexerAiterImpl(MiniMaxM3IndexerImpl):
362 + """AITER fp8 score + top-k for both prefill and decode."""
363 +
364 + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerAiterBackend
365 +
366 + def __init__(self, **kwargs) -> None:
367 + super().__init__(**kwargs)
368 + # Both passes are v_mfma_f32_16x16x32_fp8_fp8 with no bf16 instantiation,
369 + # so an index cache of any other dtype would be read as e4m3 bytes.
370 + # The selector will not get here, but nothing else may either.
371 + if self.indexer_kv_dtype not in ("fp8", "fp8_e4m3"):
372 + raise ValueError(
373 + "The AITER indexer requires an fp8 e4m3 index cache, got "
374 + f"indexer_kv_dtype={self.indexer_kv_dtype!r}"
375 + )
376 + # Shared, stable-address page table + per-row context bound the top-k
377 + # emits for the attend. Owned by the model so one allocation serves
378 + # every layer; left None when it reserved none, in which case the
379 + # attend rebuilds the table itself.
380 + self.sparse_bt_buffer: torch.Tensor | None = None
381 + self.sparse_ctx_buffer: torch.Tensor | None = None
382 +
383 + @property
384 + def pages_per_block(self) -> int:
385 + """Physical pages one selected block expands into for the attend."""
386 + return self.block_size // ASM_PAGE_SIZE
387 +
388 + def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]:
389 + """The page table and context rows covering score rows ``[lo, hi)``.
390 +
391 + One table row per (token, kv head), head minor, which is the order
392 + ``pa_decode_gluon`` reads once it flattens the cache. Slices of the
393 + shared buffers rather than copies, so the attend sees the writes; a
394 + model that reserved no buffers gets throwaway ones, and the attend
395 + rebuilds the table itself in that case.
396 + """
397 + kv_heads = self.num_kv_heads
398 + if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None:
399 + rows = (hi - lo) * kv_heads
400 + device = self.index_cache.kv_cache.device
401 + return (
402 + torch.empty(
403 + (rows, self.topk_blocks * self.pages_per_block),
404 + dtype=torch.int32,
405 + device=device,
406 + ),
407 + torch.empty(rows, dtype=torch.int32, device=device),
408 + )
409 + return (
410 + self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads],
411 + self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads],
412 + )
413 +
414 + def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor:
415 + """Score buffer for ``rows`` query rows.
416 +
417 + Left uninitialized on purpose: the score pass writes every block up to
418 + the longest row it covers, and the top-k reads only the blocks its own
419 + row can see, so nothing downstream observes the padded tail. Filling it
420 + would cost a write over the whole [H, rows, width] extent, which at long
421 + context is the largest tensor in the indexer.
422 + """
423 + return torch.empty(
424 + (
425 + self.num_index_heads,
426 + rows,
427 + score_block_width(max_seq_len, self.block_size),
428 + ),
429 + dtype=torch.float32,
430 + device=self.index_cache.kv_cache.device,
431 + )
432 +
433 + def forward(
434 + self,
435 + index_query: torch.Tensor,
436 + *,
437 + decode_page16_block_table: torch.Tensor | None = None,
438 + prefill_page16_block_table: torch.Tensor | None = None,
439 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
440 + from aiter.ops.msa_attention import (
441 + pa_sparse_block_score_decode,
442 + pa_sparse_block_score_prefill,
443 + pa_sparse_block_topk,
444 + )
445 +
446 + attn_metadata = get_forward_context().attn_metadata
447 + if not isinstance(attn_metadata, dict):
448 + return None, None # profiling run; caches unbound
449 + md = attn_metadata[self.index_cache.prefix]
450 + assert isinstance(md, MiniMaxM3IndexerAiterMetadata)
451 + # The emitted page table addresses the main cache, whose blocks are a
452 + # different group's than the index cache's, so the top-k resolves the
453 + # selection through the attend's block table -- only the score pass reads
454 + # the index cache and takes the indexer's. It has to be the page-16
455 + # rebase of that table, which the attend's own metadata builder does
456 + # once per step; the indexer's metadata cannot reach another group's
457 + # blocks, so the layer reads it off the attend and hands it in.
458 + num_tokens = md.num_actual_tokens
459 + nd = md.num_decode_tokens
460 + iq = index_query[:num_tokens].view(
461 + -1, self.num_index_heads, self.index_head_dim
462 + )
463 + kv = self.index_cache.kv_cache
464 +
465 + # Both sides write into the single shared persistent buffer (decode at
466 + # [:, :nd], prefill at [:, nd:]) and return views into it. The top-k
467 + # takes the head/row strides, so these slices need no copy back.
468 + buf = self.topk_indices_buffer
469 + if buf is None:
470 + buf = torch.empty(
471 + (self.num_index_heads, num_tokens, self.topk_blocks),
472 + dtype=torch.int32,
473 + device=iq.device,
474 + )
475 +
476 + decode_topk: torch.Tensor | None = None
477 + prefill_topk: torch.Tensor | None = None
478 +
479 + if md.num_decodes > 0:
480 + d = md.decode
481 + assert d is not None
482 + assert decode_page16_block_table is not None, (
483 + "the AITER indexer's top-k emits the attend's page table and "
484 + "needs the page-16 rebase of the attend's decode block table"
485 + )
486 + score = self._new_score(nd, d.max_seq_len)
487 + pa_sparse_block_score_decode(
488 + iq[:nd],
489 + kv,
490 + score,
491 + d.block_table,
492 + d.seq_lens,
493 + init_blocks=self.init_blocks,
494 + local_blocks=self.local_blocks,
495 + query_len=d.decode_query_len,
496 + max_seq_len=d.max_seq_len,
497 + )
498 + decode_topk = buf[:, :nd, :]
499 + sparse_bt, sparse_ctx = self._table_rows(0, nd)
500 + # A decode row's causal length is seq_len - query_len + token + 1,
501 + # which the kernel derives itself, so the emitted table covers
502 + # speculative rows without any extra per-row shape.
503 + pa_sparse_block_topk(
504 + score,
505 + decode_topk,
506 + decode_page16_block_table,
507 + d.seq_lens,
508 + sparse_bt,
509 + sparse_ctx,
510 + max_seq_len=d.max_seq_len,
511 + block_size=self.block_size,
512 + query_len=d.decode_query_len,
513 + num_kv_heads=self.num_kv_heads,
514 + pages_per_block=self.pages_per_block,
515 + )
516 +
517 + if md.num_prefills > 0:
518 + p = md.prefill
519 + assert p is not None
520 + assert prefill_page16_block_table is not None, (
521 + "the AITER indexer's top-k emits the attend's page table and "
522 + "needs the page-16 rebase of the attend's prefill block table"
523 + )
524 + assert md.prefill_num_valid_pages is not None
525 + assert md.prefill_row_req_id is not None
526 + assert md.prefill_kv_lens is not None
527 + score = self._new_score(num_tokens - nd, p.max_seq_len)
528 + pa_sparse_block_score_prefill(
529 + iq[nd:],
530 + kv,
531 + score,
532 + p.block_table,
533 + p.cu_seqlens_q,
534 + p.seq_lens,
535 + init_blocks=self.init_blocks,
536 + local_blocks=self.local_blocks,
537 + max_query_len=p.max_query_len,
538 + max_seq_len=p.max_seq_len,
539 + )
540 + prefill_topk = buf[:, nd:num_tokens, :]
541 + sparse_bt, sparse_ctx = self._table_rows(nd, num_tokens)
542 + # Ragged rows carry their request and causal length explicitly: the
543 + # block count alone cannot place the tail block the table ends on.
544 + pa_sparse_block_topk(
545 + score,
546 + prefill_topk,
547 + prefill_page16_block_table,
548 + p.seq_lens,
549 + sparse_bt,
550 + sparse_ctx,
551 + max_seq_len=p.max_seq_len,
552 + block_size=self.block_size,
553 + num_valid_pages=md.prefill_num_valid_pages,
554 + row_req_id=md.prefill_row_req_id,
555 + kv_lens=md.prefill_kv_lens,
556 + num_kv_heads=self.num_kv_heads,
557 + pages_per_block=self.pages_per_block,
558 + )
559 +
560 + return decode_topk, prefill_topk
561 +
562 +
563 +def select_aiter_indexer_impl_cls(
564 + *,
565 + topk_blocks: int,
566 + sparse_block_size: int,
567 + num_index_heads: int,
568 + index_head_dim: int,
569 + indexer_kv_dtype: IndexerKVDType,
570 + score_type: str = "max",
571 +) -> type[MiniMaxM3IndexerAiterImpl] | None:
572 + """The AITER indexer impl if this config can use it, else None.
573 +
574 + ``None`` sends the caller to the platform-neutral ``MiniMaxM3Indexer``,
575 + which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8
576 + index cache that lands here has nowhere to go.
577 + """
578 + reason = aiter_indexer_unsupported_reason(
579 + topk_blocks=topk_blocks,
580 + sparse_block_size=sparse_block_size,
581 + num_index_heads=num_index_heads,
582 + index_head_dim=index_head_dim,
583 + indexer_kv_dtype=indexer_kv_dtype,
584 + max_model_len=get_current_vllm_config().model_config.max_model_len,
585 + max_decode_query_len=aiter_indexer_max_decode_query_len(
586 + get_current_vllm_config()
587 + ),
588 + score_type=score_type,
589 + )
590 + if reason is not None:
591 + logger.info_once("MiniMax M3 indexer: AITER unavailable (%s)", reason)
592 + return None
593 + logger.info_once(
594 + "MiniMax M3 indexer: selected AITER (fp8 MFMA score + top-k) "
595 + "[topk_blocks=%d, indexer_kv_dtype=%s]",
596 + topk_blocks,
597 + indexer_kv_dtype,
598 + )
599 + return MiniMaxM3IndexerAiterImpl
600 +
601 +
602 +class MiniMaxM3AiterIndexer(nn.Module):
603 + """``MiniMaxM3Indexer``'s surface over the AITER impl.
604 +
605 + The platform-neutral wrapper picks its impl through ``common``'s selector
606 + and forwards a Triton-only set of fused-table arguments, neither of which
607 + can reach this impl without editing ``common``. This holds the same three
608 + members the attention layer uses -- ``impl``, ``index_cache``,
609 + ``num_index_heads`` -- and forwards the one argument AITER needs instead.
610 + """
611 +
612 + def __init__(
613 + self,
614 + *,
615 + impl_cls: type[MiniMaxM3IndexerAiterImpl],
616 + sparse_bt_buffer: torch.Tensor | None = None,
617 + sparse_ctx_buffer: torch.Tensor | None = None,
618 + **impl_kwargs,
619 + ) -> None:
620 + super().__init__()
621 + self.impl = impl_cls(**impl_kwargs)
622 + # Assigned rather than passed: the impl's base ``__init__`` lives in
623 + # common and takes no table buffers.
624 + self.impl.sparse_bt_buffer = sparse_bt_buffer
625 + self.impl.sparse_ctx_buffer = sparse_ctx_buffer
626 +
627 + @property
628 + def index_cache(self) -> MiniMaxM3IndexerCache:
629 + return self.impl.index_cache
630 +
631 + @property
632 + def num_index_heads(self) -> int:
633 + return self.impl.num_index_heads
634 +
635 + def forward(
636 + self,
637 + index_query: torch.Tensor,
638 + *,
639 + decode_page16_block_table: torch.Tensor | None = None,
640 + prefill_page16_block_table: torch.Tensor | None = None,
641 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
642 + return self.impl(
643 + index_query,
644 + decode_page16_block_table=decode_page16_block_table,
645 + prefill_page16_block_table=prefill_page16_block_table,
646 + )
vllm/models/minimax_m3/amd/model.py
+158 -36
@@ -87,17 +87,24 @@ from vllm.model_executor.models.utils import (
87 maybe_prefix,
88 )
89 from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
90 +from vllm.models.minimax_m3.amd.indexer_aiter import (
91 + MiniMaxM3AiterIndexer,
92 + select_aiter_indexer_impl_cls,
93 +)
94 from vllm.models.minimax_m3.amd.ops import (
95 gemma_fused_add_rmsnorm,
96 gemma_rmsnorm,
97 swiglu_oai_split,
98 )
99 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
100 + ASM_PAGE_SIZE,
101 minimax_m3_alloc_sparse_block_table,
102 minimax_m3_sparse_block_page_stride,
103 )
104 from vllm.models.minimax_m3.amd.sparse_attention_msa import (
105 + MiniMaxM3SparseAiterPADecodeMetadata,
106 MiniMaxM3SparseAiterPAImpl,
107 + MiniMaxM3SparseAiterPAPrefillMetadata,
108 )
109 from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
110 from vllm.models.minimax_m3.common.mm_preprocess import (
@@ -111,7 +118,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
118 MiniMaxM3SparseMetadata,
119 minimax_m3_rebase_slots_to_page16,
120 minimax_m3_use_aiter_sparse_pa,
114 - select_main_impl_cls,
121 + select_main_backend_and_impl_cls,
122 )
123 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
124 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -596,6 +603,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
603 prefix: str = "",
604 cache_config: CacheConfig | None = None,
605 topk_indices_buffer: torch.Tensor | None = None,
606 + sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
607 ) -> None:
608 super().__init__()
609 self.hidden_size = config.hidden_size
@@ -684,18 +692,45 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
692 # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
693 set_default_quant_scales(self, register_buffer=True)
694
695 + # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main cache
696 + # (--attention-config '{"indexer_kv_dtype": ...}'). fp8 e4m3 is what the
697 + # AITER indexer needs; bf16 keeps the Triton indexer.
698 + self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
699 + "bf16"
700 + )
701 +
702 # Shared top-k buffer: the indexer writes the selected blocks into it and
703 # the attend impl reads them back (no Python value crosses the break).
704 self.topk_indices_buffer = topk_indices_buffer
690 - self.attn_backend = MiniMaxM3SparseBackend
691 - # Indexer and main attention are separate impls. On ROCm the SM100 gate
692 - # is always False, so both pick Triton and the index cache stays bf16.
693 - # impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
694 - self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
705 + # Indexer and main attention are separate impls. The AITER indexer is
706 + # ROCm-only, so it is selected here rather than inside the neutral
707 + # MiniMaxM3Indexer, which knows nothing about it and would pick Triton.
708 + indexer_impl_cls = select_aiter_indexer_impl_cls(
709 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
710 + sparse_block_size=sparse_cfg["sparse_block_size"],
711 + num_index_heads=self.num_idx_heads,
712 + index_head_dim=self.idx_head_dim,
713 + indexer_kv_dtype=self.indexer_kv_dtype,
714 + score_type=sparse_cfg.get("sparse_score_type", "max"),
715 + )
716 + # The attend gate below needs this: an emitted table is what lets it
717 + # serve more than one KV head per rank. Buffers to write the table into
718 + # only arrive when the model already resolved the attend to the AITER
719 + # path the table addresses, so that is not rechecked here.
720 + self.indexer_emits_table = (
721 + indexer_impl_cls is not None and sparse_table_buffers is not None
722 + )
723 + # The backend names the metadata builder, which for the AITER path is
724 + # the one that rebases the block table the indexer's top-k resolves
725 + # through, so both have to come from the same selection.
726 + self.attn_backend, main_impl_cls = select_main_backend_and_impl_cls(
727 topk_blocks=sparse_cfg["sparse_topk_blocks"],
728 kv_cache_dtype=self.kv_cache_dtype,
729 num_kv_heads=self.num_kv_heads,
698 - )(
730 + emits_sparse_block_table=self.indexer_emits_table,
731 + )
732 + # impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
733 + self.impl: MiniMaxM3SparseImpl = main_impl_cls( # type: ignore[assignment]
734 self.num_heads,
735 self.head_dim,
736 self.scaling,
@@ -704,14 +739,14 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
739 topk_blocks=sparse_cfg["sparse_topk_blocks"],
740 sparse_block_size=sparse_cfg["sparse_block_size"],
741 )
707 - self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(self.num_kv_heads)
742 + self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
743 + self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table
744 + )
745 self.kv_cache_k = torch.tensor([])
746 self.kv_cache_v = torch.tensor([])
747 self._aiter_sparse_pa_cache_data_ptr = 0
748 self._aiter_sparse_pa_block_page_stride = 0
712 - # Self-contained nn.Module: owns its side cache, selects its impl in init
713 - # (Triton on ROCm, where the SM100 gate is always False).
714 - self.indexer = MiniMaxM3Indexer(
749 + indexer_kwargs = dict(
750 num_kv_heads=self.num_kv_heads,
751 scale=self.scaling,
752 topk_blocks=sparse_cfg["sparse_topk_blocks"],
@@ -723,8 +758,24 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
758 local_blocks=sparse_cfg.get("sparse_local_block", 0),
759 score_type=sparse_cfg.get("sparse_score_type", "max"),
760 cache_config=cache_config,
761 + indexer_kv_dtype=self.indexer_kv_dtype,
762 topk_indices_buffer=topk_indices_buffer,
763 )
764 + self.sparse_bt_buffer: torch.Tensor | None = None
765 + self.sparse_ctx_buffer: torch.Tensor | None = None
766 + if indexer_impl_cls is None:
767 + # Self-contained nn.Module: owns its side cache, selects its impl.
768 + self.indexer = MiniMaxM3Indexer(**indexer_kwargs)
769 + else:
770 + if self.indexer_emits_table:
771 + assert sparse_table_buffers is not None
772 + self.sparse_bt_buffer, self.sparse_ctx_buffer = sparse_table_buffers
773 + self.indexer = MiniMaxM3AiterIndexer(
774 + impl_cls=indexer_impl_cls,
775 + sparse_bt_buffer=self.sparse_bt_buffer,
776 + sparse_ctx_buffer=self.sparse_ctx_buffer,
777 + **indexer_kwargs,
778 + )
779
780 # Register the main K/V cache so the KV-cache manager allocates it.
781 compilation_config = vllm_config.compilation_config
@@ -952,7 +1003,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
1003 )
1004 else:
1005 index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
955 - index_q = qkv.new_empty((num_tokens, self.index_q_size))
1006 + # index_q matches the index-K cache dtype (e4m3 for the AITER fp8
1007 + # score path); the fused kernel emits fp8 directly when this buffer
1008 + # is e4m3.
1009 + index_q = qkv.new_empty(
1010 + (num_tokens, self.index_q_size),
1011 + dtype=self.indexer.index_cache.dtype,
1012 + )
1013 if self.use_aiter_sparse_pa:
1014 ops.fused_minimax_m3_qknorm_rope_kv_insert(
1015 qkv,
@@ -1034,31 +1091,56 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
1091 decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None
1092 if not self.skip_index_topk:
1093 assert index_query is not None
1037 - attn_metadata = get_forward_context().attn_metadata
1038 - if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
1039 - main_md = attn_metadata[self.layer_name]
1040 - assert isinstance(main_md, MiniMaxM3SparseMetadata)
1041 - if main_md.num_decodes > 0:
1042 - d = main_md.decode
1043 - assert d is not None
1044 - topk = self.topk_indices_buffer
1045 - assert topk is not None
1046 - decode_sparse_table = minimax_m3_alloc_sparse_block_table(
1047 - topk[:, : main_md.num_decode_tokens, :]
1048 - )
1049 - block_page_stride = self._aiter_sparse_pa_block_page_stride
1050 - assert block_page_stride > 0
1051 - self.indexer(
1052 - index_query,
1053 - attention_block_table=d.block_table,
1054 - sparse_block_table_out=decode_sparse_table[0],
1055 - sparse_context_lens_out=decode_sparse_table[1],
1056 - block_page_stride=block_page_stride,
1057 - )
1094 + # The AITER indexer emits the attend table into its persistent
1095 + # buffers, resolving its selection through the page-16 rebase of the
1096 + # attend's block table. The Triton indexer can instead fuse decode
1097 + # top-k with main's per-forward sparse-table allocation.
1098 + if self.indexer_emits_table:
1099 + attn_metadata = get_forward_context().attn_metadata
1100 + decode_page16 = prefill_page16 = None
1101 + if isinstance(attn_metadata, dict):
1102 + main_md = attn_metadata[self.layer_name]
1103 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
1104 + # Rebased once per step by the AITER attend's builder, which
1105 + # this path selected along with the impl.
1106 + d, p = main_md.decode, main_md.prefill
1107 + if d is not None:
1108 + assert isinstance(d, MiniMaxM3SparseAiterPADecodeMetadata)
1109 + decode_page16 = d.page16_block_table
1110 + if p is not None:
1111 + assert isinstance(p, MiniMaxM3SparseAiterPAPrefillMetadata)
1112 + prefill_page16 = p.page16_block_table
1113 + self.indexer(
1114 + index_query,
1115 + decode_page16_block_table=decode_page16,
1116 + prefill_page16_block_table=prefill_page16,
1117 + )
1118 + else:
1119 + attn_metadata = get_forward_context().attn_metadata
1120 + if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
1121 + main_md = attn_metadata[self.layer_name]
1122 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
1123 + if main_md.num_decodes > 0:
1124 + d = main_md.decode
1125 + assert d is not None
1126 + topk = self.topk_indices_buffer
1127 + assert topk is not None
1128 + decode_sparse_table = minimax_m3_alloc_sparse_block_table(
1129 + topk[:, : main_md.num_decode_tokens, :]
1130 + )
1131 + block_page_stride = self._aiter_sparse_pa_block_page_stride
1132 + assert block_page_stride > 0
1133 + self.indexer(
1134 + index_query,
1135 + attention_block_table=d.block_table,
1136 + sparse_block_table_out=decode_sparse_table[0],
1137 + sparse_context_lens_out=decode_sparse_table[1],
1138 + block_page_stride=block_page_stride,
1139 + )
1140 + else:
1141 + self.indexer(index_query)
1142 else:
1143 self.indexer(index_query)
1060 - else:
1061 - self.indexer(index_query)
1144 if self.use_aiter_sparse_pa:
1145 assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
1146 return self.impl.forward(
@@ -1081,6 +1163,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
1163 force_sparse_attn: bool = False,
1164 force_moe: bool = False,
1165 topk_indices_buffer: torch.Tensor | None = None,
1166 + sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
1167 ) -> None:
1168 super().__init__()
1169 self.hidden_size = config.hidden_size
@@ -1101,6 +1184,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
1184 prefix=f"{prefix}.self_attn",
1185 cache_config=cache_config,
1186 topk_indices_buffer=topk_indices_buffer,
1187 + sparse_table_buffers=sparse_table_buffers,
1188 )
1189 else:
1190 self.self_attn = MiniMaxM3Attention(
@@ -1189,15 +1273,52 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1273 # layers (mirrors DeepseekV4); the indexer writes its per-head decode/
1274 # prefill block selection into it, the attend reads it back.
1275 sparse_cfg = getattr(config, "sparse_attention_config", None)
1276 + self.sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] | None = None
1277 if sparse_cfg is not None:
1278 tp_size = get_tensor_model_parallel_world_size()
1279 num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
1280 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1281 self.topk_indices_buffer = torch.empty(
1282 num_index_heads,
1197 - vllm_config.scheduler_config.max_num_batched_tokens,
1283 + max_tokens,
1284 sparse_cfg["sparse_topk_blocks"],
1285 dtype=torch.int32,
1286 )
1287 +
1288 + num_kv_heads = max(1, config.num_key_value_heads // tp_size)
1289 + indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
1290 + "bf16"
1291 + )
1292 + # Probed with the head count the layer will pass (it derives
1293 + # num_idx_heads from num_kv_heads, not from sparse_num_index_heads),
1294 + # so this cannot land on the other side of the MFMA column limit
1295 + # from the selection the layer makes.
1296 + emits_table = (
1297 + select_aiter_indexer_impl_cls(
1298 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
1299 + sparse_block_size=sparse_cfg["sparse_block_size"],
1300 + num_index_heads=num_kv_heads,
1301 + index_head_dim=sparse_cfg["sparse_index_dim"],
1302 + indexer_kv_dtype=indexer_kv_dtype,
1303 + score_type=sparse_cfg.get("sparse_score_type", "max"),
1304 + )
1305 + is not None
1306 + )
1307 + if emits_table and minimax_m3_use_aiter_sparse_pa(
1308 + num_kv_heads, emits_sparse_block_table=True
1309 + ):
1310 + # One row per (token, kv head): the table addresses the attend's
1311 + # cache, not the index cache.
1312 + rows = max_tokens * num_kv_heads
1313 + self.sparse_table_buffers = (
1314 + torch.empty(
1315 + rows,
1316 + sparse_cfg["sparse_topk_blocks"]
1317 + * (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE),
1318 + dtype=torch.int32,
1319 + ),
1320 + torch.empty(rows, dtype=torch.int32),
1321 + )
1322 else:
1323 self.topk_indices_buffer = None
1324
@@ -1209,6 +1330,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1330 cache_config=cache_config,
1331 quant_config=quant_config,
1332 topk_indices_buffer=self.topk_indices_buffer,
1333 + sparse_table_buffers=self.sparse_table_buffers,
1334 ),
1335 prefix=f"{prefix}.layers",
1336 )
vllm/models/minimax_m3/amd/ops/sparse_pa.py
+45 -8
@@ -312,7 +312,9 @@ def _insert_index_cache_kernel(
312 )
313 mask = (slot >= 0) & (offs_d < HEAD_DIM)
314 value = tl.load(src, mask=offs_d < HEAD_DIM, other=0.0)
315 - tl.store(dst, value, mask=mask)
315 + # The side cache carries its own dtype (bf16, or e4m3 for the fp8 indexer),
316 + # so convert rather than leaning on the store's implicit cast.
317 + tl.store(dst, value.to(dst.dtype.element_ty), mask=mask)
318
319
320 @torch.no_grad()
@@ -368,6 +370,32 @@ def minimax_m3_sparse_block_page_stride(
370 return PAGES_PER_SPARSE_BLOCK * (2 if _sides_are_packed(k_cache, v_cache) else 1)
371
372
373 +# K/V sides sharing one block, which is how many sides' pages a block spans.
374 +# The attend publishes two head slots and the indexer's side cache one, and
375 +# specs that mix HNC shapes narrow the resolved layout to a block-compact one
376 +# (LHBNC is not), so this path always lands on an interleaved block.
377 +PAGE16_SIDES_PER_BLOCK = 2
378 +
379 +
380 +def minimax_m3_rebase_block_table_to_page16(
381 + block_table: torch.Tensor,
382 + out: torch.Tensor | None = None,
383 +) -> torch.Tensor:
384 + """Rebase a logical page table onto AITER's page-16 page numbering.
385 +
386 + The page ids AITER's top-k emits for the attend are
387 + ``block_table[blk] * pages_per_block + j``, and ``pages_per_block`` is fixed
388 + at AITER build time to one side's pages. An interleaved block spans both
389 + sides, so the only way to reach its pages through that kernel is to hand it
390 + a table already scaled to the wider stride. The caller does this once per
391 + step, since every sparse layer resolves its selection through the same
392 + table.
393 + """
394 + if out is None:
395 + out = torch.empty_like(block_table)
396 + return torch.mul(block_table, PAGE16_SIDES_PER_BLOCK, out=out)
397 +
398 +
399 def _gluon_scale_arg(
400 scale: torch.Tensor | None,
401 *,
@@ -461,14 +489,23 @@ def minimax_m3_sparse_attn_prefill_aiter(
489 output: torch.Tensor,
490 k_scale: torch.Tensor | None = None,
491 v_scale: torch.Tensor | None = None,
492 + sparse_bt: torch.Tensor | None = None,
493 + sparse_ctx: torch.Tensor | None = None,
494 ) -> None:
465 - sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
466 - topk_idx,
467 - block_table,
468 - query_req_id,
469 - query_abs_pos,
470 - minimax_m3_sparse_block_page_stride(k_cache, v_cache),
471 - )
495 + if (sparse_bt is None) != (sparse_ctx is None):
496 + raise ValueError(
497 + "MiniMax-M3 prepared sparse block table and context lengths "
498 + "must be provided together"
499 + )
500 + if sparse_bt is None:
501 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
502 + topk_idx,
503 + block_table,
504 + query_req_id,
505 + query_abs_pos,
506 + minimax_m3_sparse_block_page_stride(k_cache, v_cache),
507 + )
508 + assert sparse_ctx is not None
509 _run_gluon_decode(
510 q,
511 k_cache,
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+152 -5
@@ -2,16 +2,118 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """MSA AITER block-sparse attend for MiniMax M3."""
4
5 +from dataclasses import dataclass
6 +
7 import torch
8
9 +from vllm.config import VllmConfig
10 from vllm.forward_context import get_forward_context
11 from vllm.models.minimax_m3.common.sparse_attention import (
12 + MiniMaxM3SparseBackend,
13 + MiniMaxM3SparseDecodeMetadata,
14 MiniMaxM3SparseImpl,
15 MiniMaxM3SparseMetadata,
16 + MiniMaxM3SparseMetadataBuilder,
17 + MiniMaxM3SparsePrefillMetadata,
18 )
19 +from vllm.utils.math_utils import cdiv
20 from vllm.v1.attention.backend import (
21 AttentionLayer,
22 )
23 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
24 +from vllm.v1.kv_cache_interface import AttentionSpec
25 +
26 +
27 +class MiniMaxM3SparseAiterPABackend(MiniMaxM3SparseBackend):
28 + """MiniMax M3 backend carrying the AITER page-16 block-table rebase."""
29 +
30 + @staticmethod
31 + def get_builder_cls() -> type["MiniMaxM3SparseAiterPAMetadataBuilder"]:
32 + return MiniMaxM3SparseAiterPAMetadataBuilder
33 +
34 +
35 +@dataclass
36 +class MiniMaxM3SparseAiterPAPrefillMetadata(MiniMaxM3SparsePrefillMetadata):
37 + # ``block_table`` rebased onto AITER's page-16 numbering; see
38 + # ``MiniMaxM3SparseAiterPAMetadataBuilder.build``.
39 + page16_block_table: torch.Tensor | None = None
40 +
41 +
42 +@dataclass
43 +class MiniMaxM3SparseAiterPADecodeMetadata(MiniMaxM3SparseDecodeMetadata):
44 + page16_block_table: torch.Tensor | None = None
45 +
46 +
47 +class MiniMaxM3SparseAiterPAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
48 + """Adds the page-16 rebase of the block table the indexer's top-k needs.
49 +
50 + The AITER indexer's top-k emits the attend's page table, expanding a
51 + selected block into a compile-time number of pages -- one side's worth,
52 + while an interleaved block holds both -- so it has to resolve the selection
53 + through a table pre-scaled to the wider stride. Every sparse layer selects
54 + through the same table, so it is rebased here once per step rather than
55 + once per layer, alongside the slot mapping the base rebases for the writer.
56 + """
57 +
58 + def __init__(
59 + self,
60 + kv_cache_spec: AttentionSpec,
61 + layer_names: list[str],
62 + vllm_config: VllmConfig,
63 + device: torch.device,
64 + ) -> None:
65 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
66 + # Stable address for the same reason as the base's own buffers: a
67 + # captured decode replay writes through the address it captured.
68 + self.page16_block_table_buffer = torch.empty(
69 + (
70 + vllm_config.scheduler_config.max_num_seqs,
71 + cdiv(vllm_config.model_config.max_model_len, kv_cache_spec.block_size),
72 + ),
73 + dtype=torch.int32,
74 + device=device,
75 + )
76 +
77 + def build(
78 + self,
79 + common_prefix_len: int,
80 + common_attn_metadata: CommonAttentionMetadata,
81 + fast_build: bool = False,
82 + ) -> MiniMaxM3SparseMetadata:
83 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
84 + minimax_m3_rebase_block_table_to_page16,
85 + )
86 +
87 + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
88 + block_table = common_attn_metadata.block_table_tensor
89 + num_reqs, cols = block_table.shape
90 + assert (
91 + num_reqs <= self.page16_block_table_buffer.shape[0]
92 + and cols <= self.page16_block_table_buffer.shape[1]
93 + ), (
94 + f"block table {(num_reqs, cols)} exceeds the page-16 rebase buffer "
95 + f"{tuple(self.page16_block_table_buffer.shape)}"
96 + )
97 + page16 = minimax_m3_rebase_block_table_to_page16(
98 + block_table,
99 + out=self.page16_block_table_buffer[:num_reqs, :cols],
100 + )
101 +
102 + # Decode-first batch, so each side takes the rows its own block table
103 + # was sliced from. Reconstructed from the base's fields rather than
104 + # listed out, so a field added there does not silently drop here.
105 + nd = metadata.num_decodes
106 + if metadata.decode is not None:
107 + metadata.decode = MiniMaxM3SparseAiterPADecodeMetadata(
108 + **vars(metadata.decode),
109 + page16_block_table=page16[:nd],
110 + )
111 + if metadata.prefill is not None:
112 + metadata.prefill = MiniMaxM3SparseAiterPAPrefillMetadata(
113 + **vars(metadata.prefill),
114 + page16_block_table=page16[nd:],
115 + )
116 + return metadata
117
118
119 class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
@@ -28,8 +130,11 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
130 decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None,
131 ) -> torch.Tensor:
132 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
133 + PAGE16_SIDES_PER_BLOCK,
134 + PAGES_PER_SPARSE_BLOCK,
135 minimax_m3_sparse_attn_decode_aiter,
136 minimax_m3_sparse_attn_prefill_aiter,
137 + minimax_m3_sparse_block_page_stride,
138 )
139
140 attn_metadata = get_forward_context().attn_metadata
@@ -42,10 +147,16 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
147 num_tokens = main_md.num_actual_tokens
148 topk = layer.topk_indices_buffer # type: ignore[attr-defined]
149 assert topk is not None
45 - if self.num_kv_heads != 1:
150 + # Set only when the layer's indexer emitted the page table alongside its
151 + # selection; otherwise these stay None and the table is built below.
152 + sparse_bt_buf = getattr(layer, "sparse_bt_buffer", None)
153 + sparse_ctx_buf = getattr(layer, "sparse_ctx_buffer", None)
154 + kvh = self.num_kv_heads
155 + if kvh != 1 and sparse_bt_buf is None:
156 raise NotImplementedError(
47 - "MiniMax-M3 AITER sparse PA currently requires per-rank "
48 - f"num_kv_heads == 1, got {self.num_kv_heads}"
157 + "MiniMax-M3 AITER sparse PA needs the page table the indexer's "
158 + f"top-k emits to serve per-rank num_kv_heads == {kvh}; the "
159 + "Triton builders it falls back to address one head's cache."
160 )
161
162 hd = self.head_size
@@ -55,9 +166,31 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
166 k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
167 v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
168
169 + if sparse_bt_buf is not None:
170 + # An emitted table was scaled by the metadata builder, which reads
171 + # the packing off the layout rather than off these tensors. The two
172 + # only disagree if the resolved layout gives each side its own
173 + # plane, and a silent disagreement reads unrelated pages.
174 + stride = minimax_m3_sparse_block_page_stride(k_cache, v_cache)
175 + expected = PAGES_PER_SPARSE_BLOCK * PAGE16_SIDES_PER_BLOCK
176 + if stride != expected:
177 + raise RuntimeError(
178 + "MiniMax-M3 AITER sparse PA: the indexer emitted a page "
179 + f"table for a block spanning {expected} pages, but this "
180 + f"KV cache lays a block out over {stride}. The page-16 "
181 + "rebase assumes both K/V sides share a block."
182 + )
183 +
184 if main_md.num_decodes > 0:
185 d = main_md.decode
186 assert d is not None
187 + prepared_decode_table = decode_sparse_table
188 + if sparse_bt_buf is not None:
189 + assert sparse_ctx_buf is not None
190 + prepared_decode_table = (
191 + sparse_bt_buf[: nd * kvh],
192 + sparse_ctx_buf[: nd * kvh],
193 + )
194 minimax_m3_sparse_attn_decode_aiter(
195 q[:nd],
196 k_cache,
@@ -72,10 +205,14 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
205 v_scale=v_scale,
206 decode_query_len=d.decode_query_len,
207 sparse_block_table=(
75 - decode_sparse_table[0] if decode_sparse_table is not None else None
208 + prepared_decode_table[0]
209 + if prepared_decode_table is not None
210 + else None
211 ),
212 sparse_context_lens=(
78 - decode_sparse_table[1] if decode_sparse_table is not None else None
213 + prepared_decode_table[1]
214 + if prepared_decode_table is not None
215 + else None
216 ),
217 )
218
@@ -96,5 +233,15 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
233 out[nd:],
234 k_scale=k_scale,
235 v_scale=v_scale,
236 + sparse_bt=(
237 + None
238 + if sparse_bt_buf is None
239 + else sparse_bt_buf[nd * kvh : num_tokens * kvh]
240 + ),
241 + sparse_ctx=(
242 + None
243 + if sparse_ctx_buf is None
244 + else sparse_ctx_buf[nd * kvh : num_tokens * kvh]
245 + ),
246 )
247 return output
vllm/models/minimax_m3/common/sparse_attention.py
+24 -9
@@ -109,13 +109,22 @@ def minimax_m3_query_token_positions(
109 return query_req_id, query_abs_pos
110
111
112 -def minimax_m3_use_aiter_sparse_pa(num_kv_heads: int) -> bool:
113 - """Whether to use the ROCm AITER page-16 sparse PA prototype."""
112 +def minimax_m3_use_aiter_sparse_pa(
113 + num_kv_heads: int, *, emits_sparse_block_table: bool = False
114 +) -> bool:
115 + """Whether to use the ROCm AITER page-16 sparse PA prototype.
116 +
117 + More than one KV head per rank is only served when the layer's indexer
118 + emits the attend's page table itself, since the Triton builders this path
119 + otherwise falls back to address one head's cache. Whether an indexer does
120 + that is a ROCm-side fact, so its caller passes it in.
121 + """
122 requested = _minimax_m3_aiter_sparse_pa_requested()
115 - if requested and num_kv_heads != 1:
123 + if requested and num_kv_heads != 1 and not emits_sparse_block_table:
124 raise ValueError(
125 "MiniMax M3 AITER sparse paged attention requires "
118 - f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}."
126 + f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}, "
127 + "unless the indexer emits the attend's page table."
128 )
129 return requested
130
@@ -252,9 +261,9 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
261 # Every sparse layer shares one slot mapping, so the AITER page-16
262 # rebase is done here once per step instead of once per layer. Stable
263 # buffer for the same reason as the context lengths above.
255 - self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
256 - kv_cache_spec.num_kv_heads
257 - )
264 + # The request itself, not the gate: the buffer is the same whatever the
265 + # head count, and a builder cannot see which indexer its layers picked.
266 + self.use_aiter_sparse_pa = _minimax_m3_aiter_sparse_pa_requested()
267 self.page16_slot_mapping_buffer: torch.Tensor | None = None
268 if self.use_aiter_sparse_pa:
269 self.page16_slot_mapping_buffer = torch.empty(
@@ -514,6 +523,7 @@ def select_main_backend_and_impl_cls(
523 topk_blocks: int,
524 kv_cache_dtype: str,
525 num_kv_heads: int,
526 + emits_sparse_block_table: bool = False,
527 ) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]:
528 """Pick the main attention backend and implementation.
529
@@ -523,7 +533,9 @@ def select_main_backend_and_impl_cls(
533 back to Triton. The MSA modules are imported lazily to avoid import errors
534 on unsupported platforms.
535 """
526 - use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(num_kv_heads)
536 + use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
537 + num_kv_heads, emits_sparse_block_table=emits_sparse_block_table
538 + )
539 use_msa = (
540 current_platform.is_cuda()
541 and current_platform.is_device_capability_family(100)
@@ -541,10 +553,11 @@ def select_main_backend_and_impl_cls(
553 )
554 if use_aiter_sparse_pa:
555 from vllm.models.minimax_m3.amd.sparse_attention_msa import (
556 + MiniMaxM3SparseAiterPABackend,
557 MiniMaxM3SparseAiterPAImpl,
558 )
559
547 - return MiniMaxM3SparseBackend, MiniMaxM3SparseAiterPAImpl
560 + return MiniMaxM3SparseAiterPABackend, MiniMaxM3SparseAiterPAImpl
561 if use_msa:
562 from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
563 MiniMaxM3SparseMSABackend,
@@ -560,10 +573,12 @@ def select_main_impl_cls(
573 topk_blocks: int,
574 kv_cache_dtype: str,
575 num_kv_heads: int,
576 + emits_sparse_block_table: bool = False,
577 ) -> type[MiniMaxM3SparseImpl]:
578 """Backward-compatible implementation-only selector."""
579 return select_main_backend_and_impl_cls(
580 topk_blocks=topk_blocks,
581 kv_cache_dtype=kv_cache_dtype,
582 num_kv_heads=num_kv_heads,
583 + emits_sparse_block_table=emits_sparse_block_table,
584 )[1]