| 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 ( |
| 56 | AttentionBackend, |
| 57 | CommonAttentionMetadata, |
| 58 | max_decode_query_len, |
| 59 | ) |
| 60 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 61 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 62 | |
| 63 | logger = init_logger(__name__) |
| 64 | |
| 65 | # Compiled AITER MiniMax-M3 score/top-k contract |
| 66 | MSA_TOPK_BLOCKS = 16 |
| 67 | MSA_SCORE_TYPE = "max" |
| 68 | MSA_SPARSE_BLOCK_SIZE = 128 |
| 69 | MSA_INDEX_HEAD_DIM = 128 |
| 70 | # Wave width the top-k is written against: it gives one lane per output slot and |
| 71 | # reads the score row in wave-wide strips. |
| 72 | WAVE_SIZE = 64 |
| 73 | # Per-lane register slots the top-k can hold, which is what caps the context: a |
| 74 | # row may span at most SLOTS_MAX * WAVE_SIZE blocks. |
| 75 | SLOTS_MAX = 128 |
| 76 | MAX_SUPPORTED_BLOCKS = SLOTS_MAX * WAVE_SIZE |
| 77 | # MFMA columns the score pass has, one per (query token, index head) pair. |
| 78 | MFMA_COLS = 16 |
| 79 | # The fp8 MFMA the score kernels are built on exists on these targets only. |
| 80 | SUPPORTED_ARCHS = ("gfx950",) |
| 81 | |
| 82 | |
| 83 | def _pow2_ceil(n: int) -> int: |
| 84 | return 1 << (n - 1).bit_length() if n >= 1 else 1 |
| 85 | |
| 86 | |
| 87 | def score_block_width(max_seq_len: int, block_size: int) -> int: |
| 88 | """Block-axis width the top-k requires of the score buffer. |
| 89 | |
| 90 | Lanes read whole wave-wide strips with no tail guard and each holds a |
| 91 | power-of-two count of them, so the axis is padded past the block count. |
| 92 | """ |
| 93 | max_blk = math.ceil(max(max_seq_len, 1) / block_size) |
| 94 | return _pow2_ceil(math.ceil(max_blk / WAVE_SIZE)) * WAVE_SIZE |
| 95 | |
| 96 | |
| 97 | def aiter_indexer_max_decode_query_len(vllm_config: VllmConfig) -> int: |
| 98 | """Longest query a decode row can carry, which spec decode is what sets. |
| 99 | |
| 100 | Mirrors ``_init_reorder_batch_threshold(1, supports_spec_as_decode=True)``, |
| 101 | since that is what the builder splits the batch on and therefore what the |
| 102 | decode kernel will actually be handed. |
| 103 | """ |
| 104 | return max_decode_query_len(vllm_config) |
| 105 | |
| 106 | |
| 107 | @cache |
| 108 | def aiter_msa_kernels_unavailable_reason() -> str | None: |
| 109 | """Return why the AITER MSA score/top-k ops cannot be imported, or None. |
| 110 | |
| 111 | They are a recent addition, so an AITER that predates them imports fine |
| 112 | while these three names do not exist, and the failure would otherwise |
| 113 | surface as an ImportError from the middle of a forward. ``compile_ops`` |
| 114 | binds them lazily, so this costs the module import only -- the kernel build |
| 115 | itself still happens on the first call. |
| 116 | """ |
| 117 | try: |
| 118 | from aiter.ops.msa_attention import ( # noqa: F401 |
| 119 | pa_sparse_block_score_decode, |
| 120 | pa_sparse_block_score_prefill, |
| 121 | pa_sparse_block_topk, |
| 122 | ) |
| 123 | except ImportError as exc: |
| 124 | return f"AITER cannot supply the MSA score/top-k kernels ({exc})" |
| 125 | return None |
| 126 | |
| 127 | |
| 128 | def aiter_indexer_unsupported_reason( |
| 129 | *, |
| 130 | topk_blocks: int, |
| 131 | sparse_block_size: int, |
| 132 | num_index_heads: int, |
| 133 | index_head_dim: int, |
| 134 | indexer_kv_dtype: IndexerKVDType, |
| 135 | max_model_len: int, |
| 136 | max_decode_query_len: int = 1, |
| 137 | score_type: str = "max", |
| 138 | ) -> str | None: |
| 139 | """Return why this config cannot use the AITER indexer, or None if it can. |
| 140 | |
| 141 | Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache |
| 142 | dtype, the compiled score/top-k contract, MFMA column limits, max context |
| 143 | in blocks, and whether AITER exposes the MSA kernels. |
| 144 | ``select_aiter_indexer_impl_cls`` logs the string and falls back when it is |
| 145 | not None. |
| 146 | """ |
| 147 | if not current_platform.is_rocm(): |
| 148 | return ( |
| 149 | "needs ROCm for the AITER fp8 MFMA score/top-k, " |
| 150 | f"got platform={current_platform.device_type!r}" |
| 151 | ) |
| 152 | if not _minimax_m3_aiter_sparse_pa_requested(): |
| 153 | # The top-k emits the attend's page table in page-16 numbering, which |
| 154 | # only addresses the interleaved cache the AITER attend reads. Paired |
| 155 | # with any other attend there is nowhere valid to write it, so this |
| 156 | # indexer is not usable on its own. |
| 157 | return ( |
| 158 | "needs the AITER sparse PA attend, whose page table its top-k " |
| 159 | "emits (rocm_aiter_ops + shuffle KV cache layout)" |
| 160 | ) |
| 161 | if indexer_kv_dtype not in ("fp8", "fp8_e4m3"): |
| 162 | # The score kernels are fp8 MFMA; there is no bf16 instantiation. |
| 163 | return ( |
| 164 | f"needs an fp8 e4m3 index cache, got indexer_kv_dtype={indexer_kv_dtype!r}" |
| 165 | ) |
| 166 | from vllm.platforms.rocm import on_gfx950 |
| 167 | |
| 168 | if not on_gfx950(): |
| 169 | return f"needs {' or '.join(SUPPORTED_ARCHS)} for the fp8 MFMA" |
| 170 | if score_type != MSA_SCORE_TYPE: |
| 171 | return f"needs score_type={MSA_SCORE_TYPE!r}, got score_type={score_type!r}" |
| 172 | if topk_blocks != MSA_TOPK_BLOCKS: |
| 173 | return f"needs topk_blocks={MSA_TOPK_BLOCKS}, got topk_blocks={topk_blocks}" |
| 174 | if sparse_block_size != MSA_SPARSE_BLOCK_SIZE: |
| 175 | return ( |
| 176 | f"needs sparse_block_size={MSA_SPARSE_BLOCK_SIZE}, " |
| 177 | f"got sparse_block_size={sparse_block_size}" |
| 178 | ) |
| 179 | if index_head_dim != MSA_INDEX_HEAD_DIM: |
| 180 | return ( |
| 181 | f"needs index_head_dim={MSA_INDEX_HEAD_DIM}, " |
| 182 | f"got index_head_dim={index_head_dim}" |
| 183 | ) |
| 184 | if num_index_heads > MFMA_COLS: |
| 185 | return f"num_index_heads={num_index_heads} exceeds the {MFMA_COLS} MFMA columns" |
| 186 | # A decode row's whole query shares one MFMA tile, one column per (token, |
| 187 | # head) pair, so spec decode trades columns against the head count. |
| 188 | if num_index_heads * max_decode_query_len > MFMA_COLS: |
| 189 | return ( |
| 190 | f"num_index_heads={num_index_heads} x max_decode_query_len=" |
| 191 | f"{max_decode_query_len} exceeds the {MFMA_COLS} MFMA columns" |
| 192 | ) |
| 193 | max_blocks = math.ceil(max_model_len / sparse_block_size) |
| 194 | if max_blocks > MAX_SUPPORTED_BLOCKS: |
| 195 | return ( |
| 196 | f"max_model_len={max_model_len} needs {max_blocks} blocks per row, " |
| 197 | f"more than the top-k's {MAX_SUPPORTED_BLOCKS}" |
| 198 | ) |
| 199 | return aiter_msa_kernels_unavailable_reason() |
| 200 | |
| 201 | |
| 202 | class MiniMaxM3IndexerAiterBackend(MiniMaxM3IndexerBackend): |
| 203 | """Indexer side-cache backend selecting the AITER builder.""" |
| 204 | |
| 205 | @staticmethod |
| 206 | def get_builder_cls() -> type["MiniMaxM3IndexerAiterMetadataBuilder"]: |
| 207 | return MiniMaxM3IndexerAiterMetadataBuilder |
| 208 | |
| 209 | |
| 210 | @dataclass |
| 211 | class MiniMaxM3IndexerAiterMetadata(MiniMaxM3IndexerMetadata): |
| 212 | """Adds the per-row shape the ragged top-k needs. |
| 213 | |
| 214 | The uniform decode rows are recovered inside the kernel from ``seq_lens`` and |
| 215 | the shared query length, which also clamps cudagraph padding rows to nothing. |
| 216 | Prefill rows have no such shape, so it is materialized here once per forward |
| 217 | and shared by every layer -- which is also where the emitted page table's |
| 218 | tail block comes from, since a block count alone does not say how many tokens |
| 219 | the last block holds. |
| 220 | """ |
| 221 | |
| 222 | # [num_prefill_tokens] int32, cdiv(position + 1, sparse_block_size) per prefill row. |
| 223 | prefill_num_valid_pages: torch.Tensor | None = None |
| 224 | # [num_prefill_tokens] int32, the request each prefill row belongs to. |
| 225 | prefill_row_req_id: torch.Tensor | None = None |
| 226 | # [num_prefill_tokens] int32, causal token count (position + 1) per row. |
| 227 | prefill_kv_lens: torch.Tensor | None = None |
| 228 | |
| 229 | |
| 230 | class MiniMaxM3IndexerAiterMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): |
| 231 | """The Triton indexer's metadata plus the prefill rows' causal shape.""" |
| 232 | |
| 233 | def __init__( |
| 234 | self, |
| 235 | kv_cache_spec: AttentionSpec, |
| 236 | layer_names: list[str], |
| 237 | vllm_config: VllmConfig, |
| 238 | device: torch.device, |
| 239 | ) -> None: |
| 240 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 241 | hf_config = vllm_config.model_config.hf_config |
| 242 | text_config = getattr(hf_config, "text_config", hf_config) |
| 243 | self.sparse_block_size = int( |
| 244 | text_config.sparse_attention_config["sparse_block_size"] |
| 245 | ) |
| 246 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 247 | # Companions to the base's num_valid_pages_buffer, for the two vectors |
| 248 | # only the emitted table needs. |
| 249 | self.row_req_id_buffer = torch.empty( |
| 250 | max_tokens, dtype=torch.int32, device=device |
| 251 | ) |
| 252 | self.kv_lens_buffer = torch.empty(max_tokens, dtype=torch.int32, device=device) |
| 253 | |
| 254 | def build( |
| 255 | self, |
| 256 | common_prefix_len: int, |
| 257 | common_attn_metadata: CommonAttentionMetadata, |
| 258 | fast_build: bool = False, |
| 259 | ) -> MiniMaxM3IndexerAiterMetadata: |
| 260 | num_reqs = common_attn_metadata.num_reqs |
| 261 | num_tokens = common_attn_metadata.num_actual_tokens |
| 262 | query_start_loc = common_attn_metadata.query_start_loc |
| 263 | seq_lens = common_attn_metadata.seq_lens |
| 264 | block_table = common_attn_metadata.block_table_tensor |
| 265 | |
| 266 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 267 | split_decodes_and_prefills( |
| 268 | common_attn_metadata, |
| 269 | decode_threshold=self.reorder_batch_threshold, |
| 270 | require_uniform=True, |
| 271 | ) |
| 272 | ) |
| 273 | assert num_decodes + num_prefills == num_reqs |
| 274 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 275 | |
| 276 | # Decode-first batch: context lengths into the stable cudagraph buffer. |
| 277 | context_lens = self.context_len_buffer[:num_reqs] |
| 278 | context_lens.copy_( |
| 279 | common_attn_metadata.compute_num_computed_tokens(), non_blocking=True |
| 280 | ) |
| 281 | |
| 282 | prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None |
| 283 | prefill_num_valid_pages: torch.Tensor | None = None |
| 284 | prefill_row_req_id: torch.Tensor | None = None |
| 285 | prefill_kv_lens: torch.Tensor | None = None |
| 286 | if num_prefills > 0: |
| 287 | cu_seqlens_q = (query_start_loc[num_decodes:] - num_decode_tokens).to( |
| 288 | torch.int32 |
| 289 | ) |
| 290 | prefill_metadata = MiniMaxM3IndexerPrefillMetadata( |
| 291 | cu_seqlens_q=cu_seqlens_q, |
| 292 | seq_lens=seq_lens[num_decodes:], |
| 293 | context_lens=context_lens[num_decodes:], |
| 294 | block_table=block_table[num_decodes:], |
| 295 | max_query_len=common_attn_metadata.max_query_len, |
| 296 | max_seq_len=common_attn_metadata.max_seq_len, |
| 297 | ) |
| 298 | # A prefill row sees its own position, so its causal length and block |
| 299 | # count both follow from that alone; the request it belongs to comes |
| 300 | # from the query offsets. Prefill batches are never captured, so the |
| 301 | # stable buffers are only being reused here, not required. |
| 302 | positions = common_attn_metadata.positions |
| 303 | assert positions is not None |
| 304 | row_positions = positions[num_decode_tokens:num_tokens] |
| 305 | prefill_num_valid_pages = self.num_valid_pages_buffer[ |
| 306 | num_decode_tokens:num_tokens |
| 307 | ] |
| 308 | prefill_num_valid_pages.copy_( |
| 309 | row_positions // self.sparse_block_size + 1, non_blocking=True |
| 310 | ) |
| 311 | prefill_kv_lens = self.kv_lens_buffer[num_decode_tokens:num_tokens] |
| 312 | prefill_kv_lens.copy_(row_positions + 1, non_blocking=True) |
| 313 | prefill_row_req_id = self.row_req_id_buffer[num_decode_tokens:num_tokens] |
| 314 | prefill_row_req_id.copy_( |
| 315 | torch.searchsorted( |
| 316 | cu_seqlens_q[1:].contiguous(), |
| 317 | torch.arange( |
| 318 | num_prefill_tokens, |
| 319 | dtype=torch.int32, |
| 320 | device=cu_seqlens_q.device, |
| 321 | ), |
| 322 | right=True, |
| 323 | ), |
| 324 | non_blocking=True, |
| 325 | ) |
| 326 | |
| 327 | decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None |
| 328 | if num_decodes > 0: |
| 329 | qsl_cpu = common_attn_metadata.query_start_loc_cpu |
| 330 | query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] |
| 331 | decode_query_len = int(query_lens_cpu[0].item()) |
| 332 | assert decode_query_len > 0 |
| 333 | assert torch.all( |
| 334 | (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) |
| 335 | ) |
| 336 | assert num_decode_tokens == num_decodes * decode_query_len |
| 337 | decode_metadata = MiniMaxM3IndexerDecodeMetadata( |
| 338 | seq_lens=seq_lens[:num_decodes], |
| 339 | block_table=block_table[:num_decodes], |
| 340 | max_seq_len=common_attn_metadata.max_seq_len, |
| 341 | decode_query_len=decode_query_len, |
| 342 | max_decode_query_len=self.max_decode_query_len, |
| 343 | ) |
| 344 | |
| 345 | return MiniMaxM3IndexerAiterMetadata( |
| 346 | seq_lens=seq_lens, |
| 347 | max_seq_len=common_attn_metadata.max_seq_len, |
| 348 | slot_mapping=common_attn_metadata.slot_mapping, |
| 349 | num_actual_tokens=num_tokens, |
| 350 | num_decodes=num_decodes, |
| 351 | num_decode_tokens=num_decode_tokens, |
| 352 | num_prefills=num_prefills, |
| 353 | num_prefill_tokens=num_prefill_tokens, |
| 354 | prefill=prefill_metadata, |
| 355 | decode=decode_metadata, |
| 356 | prefill_num_valid_pages=prefill_num_valid_pages, |
| 357 | prefill_row_req_id=prefill_row_req_id, |
| 358 | prefill_kv_lens=prefill_kv_lens, |
| 359 | ) |
| 360 | |
| 361 | |
| 362 | class MiniMaxM3IndexerAiterImpl(MiniMaxM3IndexerImpl): |
| 363 | """AITER fp8 score + top-k for both prefill and decode.""" |
| 364 | |
| 365 | indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerAiterBackend |
| 366 | |
| 367 | def __init__(self, **kwargs) -> None: |
| 368 | super().__init__(**kwargs) |
| 369 | # Both passes are v_mfma_f32_16x16x32_fp8_fp8 with no bf16 instantiation, |
| 370 | # so an index cache of any other dtype would be read as e4m3 bytes. |
| 371 | # The selector will not get here, but nothing else may either. |
| 372 | if self.indexer_kv_dtype not in ("fp8", "fp8_e4m3"): |
| 373 | raise ValueError( |
| 374 | "The AITER indexer requires an fp8 e4m3 index cache, got " |
| 375 | f"indexer_kv_dtype={self.indexer_kv_dtype!r}" |
| 376 | ) |
| 377 | # Shared, stable-address page table + per-row context bound the top-k |
| 378 | # emits for the attend. Owned by the model so one allocation serves |
| 379 | # every layer; left None when it reserved none, in which case the |
| 380 | # attend rebuilds the table itself. |
| 381 | self.sparse_bt_buffer: torch.Tensor | None = None |
| 382 | self.sparse_ctx_buffer: torch.Tensor | None = None |
| 383 | |
| 384 | @property |
| 385 | def pages_per_block(self) -> int: |
| 386 | """Physical pages one selected block expands into for the attend.""" |
| 387 | return self.block_size // ASM_PAGE_SIZE |
| 388 | |
| 389 | def _table_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 390 | """The page table and context rows covering score rows ``[lo, hi)``. |
| 391 | |
| 392 | One table row per (token, kv head), head minor, which is the order |
| 393 | ``pa_decode_gluon`` reads once it flattens the cache. Slices of the |
| 394 | shared buffers rather than copies, so the attend sees the writes; a |
| 395 | model that reserved no buffers gets throwaway ones, and the attend |
| 396 | rebuilds the table itself in that case. |
| 397 | """ |
| 398 | kv_heads = self.num_kv_heads |
| 399 | if self.sparse_bt_buffer is None or self.sparse_ctx_buffer is None: |
| 400 | rows = (hi - lo) * kv_heads |
| 401 | device = self.index_cache.kv_cache.device |
| 402 | return ( |
| 403 | torch.empty( |
| 404 | (rows, self.topk_blocks * self.pages_per_block), |
| 405 | dtype=torch.int32, |
| 406 | device=device, |
| 407 | ), |
| 408 | torch.empty(rows, dtype=torch.int32, device=device), |
| 409 | ) |
| 410 | return ( |
| 411 | self.sparse_bt_buffer[lo * kv_heads : hi * kv_heads], |
| 412 | self.sparse_ctx_buffer[lo * kv_heads : hi * kv_heads], |
| 413 | ) |
| 414 | |
| 415 | def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor: |
| 416 | """Score buffer for ``rows`` query rows. |
| 417 | |
| 418 | Left uninitialized on purpose: the score pass writes every block up to |
| 419 | the longest row it covers, and the top-k reads only the blocks its own |
| 420 | row can see, so nothing downstream observes the padded tail. Filling it |
| 421 | would cost a write over the whole [H, rows, width] extent, which at long |
| 422 | context is the largest tensor in the indexer. |
| 423 | """ |
| 424 | return torch.empty( |
| 425 | ( |
| 426 | self.num_index_heads, |
| 427 | rows, |
| 428 | score_block_width(max_seq_len, self.block_size), |
| 429 | ), |
| 430 | dtype=torch.float32, |
| 431 | device=self.index_cache.kv_cache.device, |
| 432 | ) |
| 433 | |
| 434 | def forward( |
| 435 | self, |
| 436 | index_query: torch.Tensor, |
| 437 | *, |
| 438 | decode_page16_block_table: torch.Tensor | None = None, |
| 439 | prefill_page16_block_table: torch.Tensor | None = None, |
| 440 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 441 | from aiter.ops.msa_attention import ( |
| 442 | pa_sparse_block_score_decode, |
| 443 | pa_sparse_block_score_prefill, |
| 444 | pa_sparse_block_topk, |
| 445 | ) |
| 446 | |
| 447 | attn_metadata = get_forward_context().attn_metadata |
| 448 | if not isinstance(attn_metadata, dict): |
| 449 | return None, None # profiling run; caches unbound |
| 450 | md = attn_metadata[self.index_cache.prefix] |
| 451 | assert isinstance(md, MiniMaxM3IndexerAiterMetadata) |
| 452 | # The emitted page table addresses the main cache, whose blocks are a |
| 453 | # different group's than the index cache's, so the top-k resolves the |
| 454 | # selection through the attend's block table -- only the score pass reads |
| 455 | # the index cache and takes the indexer's. It has to be the page-16 |
| 456 | # rebase of that table, which the attend's own metadata builder does |
| 457 | # once per step; the indexer's metadata cannot reach another group's |
| 458 | # blocks, so the layer reads it off the attend and hands it in. |
| 459 | num_tokens = md.num_actual_tokens |
| 460 | nd = md.num_decode_tokens |
| 461 | iq = index_query[:num_tokens].view( |
| 462 | -1, self.num_index_heads, self.index_head_dim |
| 463 | ) |
| 464 | kv = self.index_cache.kv_cache |
| 465 | |
| 466 | # Both sides write into the single shared persistent buffer (decode at |
| 467 | # [:, :nd], prefill at [:, nd:]) and return views into it. The top-k |
| 468 | # takes the head/row strides, so these slices need no copy back. |
| 469 | buf = self.topk_indices_buffer |
| 470 | if buf is None: |
| 471 | buf = torch.empty( |
| 472 | (self.num_index_heads, num_tokens, self.topk_blocks), |
| 473 | dtype=torch.int32, |
| 474 | device=iq.device, |
| 475 | ) |
| 476 | |
| 477 | decode_topk: torch.Tensor | None = None |
| 478 | prefill_topk: torch.Tensor | None = None |
| 479 | |
| 480 | if md.num_decodes > 0: |
| 481 | d = md.decode |
| 482 | assert d is not None |
| 483 | assert decode_page16_block_table is not None, ( |
| 484 | "the AITER indexer's top-k emits the attend's page table and " |
| 485 | "needs the page-16 rebase of the attend's decode block table" |
| 486 | ) |
| 487 | score = self._new_score(nd, d.max_seq_len) |
| 488 | pa_sparse_block_score_decode( |
| 489 | iq[:nd], |
| 490 | kv, |
| 491 | score, |
| 492 | d.block_table, |
| 493 | d.seq_lens, |
| 494 | init_blocks=self.init_blocks, |
| 495 | local_blocks=self.local_blocks, |
| 496 | query_len=d.decode_query_len, |
| 497 | max_seq_len=d.max_seq_len, |
| 498 | ) |
| 499 | decode_topk = buf[:, :nd, :] |
| 500 | sparse_bt, sparse_ctx = self._table_rows(0, nd) |
| 501 | # A decode row's causal length is seq_len - query_len + token + 1, |
| 502 | # which the kernel derives itself, so the emitted table covers |
| 503 | # speculative rows without any extra per-row shape. |
| 504 | pa_sparse_block_topk( |
| 505 | score, |
| 506 | decode_topk, |
| 507 | decode_page16_block_table, |
| 508 | d.seq_lens, |
| 509 | sparse_bt, |
| 510 | sparse_ctx, |
| 511 | max_seq_len=d.max_seq_len, |
| 512 | block_size=self.block_size, |
| 513 | query_len=d.decode_query_len, |
| 514 | num_kv_heads=self.num_kv_heads, |
| 515 | pages_per_block=self.pages_per_block, |
| 516 | ) |
| 517 | |
| 518 | if md.num_prefills > 0: |
| 519 | p = md.prefill |
| 520 | assert p is not None |
| 521 | assert prefill_page16_block_table is not None, ( |
| 522 | "the AITER indexer's top-k emits the attend's page table and " |
| 523 | "needs the page-16 rebase of the attend's prefill block table" |
| 524 | ) |
| 525 | assert md.prefill_num_valid_pages is not None |
| 526 | assert md.prefill_row_req_id is not None |
| 527 | assert md.prefill_kv_lens is not None |
| 528 | score = self._new_score(num_tokens - nd, p.max_seq_len) |
| 529 | pa_sparse_block_score_prefill( |
| 530 | iq[nd:], |
| 531 | kv, |
| 532 | score, |
| 533 | p.block_table, |
| 534 | p.cu_seqlens_q, |
| 535 | p.seq_lens, |
| 536 | init_blocks=self.init_blocks, |
| 537 | local_blocks=self.local_blocks, |
| 538 | max_query_len=p.max_query_len, |
| 539 | max_seq_len=p.max_seq_len, |
| 540 | ) |
| 541 | prefill_topk = buf[:, nd:num_tokens, :] |
| 542 | sparse_bt, sparse_ctx = self._table_rows(nd, num_tokens) |
| 543 | # Ragged rows carry their request and causal length explicitly: the |
| 544 | # block count alone cannot place the tail block the table ends on. |
| 545 | pa_sparse_block_topk( |
| 546 | score, |
| 547 | prefill_topk, |
| 548 | prefill_page16_block_table, |
| 549 | p.seq_lens, |
| 550 | sparse_bt, |
| 551 | sparse_ctx, |
| 552 | max_seq_len=p.max_seq_len, |
| 553 | block_size=self.block_size, |
| 554 | num_valid_pages=md.prefill_num_valid_pages, |
| 555 | row_req_id=md.prefill_row_req_id, |
| 556 | kv_lens=md.prefill_kv_lens, |
| 557 | num_kv_heads=self.num_kv_heads, |
| 558 | pages_per_block=self.pages_per_block, |
| 559 | ) |
| 560 | |
| 561 | return decode_topk, prefill_topk |
| 562 | |
| 563 | |
| 564 | def select_aiter_indexer_impl_cls( |
| 565 | *, |
| 566 | topk_blocks: int, |
| 567 | sparse_block_size: int, |
| 568 | num_index_heads: int, |
| 569 | index_head_dim: int, |
| 570 | indexer_kv_dtype: IndexerKVDType, |
| 571 | score_type: str = "max", |
| 572 | ) -> type[MiniMaxM3IndexerAiterImpl] | None: |
| 573 | """The AITER indexer impl if this config can use it, else None. |
| 574 | |
| 575 | ``None`` sends the caller to the platform-neutral ``MiniMaxM3Indexer``, |
| 576 | which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8 |
| 577 | index cache that lands here has nowhere to go. |
| 578 | """ |
| 579 | reason = aiter_indexer_unsupported_reason( |
| 580 | topk_blocks=topk_blocks, |
| 581 | sparse_block_size=sparse_block_size, |
| 582 | num_index_heads=num_index_heads, |
| 583 | index_head_dim=index_head_dim, |
| 584 | indexer_kv_dtype=indexer_kv_dtype, |
| 585 | max_model_len=get_current_vllm_config().model_config.max_model_len, |
| 586 | max_decode_query_len=aiter_indexer_max_decode_query_len( |
| 587 | get_current_vllm_config() |
| 588 | ), |
| 589 | score_type=score_type, |
| 590 | ) |
| 591 | if reason is not None: |
| 592 | logger.info_once("MiniMax M3 indexer: AITER unavailable (%s)", reason) |
| 593 | return None |
| 594 | logger.info_once( |
| 595 | "MiniMax M3 indexer: selected AITER (fp8 MFMA score + top-k) " |
| 596 | "[topk_blocks=%d, indexer_kv_dtype=%s]", |
| 597 | topk_blocks, |
| 598 | indexer_kv_dtype, |
| 599 | ) |
| 600 | return MiniMaxM3IndexerAiterImpl |
| 601 | |
| 602 | |
| 603 | class MiniMaxM3AiterIndexer(nn.Module): |
| 604 | """``MiniMaxM3Indexer``'s surface over the AITER impl. |
| 605 | |
| 606 | The platform-neutral wrapper picks its impl through ``common``'s selector |
| 607 | and forwards a Triton-only set of fused-table arguments, neither of which |
| 608 | can reach this impl without editing ``common``. This holds the same three |
| 609 | members the attention layer uses -- ``impl``, ``index_cache``, |
| 610 | ``num_index_heads`` -- and forwards the one argument AITER needs instead. |
| 611 | """ |
| 612 | |
| 613 | def __init__( |
| 614 | self, |
| 615 | *, |
| 616 | impl_cls: type[MiniMaxM3IndexerAiterImpl], |
| 617 | sparse_bt_buffer: torch.Tensor | None = None, |
| 618 | sparse_ctx_buffer: torch.Tensor | None = None, |
| 619 | **impl_kwargs, |
| 620 | ) -> None: |
| 621 | super().__init__() |
| 622 | self.impl = impl_cls(**impl_kwargs) |
| 623 | # Assigned rather than passed: the impl's base ``__init__`` lives in |
| 624 | # common and takes no table buffers. |
| 625 | self.impl.sparse_bt_buffer = sparse_bt_buffer |
| 626 | self.impl.sparse_ctx_buffer = sparse_ctx_buffer |
| 627 | |
| 628 | @property |
| 629 | def index_cache(self) -> MiniMaxM3IndexerCache: |
| 630 | return self.impl.index_cache |
| 631 | |
| 632 | @property |
| 633 | def num_index_heads(self) -> int: |
| 634 | return self.impl.num_index_heads |
| 635 | |
| 636 | def forward( |
| 637 | self, |
| 638 | index_query: torch.Tensor, |
| 639 | *, |
| 640 | decode_page16_block_table: torch.Tensor | None = None, |
| 641 | prefill_page16_block_table: torch.Tensor | None = None, |
| 642 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 643 | return self.impl( |
| 644 | index_query, |
| 645 | decode_page16_block_table=decode_page16_block_table, |
| 646 | prefill_page16_block_table=prefill_page16_block_table, |
| 647 | ) |