| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Main block-sparse GQA attention for MiniMax M3 sparse layers. |
| 4 | |
| 5 | The lightning indexer (``indexer.py``) selects the top-k KV blocks (written into |
| 6 | the shared ``layer.topk_indices_buffer``); this module holds the main attention |
| 7 | that attends only to those blocks: the paged K/V cache backend, its metadata + |
| 8 | builder, and the impl that reads the indexer's top-k from that buffer. The Triton |
| 9 | attend kernel lives here; the SM100 (MSA) |
| 10 | ``build_k2q_csr`` + ``sparse_atten_func`` attend lives in |
| 11 | ``nvidia/sparse_attention_msa.py``. |
| 12 | |
| 13 | ``MiniMaxM3SparseBackend`` and ``MiniMaxM3SparseMetadata`` are referenced by the |
| 14 | attention-backend registry (by dotted path) and by spec-decode, so they must |
| 15 | keep these names and stay in this module. |
| 16 | """ |
| 17 | |
| 18 | from dataclasses import dataclass |
| 19 | from typing import ClassVar |
| 20 | |
| 21 | import torch |
| 22 | |
| 23 | from vllm._aiter_ops import rocm_aiter_ops |
| 24 | from vllm.config import VllmConfig |
| 25 | from vllm.config.cache import CacheDType |
| 26 | from vllm.forward_context import get_forward_context |
| 27 | from vllm.logger import init_logger |
| 28 | from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE |
| 29 | from vllm.platforms import current_platform |
| 30 | |
| 31 | # AMD/ROCm uses the gfx942/gfx950-optimized block-sparse kernels in amd.ops; |
| 32 | # every other platform uses the generic common.ops implementation. |
| 33 | if current_platform.is_rocm(): |
| 34 | from vllm.models.minimax_m3.amd.ops.sparse_attn import ( |
| 35 | minimax_m3_sparse_attn, |
| 36 | minimax_m3_sparse_attn_decode, |
| 37 | ) |
| 38 | else: |
| 39 | from vllm.models.minimax_m3.common.ops.sparse_attn import ( |
| 40 | minimax_m3_sparse_attn, |
| 41 | minimax_m3_sparse_attn_decode, |
| 42 | ) |
| 43 | from vllm.v1.attention.backend import ( |
| 44 | AttentionBackend, |
| 45 | AttentionCGSupport, |
| 46 | AttentionImplBase, |
| 47 | AttentionLayer, |
| 48 | AttentionMetadata, |
| 49 | AttentionMetadataBuilder, |
| 50 | CommonAttentionMetadata, |
| 51 | MultipleOf, |
| 52 | ) |
| 53 | from vllm.v1.attention.backends.utils import ( |
| 54 | split_decodes_and_prefills, |
| 55 | ) |
| 56 | from vllm.v1.kv_cache_interface import ( |
| 57 | AttentionSpec, |
| 58 | KVCacheLayout, |
| 59 | get_kv_quant_mode, |
| 60 | is_quantized_kv_cache, |
| 61 | ) |
| 62 | |
| 63 | logger = init_logger(__name__) |
| 64 | |
| 65 | |
| 66 | def _minimax_m3_aiter_sparse_pa_requested() -> bool: |
| 67 | return rocm_aiter_ops.is_enabled() and rocm_aiter_ops.is_shuffle_kv_cache_enabled() |
| 68 | |
| 69 | |
| 70 | def minimax_m3_rebase_slots_to_page16( |
| 71 | slot_mapping: torch.Tensor, |
| 72 | block_size: int, |
| 73 | out: torch.Tensor | None = None, |
| 74 | ) -> torch.Tensor: |
| 75 | """Rebase a token slot mapping onto AITER's page-16 page numbering. |
| 76 | |
| 77 | AITER's KV writer derives the destination page from ``slot // 16`` and |
| 78 | assumes pages are numbered consecutively. When the resolved layout stores |
| 79 | both K/V sides inside a block, a block spans twice as many pages, so only |
| 80 | the block component of the slot doubles -- the offset within the page must |
| 81 | not move. Clamping before the division leaves padding slots at their |
| 82 | negative sentinel, which is what tells the writer to skip them. |
| 83 | """ |
| 84 | if out is None: |
| 85 | out = torch.empty_like(slot_mapping) |
| 86 | torch.clamp(slot_mapping, min=0, out=out) |
| 87 | return ( |
| 88 | out.div_(block_size, rounding_mode="floor").mul_(block_size).add_(slot_mapping) |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def minimax_m3_query_token_positions( |
| 93 | cu_seqlens_q: torch.Tensor, |
| 94 | prefix_lens: torch.Tensor, |
| 95 | total_q: int, |
| 96 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 97 | """Map each prefill query token to its request id and absolute KV position. |
| 98 | |
| 99 | Both tensors are the same for every sparse layer in a step, so the AITER |
| 100 | sparse PA block-table builder reads them from the metadata instead of |
| 101 | rebuilding them once per layer. |
| 102 | """ |
| 103 | pos = torch.arange(total_q, dtype=torch.int32, device=cu_seqlens_q.device) |
| 104 | query_req_id = torch.searchsorted( |
| 105 | cu_seqlens_q[1:].contiguous(), pos, right=True |
| 106 | ).to(torch.int32) |
| 107 | query_abs_pos = (prefix_lens[query_req_id] + (pos - cu_seqlens_q[query_req_id])).to( |
| 108 | torch.int32 |
| 109 | ) |
| 110 | return query_req_id, query_abs_pos |
| 111 | |
| 112 | |
| 113 | def minimax_m3_use_aiter_sparse_pa( |
| 114 | num_kv_heads: int, *, emits_sparse_block_table: bool = False |
| 115 | ) -> bool: |
| 116 | """Whether to use the ROCm AITER page-16 sparse PA prototype. |
| 117 | |
| 118 | More than one KV head per rank is only served when the layer's indexer |
| 119 | emits the attend's page table itself, since the Triton builders this path |
| 120 | otherwise falls back to address one head's cache. Whether an indexer does |
| 121 | that is a ROCm-side fact, so its caller passes it in. |
| 122 | """ |
| 123 | requested = _minimax_m3_aiter_sparse_pa_requested() |
| 124 | if requested and num_kv_heads != 1 and not emits_sparse_block_table: |
| 125 | raise ValueError( |
| 126 | "MiniMax M3 AITER sparse paged attention requires " |
| 127 | f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}, " |
| 128 | "unless the indexer emits the attend's page table." |
| 129 | ) |
| 130 | return requested |
| 131 | |
| 132 | |
| 133 | class MiniMaxM3SparseBackend(AttentionBackend): |
| 134 | """Block-sparse GQA backend for MiniMax M3 sparse attention layers.""" |
| 135 | |
| 136 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16] |
| 137 | # bf16 or fp8 (e4m3/e5m2): the Triton kernels dequant fp8 before the dots. |
| 138 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 139 | "bfloat16", |
| 140 | "fp8", |
| 141 | "fp8_e4m3", |
| 142 | "fp8_e5m2", |
| 143 | ] |
| 144 | |
| 145 | @staticmethod |
| 146 | def get_name() -> str: |
| 147 | return "MINIMAX_M3_SPARSE" |
| 148 | |
| 149 | @staticmethod |
| 150 | def get_impl_cls() -> type["MiniMaxM3SparseImpl"]: |
| 151 | # Concrete impl chosen by select_main_impl_cls; base for introspection. |
| 152 | return MiniMaxM3SparseImpl |
| 153 | |
| 154 | @staticmethod |
| 155 | def get_builder_cls() -> type["MiniMaxM3SparseMetadataBuilder"]: |
| 156 | return MiniMaxM3SparseMetadataBuilder |
| 157 | |
| 158 | @classmethod |
| 159 | def get_supported_head_sizes(cls) -> list[int]: |
| 160 | return [128] |
| 161 | |
| 162 | @staticmethod |
| 163 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 164 | # Page size == sparse block size (one sparse block per KV page). |
| 165 | return [128] |
| 166 | |
| 167 | @classmethod |
| 168 | def is_sparse(cls) -> bool: |
| 169 | return True |
| 170 | |
| 171 | @classmethod |
| 172 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: |
| 173 | # AITER sparse PA reads K and V as separate page-16 tensors. LHBNC |
| 174 | # gives each side a dense plane; LBHNC interleaves the sides inside a |
| 175 | # block, which the page-id rebasing in ops/sparse_pa.py absorbs. LBHNC |
| 176 | # is listed first because these layers publish two head slots while the |
| 177 | # indexer publishes one, and mixed HNC shapes need a block-compact |
| 178 | # layout to share the block-id pool. |
| 179 | if _minimax_m3_aiter_sparse_pa_requested(): |
| 180 | return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) |
| 181 | return super().supported_kv_cache_layouts() |
| 182 | |
| 183 | |
| 184 | @dataclass |
| 185 | class MiniMaxM3SparsePrefillMetadata: |
| 186 | """Per-prefill state; ``cu_seqlens_k``/``total_kv_blocks`` feed the MSA CSR.""" |
| 187 | |
| 188 | cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0 |
| 189 | cu_seqlens_k: torch.Tensor # [num_prefills + 1] int32, cumulative KV lengths |
| 190 | seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths |
| 191 | context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens) |
| 192 | block_table: torch.Tensor |
| 193 | max_query_len: int |
| 194 | max_seq_len: int |
| 195 | total_kv_blocks: int |
| 196 | |
| 197 | # Per-query-token request id and absolute KV position, built once per step |
| 198 | # for the AITER sparse PA block-table kernel. See |
| 199 | # ``minimax_m3_query_token_positions``. |
| 200 | query_req_id: torch.Tensor | None = None |
| 201 | query_abs_pos: torch.Tensor | None = None |
| 202 | |
| 203 | |
| 204 | @dataclass |
| 205 | class MiniMaxM3SparseDecodeMetadata: |
| 206 | """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform |
| 207 | per-request query length (1, or 1 + num_speculative_tokens).""" |
| 208 | |
| 209 | seq_lens: torch.Tensor # [num_decodes] int32 |
| 210 | block_table: torch.Tensor |
| 211 | decode_query_len: int |
| 212 | |
| 213 | |
| 214 | @dataclass |
| 215 | class MiniMaxM3SparseMetadata(AttentionMetadata): |
| 216 | """Sparse-attention metadata, split into prefill and decode sub-metadata.""" |
| 217 | |
| 218 | seq_lens: torch.Tensor |
| 219 | max_seq_len: int |
| 220 | slot_mapping: torch.Tensor |
| 221 | |
| 222 | num_actual_tokens: int # total query tokens (decode-first batch) |
| 223 | |
| 224 | # Split counts (batch reordered decode-first). |
| 225 | num_decodes: int |
| 226 | num_decode_tokens: int |
| 227 | num_prefills: int |
| 228 | num_prefill_tokens: int |
| 229 | |
| 230 | prefill: MiniMaxM3SparsePrefillMetadata | None = None |
| 231 | decode: MiniMaxM3SparseDecodeMetadata | None = None |
| 232 | |
| 233 | # ``slot_mapping`` rebased onto the page-16 numbering AITER's KV writer |
| 234 | # uses when both K/V sides share a block. Only set for the AITER sparse PA |
| 235 | # path; see ``MiniMaxM3SparseMetadataBuilder.build``. |
| 236 | page16_slot_mapping: torch.Tensor | None = None |
| 237 | |
| 238 | |
| 239 | class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMetadata]): |
| 240 | # Full cudagraphs for uniform decode batches, incl. spec-decode verify |
| 241 | # batches with >1 query token/request. |
| 242 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 243 | # Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when |
| 244 | # spec decode is on; must match the indexer builder so the splits agree. |
| 245 | reorder_batch_threshold: int = 1 |
| 246 | |
| 247 | def __init__( |
| 248 | self, |
| 249 | kv_cache_spec: AttentionSpec, |
| 250 | layer_names: list[str], |
| 251 | vllm_config: VllmConfig, |
| 252 | device: torch.device, |
| 253 | ) -> None: |
| 254 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 255 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 256 | # Stable context-length buffer for decode cudagraph replays. |
| 257 | self.context_len_buffer = torch.empty( |
| 258 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 259 | dtype=torch.int32, |
| 260 | device=device, |
| 261 | ) |
| 262 | # Every sparse layer shares one slot mapping, so the AITER page-16 |
| 263 | # rebase is done here once per step instead of once per layer. Stable |
| 264 | # buffer for the same reason as the context lengths above. |
| 265 | # The request itself, not the gate: the buffer is the same whatever the |
| 266 | # head count, and a builder cannot see which indexer its layers picked. |
| 267 | self.use_aiter_sparse_pa = _minimax_m3_aiter_sparse_pa_requested() |
| 268 | self.page16_slot_mapping_buffer: torch.Tensor | None = None |
| 269 | if self.use_aiter_sparse_pa: |
| 270 | self.page16_slot_mapping_buffer = torch.empty( |
| 271 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 272 | dtype=torch.int64, |
| 273 | device=device, |
| 274 | ) |
| 275 | |
| 276 | def build( |
| 277 | self, |
| 278 | common_prefix_len: int, |
| 279 | common_attn_metadata: CommonAttentionMetadata, |
| 280 | fast_build: bool = False, |
| 281 | ) -> MiniMaxM3SparseMetadata: |
| 282 | num_reqs = common_attn_metadata.num_reqs |
| 283 | num_tokens = common_attn_metadata.num_actual_tokens |
| 284 | query_start_loc = common_attn_metadata.query_start_loc |
| 285 | seq_lens = common_attn_metadata.seq_lens |
| 286 | block_table = common_attn_metadata.block_table_tensor |
| 287 | |
| 288 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 289 | split_decodes_and_prefills( |
| 290 | common_attn_metadata, |
| 291 | decode_threshold=self.reorder_batch_threshold, |
| 292 | require_uniform=True, |
| 293 | ) |
| 294 | ) |
| 295 | assert num_decodes + num_prefills == num_reqs |
| 296 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 297 | |
| 298 | # Decode-first batch: context lengths into the stable cudagraph buffer. |
| 299 | context_lens = self.context_len_buffer[:num_reqs] |
| 300 | context_lens.copy_( |
| 301 | common_attn_metadata.compute_num_computed_tokens(), non_blocking=True |
| 302 | ) |
| 303 | |
| 304 | prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None |
| 305 | if num_prefills > 0: |
| 306 | seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound |
| 307 | assert seq_lens_cpu is not None |
| 308 | prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] |
| 309 | prefill_total_kv_blocks = ( |
| 310 | ((prefill_seq_lens_cpu + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE) |
| 311 | .sum() |
| 312 | .item() |
| 313 | ) |
| 314 | prefill_kv_lens = seq_lens[num_decodes:] |
| 315 | prefill_cu_seqlens_k = torch.empty( |
| 316 | num_prefills + 1, dtype=torch.int32, device=seq_lens.device |
| 317 | ) |
| 318 | if current_platform.is_rocm(): |
| 319 | prefill_cu_seqlens_k[:1].zero_() |
| 320 | else: |
| 321 | prefill_cu_seqlens_k[0] = 0 |
| 322 | torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:]) |
| 323 | prefill_cu_seqlens_q = ( |
| 324 | query_start_loc[num_decodes:] - num_decode_tokens |
| 325 | ).to(torch.int32) |
| 326 | prefill_context_lens = context_lens[num_decodes:] |
| 327 | query_req_id = query_abs_pos = None |
| 328 | if self.use_aiter_sparse_pa: |
| 329 | query_req_id, query_abs_pos = minimax_m3_query_token_positions( |
| 330 | prefill_cu_seqlens_q, prefill_context_lens, num_prefill_tokens |
| 331 | ) |
| 332 | prefill_metadata = MiniMaxM3SparsePrefillMetadata( |
| 333 | cu_seqlens_q=prefill_cu_seqlens_q, |
| 334 | cu_seqlens_k=prefill_cu_seqlens_k, |
| 335 | seq_lens=prefill_kv_lens, |
| 336 | context_lens=prefill_context_lens, |
| 337 | block_table=block_table[num_decodes:], |
| 338 | max_query_len=common_attn_metadata.max_query_len, |
| 339 | max_seq_len=common_attn_metadata.max_seq_len, |
| 340 | total_kv_blocks=prefill_total_kv_blocks, |
| 341 | query_req_id=query_req_id, |
| 342 | query_abs_pos=query_abs_pos, |
| 343 | ) |
| 344 | |
| 345 | decode_metadata: MiniMaxM3SparseDecodeMetadata | None = None |
| 346 | if num_decodes > 0: |
| 347 | qsl_cpu = common_attn_metadata.query_start_loc_cpu |
| 348 | query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] |
| 349 | decode_query_len = int(query_lens_cpu[0].item()) |
| 350 | assert decode_query_len > 0 |
| 351 | assert torch.all( |
| 352 | (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) |
| 353 | ) |
| 354 | assert num_decode_tokens == num_decodes * decode_query_len |
| 355 | decode_metadata = MiniMaxM3SparseDecodeMetadata( |
| 356 | seq_lens=seq_lens[:num_decodes], |
| 357 | block_table=block_table[:num_decodes], |
| 358 | decode_query_len=decode_query_len, |
| 359 | ) |
| 360 | |
| 361 | page16_slot_mapping = None |
| 362 | if self.page16_slot_mapping_buffer is not None: |
| 363 | slot_mapping = common_attn_metadata.slot_mapping |
| 364 | page16_slot_mapping = minimax_m3_rebase_slots_to_page16( |
| 365 | slot_mapping, |
| 366 | self.kv_cache_spec.block_size, |
| 367 | out=self.page16_slot_mapping_buffer[: slot_mapping.shape[0]], |
| 368 | ) |
| 369 | |
| 370 | return MiniMaxM3SparseMetadata( |
| 371 | seq_lens=seq_lens, |
| 372 | max_seq_len=common_attn_metadata.max_seq_len, |
| 373 | slot_mapping=common_attn_metadata.slot_mapping, |
| 374 | page16_slot_mapping=page16_slot_mapping, |
| 375 | num_actual_tokens=num_tokens, |
| 376 | num_decodes=num_decodes, |
| 377 | num_decode_tokens=num_decode_tokens, |
| 378 | num_prefills=num_prefills, |
| 379 | num_prefill_tokens=num_prefill_tokens, |
| 380 | prefill=prefill_metadata, |
| 381 | decode=decode_metadata, |
| 382 | ) |
| 383 | |
| 384 | |
| 385 | class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]): |
| 386 | """Abstract base for block-sparse GQA over the indexer-selected blocks. |
| 387 | |
| 388 | Inherits ``AttentionImplBase`` for a custom forward signature (the layer |
| 389 | pre-inserts K/V and runs the indexer, which writes the selected blocks into |
| 390 | the shared ``layer.topk_indices_buffer``; the attend reads them back from |
| 391 | there). The Triton and MSA subclasses each own a full ``forward`` -- no |
| 392 | shared forward code. |
| 393 | """ |
| 394 | |
| 395 | def __init__( |
| 396 | self, |
| 397 | num_heads: int, |
| 398 | head_size: int, |
| 399 | scale: float, |
| 400 | num_kv_heads: int | None = None, |
| 401 | kv_cache_dtype: str = "auto", |
| 402 | *, |
| 403 | topk_blocks: int, |
| 404 | sparse_block_size: int, |
| 405 | msa_decode_backend: str = "triton", |
| 406 | ) -> None: |
| 407 | self.num_heads = num_heads |
| 408 | self.head_size = head_size |
| 409 | self.scale = scale |
| 410 | self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads |
| 411 | self.kv_cache_dtype = kv_cache_dtype |
| 412 | self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype) |
| 413 | if "e5m2" in kv_cache_dtype: |
| 414 | self.kv_cache_fp8_dtype = ( |
| 415 | torch.float8_e5m2fnuz |
| 416 | if current_platform.is_fp8_fnuz() |
| 417 | else torch.float8_e5m2 |
| 418 | ) |
| 419 | else: |
| 420 | self.kv_cache_fp8_dtype = current_platform.fp8_dtype() |
| 421 | # Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE). |
| 422 | self.topk_blocks = topk_blocks |
| 423 | self.block_size = sparse_block_size |
| 424 | |
| 425 | def forward( |
| 426 | self, |
| 427 | layer: AttentionLayer, |
| 428 | query: torch.Tensor, |
| 429 | kv_cache: torch.Tensor, |
| 430 | output: torch.Tensor, |
| 431 | *, |
| 432 | query_fp8: torch.Tensor | None = None, |
| 433 | ) -> torch.Tensor: |
| 434 | """Attend the queries to the indexer-selected blocks. Per kernel. |
| 435 | |
| 436 | The indexer has already written the top-k block ids into |
| 437 | ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at |
| 438 | ``[:, nd:num_tokens]``); the attend reads them from there. |
| 439 | """ |
| 440 | raise NotImplementedError |
| 441 | |
| 442 | def should_use_msa_decode(self, layer_name: str) -> bool: |
| 443 | return False |
| 444 | |
| 445 | |
| 446 | class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl): |
| 447 | """Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode.""" |
| 448 | |
| 449 | def forward( |
| 450 | self, |
| 451 | layer: AttentionLayer, |
| 452 | query: torch.Tensor, |
| 453 | kv_cache: torch.Tensor, |
| 454 | output: torch.Tensor, |
| 455 | *, |
| 456 | query_fp8: torch.Tensor | None = None, |
| 457 | ) -> torch.Tensor: |
| 458 | attn_metadata = get_forward_context().attn_metadata |
| 459 | if not isinstance(attn_metadata, dict): |
| 460 | return output # profiling run; caches unbound |
| 461 | main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined] |
| 462 | assert isinstance(main_md, MiniMaxM3SparseMetadata) |
| 463 | |
| 464 | nd = main_md.num_decode_tokens |
| 465 | num_tokens = main_md.num_actual_tokens |
| 466 | # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:]. |
| 467 | topk_buffer = layer.topk_indices_buffer # type: ignore[attr-defined] |
| 468 | assert topk_buffer is not None |
| 469 | |
| 470 | topk = ( |
| 471 | topk_buffer |
| 472 | if current_platform.is_rocm() |
| 473 | else topk_buffer[:num_tokens].transpose(0, 1) |
| 474 | ) |
| 475 | assert topk is not None |
| 476 | hd = self.head_size |
| 477 | q = query[:num_tokens].view(-1, self.num_heads, hd) |
| 478 | out = output[:num_tokens].view(-1, self.num_heads, hd) |
| 479 | kv_cache = ( |
| 480 | kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache |
| 481 | ) |
| 482 | k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None |
| 483 | v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None |
| 484 | |
| 485 | # Decode [:nd]: split-K over the selected blocks (request-major chunks). |
| 486 | if main_md.num_decodes > 0: |
| 487 | d = main_md.decode |
| 488 | assert d is not None |
| 489 | minimax_m3_sparse_attn_decode( |
| 490 | q[:nd], |
| 491 | kv_cache, |
| 492 | topk[:, :nd, :], |
| 493 | d.block_table, |
| 494 | d.seq_lens, |
| 495 | self.num_kv_heads, |
| 496 | self.scale, |
| 497 | out[:nd], |
| 498 | d.decode_query_len, |
| 499 | k_scale=k_scale, |
| 500 | v_scale=v_scale, |
| 501 | ) |
| 502 | |
| 503 | # Prefill [nd:]: cu_seqlens_q already rebased to 0. |
| 504 | if main_md.num_prefills > 0: |
| 505 | p = main_md.prefill |
| 506 | assert p is not None |
| 507 | minimax_m3_sparse_attn( |
| 508 | q[nd:], |
| 509 | kv_cache, |
| 510 | topk[:, nd:num_tokens, :], |
| 511 | p.block_table, |
| 512 | p.cu_seqlens_q, |
| 513 | p.seq_lens, |
| 514 | p.context_lens, |
| 515 | p.max_query_len, |
| 516 | self.num_kv_heads, |
| 517 | self.scale, |
| 518 | out[nd:], |
| 519 | k_scale=k_scale, |
| 520 | v_scale=v_scale, |
| 521 | ) |
| 522 | return output |
| 523 | |
| 524 | |
| 525 | def select_main_backend_and_impl_cls( |
| 526 | *, |
| 527 | topk_blocks: int, |
| 528 | kv_cache_dtype: str, |
| 529 | num_kv_heads: int, |
| 530 | emits_sparse_block_table: bool = False, |
| 531 | ) -> tuple[type[MiniMaxM3SparseBackend], type[MiniMaxM3SparseImpl]]: |
| 532 | """Pick the main attention backend and implementation. |
| 533 | |
| 534 | Blackwell (SM100) uses the MSA attend for supported top-k block counts |
| 535 | when the KV cache is BF16, FP8 E4M3 or NVFP4; MI355 uses AITER sparse PA |
| 536 | with shuffle KV cache layout; Other platforms and FP8 E5M2 fall |
| 537 | back to Triton. NVFP4 is MSA-only. The MSA modules are imported lazily to |
| 538 | avoid import errors on unsupported platforms. |
| 539 | """ |
| 540 | use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa( |
| 541 | num_kv_heads, emits_sparse_block_table=emits_sparse_block_table |
| 542 | ) |
| 543 | use_msa = ( |
| 544 | current_platform.is_cuda() |
| 545 | and current_platform.is_device_capability_family(100) |
| 546 | and topk_blocks in (4, 8, 16, 32) |
| 547 | and kv_cache_dtype != "fp8_e5m2" |
| 548 | ) |
| 549 | use_nvfp4 = get_kv_quant_mode(kv_cache_dtype).is_nvfp4 |
| 550 | if use_nvfp4 and kv_cache_dtype != "nvfp4": |
| 551 | raise ValueError( |
| 552 | "MiniMax M3 sparse attention supports only the 'nvfp4' NVFP4 KV " |
| 553 | f"cache dtype, got {kv_cache_dtype}" |
| 554 | ) |
| 555 | if use_nvfp4 and (use_aiter_sparse_pa or not use_msa): |
| 556 | raise ValueError( |
| 557 | "MiniMax M3 sparse attention supports an NVFP4 KV cache only with " |
| 558 | f"the SM100 MSA backend (kv_cache_dtype={kv_cache_dtype}, " |
| 559 | f"topk_blocks={topk_blocks})" |
| 560 | ) |
| 561 | selected = ( |
| 562 | "AITER_SPARSE_PA" if use_aiter_sparse_pa else ("MSA" if use_msa else "Triton") |
| 563 | ) |
| 564 | logger.info_once( |
| 565 | "MiniMax M3 sparse attention selected %s (kv_cache_dtype=%s, topk_blocks=%s)", |
| 566 | selected, |
| 567 | kv_cache_dtype, |
| 568 | topk_blocks, |
| 569 | ) |
| 570 | if use_aiter_sparse_pa: |
| 571 | from vllm.models.minimax_m3.amd.sparse_attention_msa import ( |
| 572 | MiniMaxM3SparseAiterPABackend, |
| 573 | MiniMaxM3SparseAiterPAImpl, |
| 574 | ) |
| 575 | |
| 576 | return MiniMaxM3SparseAiterPABackend, MiniMaxM3SparseAiterPAImpl |
| 577 | if use_msa: |
| 578 | from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( |
| 579 | MiniMaxM3SparseMSABackend, |
| 580 | MiniMaxM3SparseMSAImpl, |
| 581 | MiniMaxM3SparseMSANvfp4Backend, |
| 582 | ) |
| 583 | |
| 584 | if use_nvfp4: |
| 585 | return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSAImpl |
| 586 | return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl |
| 587 | return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl |
| 588 | |
| 589 | |
| 590 | def select_main_impl_cls( |
| 591 | *, |
| 592 | topk_blocks: int, |
| 593 | kv_cache_dtype: str, |
| 594 | num_kv_heads: int, |
| 595 | emits_sparse_block_table: bool = False, |
| 596 | ) -> type[MiniMaxM3SparseImpl]: |
| 597 | """Backward-compatible implementation-only selector.""" |
| 598 | return select_main_backend_and_impl_cls( |
| 599 | topk_blocks=topk_blocks, |
| 600 | kv_cache_dtype=kv_cache_dtype, |
| 601 | num_kv_heads=num_kv_heads, |
| 602 | emits_sparse_block_table=emits_sparse_block_table, |
| 603 | )[1] |