| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 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): |
| 120 | """ROCm AITER page-16 SHUFFLE sparse paged attention.""" |
| 121 | |
| 122 | def forward( |
| 123 | self, |
| 124 | layer: AttentionLayer, |
| 125 | query: torch.Tensor, |
| 126 | kv_cache: torch.Tensor, |
| 127 | output: torch.Tensor, |
| 128 | *, |
| 129 | query_fp8: torch.Tensor | None = None, |
| 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 |
| 141 | if not isinstance(attn_metadata, dict): |
| 142 | return output |
| 143 | main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined] |
| 144 | assert isinstance(main_md, MiniMaxM3SparseMetadata) |
| 145 | |
| 146 | nd = main_md.num_decode_tokens |
| 147 | num_tokens = main_md.num_actual_tokens |
| 148 | topk = layer.topk_indices_buffer # type: ignore[attr-defined] |
| 149 | assert topk is not None |
| 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( |
| 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 |
| 163 | q = query[:num_tokens].view(-1, self.num_heads, hd) |
| 164 | out = output[:num_tokens].view(-1, self.num_heads, hd) |
| 165 | k_cache, v_cache = layer.get_aiter_sparse_pa_kv_cache() # type: ignore[attr-defined] |
| 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, |
| 197 | v_cache, |
| 198 | topk[:, :nd, :], |
| 199 | d.block_table, |
| 200 | d.seq_lens, |
| 201 | self.num_kv_heads, |
| 202 | self.scale, |
| 203 | out[:nd], |
| 204 | k_scale=k_scale, |
| 205 | v_scale=v_scale, |
| 206 | decode_query_len=d.decode_query_len, |
| 207 | sparse_block_table=( |
| 208 | prepared_decode_table[0] |
| 209 | if prepared_decode_table is not None |
| 210 | else None |
| 211 | ), |
| 212 | sparse_context_lens=( |
| 213 | prepared_decode_table[1] |
| 214 | if prepared_decode_table is not None |
| 215 | else None |
| 216 | ), |
| 217 | ) |
| 218 | |
| 219 | if main_md.num_prefills > 0: |
| 220 | p = main_md.prefill |
| 221 | assert p is not None |
| 222 | assert p.query_req_id is not None and p.query_abs_pos is not None |
| 223 | minimax_m3_sparse_attn_prefill_aiter( |
| 224 | q[nd:], |
| 225 | k_cache, |
| 226 | v_cache, |
| 227 | topk[:, nd:num_tokens, :], |
| 228 | p.block_table, |
| 229 | p.query_req_id, |
| 230 | p.query_abs_pos, |
| 231 | self.num_kv_heads, |
| 232 | self.scale, |
| 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 |