| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MSA (SM100/Blackwell) indexer impl for MiniMax M3. |
| 4 | |
| 5 | Both sides write block scores into one unified token-major buffer |
| 6 | ``[total_q, H, max_k_tiles]``, then a single ``fmha_sm100.sparse_topk_select`` |
| 7 | selects the top-k blocks for the whole batch (decode ``[:nd]`` + prefill |
| 8 | ``[nd:]``) into the shared ``topk_indices_buffer``. It bounds each row by its |
| 9 | causal page count and force-includes the init/local blocks, so the unwritten |
| 10 | tail of the buffer is pre-filled with ``-inf``. |
| 11 | |
| 12 | Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faster |
| 13 | than Triton for the wide prefill score, benchmarked ~3-5x), writing its |
| 14 | ``max_score`` straight into the buffer's prefill region (stride-aware, no copy). |
| 15 | |
| 16 | Decode scores with CuteDSL when the flattened query tile is supported and fall |
| 17 | back to Triton otherwise, writing into the decode region. Only the top-k is |
| 18 | shared with prefill. |
| 19 | |
| 20 | ``fmha_sm100`` imports are function-local so this module is import-safe on |
| 21 | AMD / non-SM100. |
| 22 | """ |
| 23 | |
| 24 | from dataclasses import dataclass |
| 25 | from typing import ClassVar |
| 26 | |
| 27 | import torch |
| 28 | |
| 29 | from vllm.config import VllmConfig |
| 30 | from vllm.forward_context import get_forward_context |
| 31 | from vllm.models.minimax_m3.common.indexer import ( |
| 32 | MiniMaxM3IndexerBackend, |
| 33 | MiniMaxM3IndexerDecodeMetadata, |
| 34 | MiniMaxM3IndexerImpl, |
| 35 | MiniMaxM3IndexerMetadata, |
| 36 | MiniMaxM3IndexerMetadataBuilder, |
| 37 | ) |
| 38 | from vllm.models.minimax_m3.common.ops.index_topk import ( |
| 39 | minimax_m3_index_decode_score, |
| 40 | ) |
| 41 | from vllm.models.minimax_m3.nvidia.ops import minimax_m3_index_decode_score_cutedsl |
| 42 | from vllm.v1.attention.backend import ( |
| 43 | AttentionBackend, |
| 44 | AttentionCGSupport, |
| 45 | CommonAttentionMetadata, |
| 46 | ) |
| 47 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 48 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 49 | |
| 50 | # Page size == sparse block size == index-K block; fmha tile id == M3 block id. |
| 51 | PAGE_SIZE = 128 |
| 52 | |
| 53 | # Fill for unwritten score tiles: -inf so they never win the top-k (score kernels |
| 54 | # only write causally-valid blocks). |
| 55 | _SCORE_SENTINEL = float("-inf") |
| 56 | |
| 57 | # Tile (KV-block) dim of the unified score buffer, hardcoded as a cudagraph |
| 58 | # capture-time constant so the decode score kernel's buffer shape is frozen |
| 59 | # across replays. 8192 tiles == 1M tokens of context; -inf padding + |
| 60 | # num_valid_pages bound each row to its causal range, so shorter replays reuse |
| 61 | # the same buffer safely. |
| 62 | MAX_K_TILES = 8192 |
| 63 | |
| 64 | |
| 65 | class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend): |
| 66 | """Indexer side-cache backend selecting the MSA builder.""" |
| 67 | |
| 68 | @staticmethod |
| 69 | def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]: |
| 70 | return MiniMaxM3IndexerMSAMetadataBuilder |
| 71 | |
| 72 | |
| 73 | @dataclass |
| 74 | class MiniMaxM3IndexerMSAPrefillMetadata: |
| 75 | """fmha score plan + Triton top-k inputs for the prefill side (eager).""" |
| 76 | |
| 77 | plan: dict # fmha_sm100 PlanInfo |
| 78 | cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0 |
| 79 | prefix_lens: torch.Tensor # [num_prefills] int32, context tokens |
| 80 | max_query_len: int |
| 81 | page_table: torch.Tensor # flat physical page indices for the prefill side |
| 82 | |
| 83 | |
| 84 | @dataclass |
| 85 | class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): |
| 86 | """Decode reuses the inherited base ``decode`` field (the Triton decode |
| 87 | metadata); ``prefill_msa`` carries the fmha score plan for the prefill side |
| 88 | (the base ``prefill`` field is unused on this path).""" |
| 89 | |
| 90 | prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None |
| 91 | # Per-forward view (``[:num_tokens]``) of the builder's persistent unified |
| 92 | # score buffer ``[total_q, H, MAX_K_TILES]``, shared by decode and prefill |
| 93 | # and reused across all layers. Pre-filled with the -inf sentinel in |
| 94 | # ``build()``; each layer overwrites its valid tiles before its own top-k. |
| 95 | unified_scores: torch.Tensor | None = None |
| 96 | # Tile (KV-block) dim of the unified score buffer (== ``MAX_K_TILES``). |
| 97 | # Forced as the fmha plan's ``max_k_tiles`` so prefill writes its max_score |
| 98 | # straight into the shared buffer. |
| 99 | max_k_tiles: int = 0 |
| 100 | # Batch-wide inputs for the single top-k over the unified buffer (decode + |
| 101 | # prefill in one call). ``cu_seqlens_q`` is the per-request query-start |
| 102 | # offsets (== query_start_loc, a stable view for cudagraph); ``prefix_lens`` |
| 103 | # is the per-request context length (== context_lens). |
| 104 | topk_cu_seqlens_q: torch.Tensor | None = None |
| 105 | topk_prefix_lens: torch.Tensor | None = None |
| 106 | topk_max_query_len: int = 0 |
| 107 | # Per-token causal page count cdiv(seq_pos+1, PAGE_SIZE), [total_q] int32: |
| 108 | # drives sparse_topk_select force_end_blocks + -1 out-of-range clamp. |
| 109 | topk_num_valid_pages: torch.Tensor | None = None |
| 110 | |
| 111 | |
| 112 | class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): |
| 113 | """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill |
| 114 | fmha plan is built eagerly (prefill batches are not captured).""" |
| 115 | |
| 116 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 117 | |
| 118 | def __init__( |
| 119 | self, |
| 120 | kv_cache_spec: AttentionSpec, |
| 121 | layer_names: list[str], |
| 122 | vllm_config: VllmConfig, |
| 123 | device: torch.device, |
| 124 | ) -> None: |
| 125 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 126 | # Persistent unified score buffer [T, H, MAX_K_TILES] shared by all |
| 127 | # indexer layers and reused across forwards. Stable address (required |
| 128 | # for the captured decode path) + fixed tile dim so the decode score |
| 129 | # kernel's shape is frozen at capture. Filled with -inf per forward in |
| 130 | # build(); the valid/padding partition is a per-forward constant, so |
| 131 | # every layer overwrites the same valid tiles before its own top-k. |
| 132 | self.unified_scores_buffer = torch.empty( |
| 133 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 134 | self.num_index_heads, |
| 135 | MAX_K_TILES, |
| 136 | dtype=torch.float32, |
| 137 | device=device, |
| 138 | ) |
| 139 | |
| 140 | def build( |
| 141 | self, |
| 142 | common_prefix_len: int, |
| 143 | common_attn_metadata: CommonAttentionMetadata, |
| 144 | fast_build: bool = False, |
| 145 | ) -> MiniMaxM3IndexerMSAMetadata: |
| 146 | num_reqs = common_attn_metadata.num_reqs |
| 147 | num_tokens = common_attn_metadata.num_actual_tokens |
| 148 | seq_lens = common_attn_metadata.seq_lens |
| 149 | block_table = common_attn_metadata.block_table_tensor |
| 150 | query_start_loc = common_attn_metadata.query_start_loc |
| 151 | |
| 152 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 153 | split_decodes_and_prefills( |
| 154 | common_attn_metadata, |
| 155 | decode_threshold=self.reorder_batch_threshold, |
| 156 | require_uniform=True, |
| 157 | ) |
| 158 | ) |
| 159 | assert num_decodes + num_prefills == num_reqs |
| 160 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 161 | |
| 162 | # Context (prefix) lengths into the stable cudagraph buffer. |
| 163 | context_lens = self.context_len_buffer[:num_reqs] |
| 164 | context_lens.copy_( |
| 165 | common_attn_metadata.compute_num_computed_tokens(), non_blocking=True |
| 166 | ) |
| 167 | |
| 168 | # Per-token causal page count for the top-k, into the stable cg buffer. |
| 169 | positions = common_attn_metadata.positions |
| 170 | assert positions is not None |
| 171 | num_valid_pages = self.num_valid_pages_buffer[:num_tokens] |
| 172 | num_valid_pages.copy_(positions[:num_tokens] // PAGE_SIZE + 1) |
| 173 | |
| 174 | # Unified score buffer: a per-forward view of the persistent buffer, |
| 175 | # reset to the -inf sentinel once here and shared by every layer (the |
| 176 | # tile dim is the capture-time constant MAX_K_TILES). |
| 177 | max_k_tiles = MAX_K_TILES |
| 178 | unified_scores = self.unified_scores_buffer[:num_tokens] |
| 179 | unified_scores.fill_(_SCORE_SENTINEL) |
| 180 | |
| 181 | decode: MiniMaxM3IndexerDecodeMetadata | None = None |
| 182 | if num_decodes > 0: |
| 183 | qsl_cpu = common_attn_metadata.query_start_loc_cpu |
| 184 | query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] |
| 185 | decode_query_len = int(query_lens_cpu[0].item()) |
| 186 | assert decode_query_len > 0 |
| 187 | assert torch.all( |
| 188 | (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) |
| 189 | ) |
| 190 | decode = MiniMaxM3IndexerDecodeMetadata( |
| 191 | seq_lens=seq_lens[:num_decodes], |
| 192 | block_table=block_table[:num_decodes], |
| 193 | max_seq_len=common_attn_metadata.max_seq_len, |
| 194 | decode_query_len=decode_query_len, |
| 195 | max_decode_query_len=self.max_decode_query_len, |
| 196 | ) |
| 197 | |
| 198 | prefill: MiniMaxM3IndexerMSAPrefillMetadata | None = None |
| 199 | if num_prefills > 0: |
| 200 | # Prefill is eager (not captured); the host lengths it needs (and the |
| 201 | # _fmha_sm100_plan .tolist() inside) make the D->H sync acceptable. |
| 202 | from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan |
| 203 | |
| 204 | lo, hi = num_decodes, num_reqs |
| 205 | qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1] |
| 206 | qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32) |
| 207 | kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32) |
| 208 | nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE |
| 209 | side_qo = qo_lens_cpu[lo:hi] |
| 210 | side_kv = kv_lens_cpu[lo:hi] |
| 211 | plan = _fmha_sm100_plan( |
| 212 | side_qo, |
| 213 | side_kv, |
| 214 | self.num_index_heads, |
| 215 | num_kv_heads=1, |
| 216 | qo_offset=side_kv - side_qo, # bottom-right causal |
| 217 | page_size=PAGE_SIZE, |
| 218 | output_maxscore=True, |
| 219 | causal=True, |
| 220 | num_kv_splits=1, |
| 221 | ) |
| 222 | # Force the plan's tile dim to the unified buffer's so prefill writes |
| 223 | # its max_score straight into unified[:, :, nd:] (the stride-aware |
| 224 | # binding shape-matches the tile dim exactly). max_k_tiles >= the |
| 225 | # plan's natural value, so the extra tiles are simply never written. |
| 226 | plan["max_k_tiles"] = max_k_tiles |
| 227 | cols = torch.arange(block_table.shape[1], device=block_table.device) |
| 228 | valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None] |
| 229 | prefill = MiniMaxM3IndexerMSAPrefillMetadata( |
| 230 | plan=plan, |
| 231 | cu_seqlens_q=(query_start_loc[lo : hi + 1] - query_start_loc[lo]).to( |
| 232 | torch.int32 |
| 233 | ), |
| 234 | prefix_lens=context_lens[lo:hi], |
| 235 | max_query_len=int(side_qo.max()), |
| 236 | page_table=block_table[lo:hi][valid].to(torch.int32), |
| 237 | ) |
| 238 | |
| 239 | return MiniMaxM3IndexerMSAMetadata( |
| 240 | seq_lens=seq_lens, |
| 241 | max_seq_len=common_attn_metadata.max_seq_len, |
| 242 | slot_mapping=common_attn_metadata.slot_mapping, |
| 243 | num_actual_tokens=num_tokens, |
| 244 | num_decodes=num_decodes, |
| 245 | num_decode_tokens=num_decode_tokens, |
| 246 | num_prefills=num_prefills, |
| 247 | num_prefill_tokens=num_prefill_tokens, |
| 248 | decode=decode, |
| 249 | prefill_msa=prefill, |
| 250 | unified_scores=unified_scores, |
| 251 | max_k_tiles=max_k_tiles, |
| 252 | topk_cu_seqlens_q=query_start_loc[: num_reqs + 1], |
| 253 | topk_prefix_lens=context_lens, |
| 254 | topk_max_query_len=common_attn_metadata.max_query_len, |
| 255 | topk_num_valid_pages=num_valid_pages, |
| 256 | ) |
| 257 | |
| 258 | |
| 259 | class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): |
| 260 | """Decode: CuteDSL/Triton score. Prefill: fmha_sm100 OnlyScore + top-k.""" |
| 261 | |
| 262 | indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend |
| 263 | |
| 264 | def forward( |
| 265 | self, |
| 266 | index_query: torch.Tensor, |
| 267 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 268 | attn_metadata = get_forward_context().attn_metadata |
| 269 | if not isinstance(attn_metadata, dict): |
| 270 | return None, None # profiling run; caches unbound |
| 271 | md = attn_metadata[self.index_cache.prefix] |
| 272 | assert isinstance(md, MiniMaxM3IndexerMSAMetadata) |
| 273 | |
| 274 | num_tokens = md.num_actual_tokens |
| 275 | nd = md.num_decode_tokens |
| 276 | index_q = index_query[:num_tokens].view( |
| 277 | -1, self.num_index_heads, self.index_head_dim |
| 278 | ) |
| 279 | kv = self.index_cache.kv_cache |
| 280 | # Shared persistent top-k output buffer; the unified top-k below writes |
| 281 | # the selected block ids into buf[:num_tokens]. |
| 282 | buf = self.topk_indices_buffer |
| 283 | assert buf is not None |
| 284 | |
| 285 | # Unified token-major score buffer [total_q, H, MAX_K_TILES]: the tile |
| 286 | # dim is innermost/contiguous, so both fmha writes (native [T,H,K]) and |
| 287 | # the block-iterating top-k reads hit contiguous tiles. Each side gets a |
| 288 | # contiguous slice on dim 0: decode [:nd], prefill [nd:]; the kernels |
| 289 | # read/write by strides. The builder allocates it once (persistent, |
| 290 | # shared by all layers) and resets it to the sentinel each forward, so |
| 291 | # the top-k never picks an unwritten tile. |
| 292 | unified_scores = md.unified_scores |
| 293 | assert unified_scores is not None |
| 294 | |
| 295 | # Decode scores -> unified[:nd] (transposed [H, nd, MK] view; the kernel |
| 296 | # writes by strides). Top-k is deferred to the single unified call below. |
| 297 | if md.decode is not None: |
| 298 | d = md.decode |
| 299 | # max_decode_query_len avoids recompiles across runtime decode sizes. |
| 300 | # Fall back when the flattened Q tile gets too wide for this kernel. |
| 301 | decode_score = ( |
| 302 | minimax_m3_index_decode_score_cutedsl |
| 303 | if self.num_index_heads * d.max_decode_query_len <= 32 |
| 304 | else minimax_m3_index_decode_score |
| 305 | ) |
| 306 | decode_score( |
| 307 | index_q[:nd], |
| 308 | kv, |
| 309 | d.block_table, |
| 310 | d.seq_lens, |
| 311 | d.max_seq_len, |
| 312 | self.init_blocks, |
| 313 | self.local_blocks, |
| 314 | self.num_kv_heads, |
| 315 | d.decode_query_len, |
| 316 | d.max_decode_query_len, |
| 317 | score_out=unified_scores[:nd].transpose(0, 1), |
| 318 | ) |
| 319 | |
| 320 | if md.prefill_msa is not None: |
| 321 | from vllm.third_party.fmha_sm100.api import _fmha_sm100 |
| 322 | |
| 323 | p = md.prefill_msa |
| 324 | # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D). |
| 325 | k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim) |
| 326 | # fmha writes its max_score natively as [nnz_p, H, max_k_tiles] into |
| 327 | # the prefill region (stride-aware; plan max_k_tiles forced to |
| 328 | # md.max_k_tiles so the shape matches exactly -> no copy). |
| 329 | _fmha_sm100( |
| 330 | index_q[nd:], |
| 331 | k_pages, |
| 332 | k_pages, # V placeholder; not read in OnlyScore |
| 333 | p.plan, |
| 334 | kv_indices=p.page_table, |
| 335 | output_o=False, |
| 336 | output_maxscore=True, |
| 337 | sm_scale=self.scale, |
| 338 | max_score=unified_scores[nd:], |
| 339 | ) |
| 340 | |
| 341 | # Single top-k over the unified buffer via fmha_sm100 sparse_topk_select |
| 342 | # (THK, no transpose) into ``buf``. num_valid_pages drives force_end_blocks |
| 343 | # (always keep each token's local block; fmha OnlyScore won't) + the -1 |
| 344 | # out-of-range clamp. |
| 345 | from vllm.third_party.fmha_sm100.api import sparse_topk_select |
| 346 | |
| 347 | sparse_topk_select( |
| 348 | unified_scores, |
| 349 | self.topk_blocks, |
| 350 | num_valid_pages=md.topk_num_valid_pages, |
| 351 | force_begin_blocks=self.init_blocks, |
| 352 | force_end_blocks=self.local_blocks, |
| 353 | output=buf[:num_tokens], |
| 354 | max_score_layout="THK", |
| 355 | ) |
| 356 | |
| 357 | # The attend reads ``buf`` directly; this return is vestigial. |
| 358 | return None, None |