| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MiniMax M3 lightning indexer: side cache, metadata, and impl. |
| 4 | |
| 5 | The indexer scores KV blocks with the index heads and selects the top-k blocks |
| 6 | (plus fixed init/local blocks) that the main block-sparse attention |
| 7 | (``sparse_attention.py``) then attends to. It owns its own side cache |
| 8 | (``MiniMaxM3IndexerCache``, one index-key vector per token), metadata, and |
| 9 | metadata builder, mirroring how DeepSeek V4 keeps the indexer separate from the |
| 10 | main attention. |
| 11 | |
| 12 | ``MiniMaxM3Indexer`` is the ``nn.Module`` the attention layer holds (like |
| 13 | ``DeepseekV4Indexer``); it picks a kernel impl in ``__init__`` (via |
| 14 | ``select_indexer_impl_cls``) and delegates ``forward`` to it. |
| 15 | """ |
| 16 | |
| 17 | from dataclasses import dataclass |
| 18 | from typing import ClassVar |
| 19 | |
| 20 | import torch |
| 21 | from torch import nn |
| 22 | |
| 23 | from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config |
| 24 | from vllm.config.attention import IndexerKVDType |
| 25 | from vllm.config.cache import CacheDType |
| 26 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 27 | from vllm.forward_context import get_forward_context |
| 28 | from vllm.logger import init_logger |
| 29 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 30 | from vllm.platforms import current_platform |
| 31 | |
| 32 | if current_platform.is_rocm(): |
| 33 | from vllm.models.minimax_m3.amd.ops.index_topk import ( |
| 34 | minimax_m3_index_decode, |
| 35 | minimax_m3_index_score, |
| 36 | minimax_m3_index_topk, |
| 37 | ) |
| 38 | else: |
| 39 | from vllm.models.minimax_m3.common.ops.index_topk import ( |
| 40 | minimax_m3_index_decode, |
| 41 | minimax_m3_index_score, |
| 42 | minimax_m3_index_topk, |
| 43 | ) |
| 44 | |
| 45 | from vllm.v1.attention.backend import ( |
| 46 | AttentionBackend, |
| 47 | AttentionCGSupport, |
| 48 | AttentionMetadata, |
| 49 | AttentionMetadataBuilder, |
| 50 | CommonAttentionMetadata, |
| 51 | MultipleOf, |
| 52 | ) |
| 53 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 54 | from vllm.v1.kv_cache_interface import ( |
| 55 | AttentionSpec, |
| 56 | KVCacheSpec, |
| 57 | MLAAttentionSpec, |
| 58 | SparseCacheRole, |
| 59 | ) |
| 60 | |
| 61 | logger = init_logger(__name__) |
| 62 | |
| 63 | |
| 64 | class MiniMaxM3IndexerBackend(AttentionBackend): |
| 65 | """Indexer side-cache backend (key-only).""" |
| 66 | |
| 67 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16] |
| 68 | # bf16 today; mirrors the main backend to keep spec validation permissive. |
| 69 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 70 | "bfloat16", |
| 71 | "fp8", |
| 72 | "fp8_e4m3", |
| 73 | "fp8_e5m2", |
| 74 | ] |
| 75 | |
| 76 | @staticmethod |
| 77 | def get_name() -> str: |
| 78 | return "MINIMAX_M3_SPARSE_INDEXER" |
| 79 | |
| 80 | @staticmethod |
| 81 | def get_impl_cls() -> type["MiniMaxM3IndexerImpl"]: |
| 82 | # Concrete impl chosen by select_indexer_impl_cls; base for introspection. |
| 83 | return MiniMaxM3IndexerImpl |
| 84 | |
| 85 | @staticmethod |
| 86 | def get_builder_cls() -> type["MiniMaxM3IndexerMetadataBuilder"]: |
| 87 | return MiniMaxM3IndexerTritonMetadataBuilder |
| 88 | |
| 89 | @classmethod |
| 90 | def get_supported_head_sizes(cls) -> list[int]: |
| 91 | return [128] |
| 92 | |
| 93 | @staticmethod |
| 94 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 95 | return [128] |
| 96 | |
| 97 | @classmethod |
| 98 | def is_sparse(cls) -> bool: |
| 99 | return True |
| 100 | |
| 101 | |
| 102 | class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase): |
| 103 | """Side KV cache for the indexer's per-token index keys (key-only). |
| 104 | |
| 105 | Registers itself in the static forward context so the KV-cache manager |
| 106 | allocates it (like ``DeepseekV32IndexerCache``). |
| 107 | """ |
| 108 | |
| 109 | def __init__( |
| 110 | self, |
| 111 | head_dim: int, |
| 112 | prefix: str, |
| 113 | cache_config: CacheConfig | None = None, |
| 114 | indexer_kv_dtype: IndexerKVDType = "bf16", |
| 115 | backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend, |
| 116 | ) -> None: |
| 117 | super().__init__() |
| 118 | if indexer_kv_dtype in ("fp8", "fp8_e4m3"): |
| 119 | cache_dtype = torch.float8_e4m3fn |
| 120 | elif indexer_kv_dtype == "bf16": |
| 121 | cache_dtype = torch.bfloat16 |
| 122 | else: |
| 123 | raise NotImplementedError( |
| 124 | f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the " |
| 125 | "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')." |
| 126 | ) |
| 127 | self.kv_cache = torch.tensor([]) |
| 128 | self.head_dim = head_dim |
| 129 | self.indexer_kv_dtype = indexer_kv_dtype |
| 130 | # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path. |
| 131 | self.dtype = cache_dtype |
| 132 | self.prefix = prefix |
| 133 | self.cache_config = cache_config |
| 134 | # Impl-chosen backend -> each impl gets its own builder (get_attn_backend). |
| 135 | self.backend_cls = backend_cls |
| 136 | compilation_config = get_current_vllm_config().compilation_config |
| 137 | if prefix in compilation_config.static_forward_context: |
| 138 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 139 | compilation_config.static_forward_context[prefix] = self |
| 140 | |
| 141 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 142 | # [B, H=1, N, C] -> [B, N, C] |
| 143 | self.kv_cache = kv_cache.squeeze(1) |
| 144 | |
| 145 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 146 | # Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V). |
| 147 | return MLAAttentionSpec( |
| 148 | block_size=vllm_config.cache_config.block_size, |
| 149 | num_kv_heads=1, |
| 150 | head_size=self.head_dim, |
| 151 | dtype=self.dtype, |
| 152 | cache_role=SparseCacheRole.INDEXER, |
| 153 | ) |
| 154 | |
| 155 | def forward(self) -> None: ... |
| 156 | |
| 157 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 158 | return self.backend_cls |
| 159 | |
| 160 | |
| 161 | @dataclass |
| 162 | class MiniMaxM3IndexerPrefillMetadata: |
| 163 | """Per-prefill index-scoring state.""" |
| 164 | |
| 165 | cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0 |
| 166 | seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths |
| 167 | context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens) |
| 168 | block_table: torch.Tensor |
| 169 | max_query_len: int |
| 170 | max_seq_len: int |
| 171 | |
| 172 | |
| 173 | @dataclass |
| 174 | class MiniMaxM3IndexerDecodeMetadata: |
| 175 | """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform |
| 176 | per-request query length (1, or 1 + num_speculative_tokens).""" |
| 177 | |
| 178 | seq_lens: torch.Tensor # [num_decodes] int32 |
| 179 | block_table: torch.Tensor |
| 180 | max_seq_len: int |
| 181 | decode_query_len: int |
| 182 | max_decode_query_len: int |
| 183 | |
| 184 | |
| 185 | @dataclass |
| 186 | class MiniMaxM3IndexerMetadata(AttentionMetadata): |
| 187 | """Indexer metadata, split into prefill and decode sub-metadata.""" |
| 188 | |
| 189 | seq_lens: torch.Tensor |
| 190 | max_seq_len: int |
| 191 | slot_mapping: torch.Tensor |
| 192 | |
| 193 | num_actual_tokens: int # total query tokens (decode-first batch) |
| 194 | |
| 195 | # Split counts; identical to the main metadata's (same reorder threshold). |
| 196 | num_decodes: int |
| 197 | num_decode_tokens: int |
| 198 | num_prefills: int |
| 199 | num_prefill_tokens: int |
| 200 | |
| 201 | prefill: MiniMaxM3IndexerPrefillMetadata | None = None |
| 202 | decode: MiniMaxM3IndexerDecodeMetadata | None = None |
| 203 | |
| 204 | |
| 205 | class MiniMaxM3IndexerMetadataBuilder( |
| 206 | AttentionMetadataBuilder[MiniMaxM3IndexerMetadata] |
| 207 | ): |
| 208 | """Abstract base: shared setup only. The Triton and MSA builders are |
| 209 | parallel subclasses that each own their full ``build`` (no shared code).""" |
| 210 | |
| 211 | # Full cudagraphs for uniform decode batches (incl. spec-decode verify). |
| 212 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 213 | # Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when |
| 214 | # spec decode is on; matches the main builder so the splits agree. |
| 215 | reorder_batch_threshold: int = 1 |
| 216 | |
| 217 | def __init__( |
| 218 | self, |
| 219 | kv_cache_spec: AttentionSpec, |
| 220 | layer_names: list[str], |
| 221 | vllm_config: VllmConfig, |
| 222 | device: torch.device, |
| 223 | ) -> None: |
| 224 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 225 | hf_config = vllm_config.model_config.hf_config |
| 226 | text_config = getattr(hf_config, "text_config", hf_config) |
| 227 | sparse_cfg = text_config.sparse_attention_config |
| 228 | # Index-query head count from model config (cache spec has 1 vec/token). |
| 229 | total_index_heads = sparse_cfg["sparse_num_index_heads"] |
| 230 | tp_size = get_tensor_model_parallel_world_size() |
| 231 | if total_index_heads >= tp_size: |
| 232 | assert total_index_heads % tp_size == 0 |
| 233 | else: |
| 234 | assert tp_size % total_index_heads == 0 |
| 235 | self.num_index_heads = max(1, total_index_heads // tp_size) |
| 236 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 237 | assert self.reorder_batch_threshold is not None |
| 238 | self.max_decode_query_len = self.reorder_batch_threshold |
| 239 | |
| 240 | # Stable context-length buffer for decode cudagraph replays. |
| 241 | self.context_len_buffer = torch.empty( |
| 242 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 243 | dtype=torch.int32, |
| 244 | device=device, |
| 245 | ) |
| 246 | # Stable per-token causal page-count buffer for decode cudagraph replays |
| 247 | # (consumed by the MSA top-k path's sparse_topk_select num_valid_pages). |
| 248 | self.num_valid_pages_buffer = torch.empty( |
| 249 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 250 | dtype=torch.int32, |
| 251 | device=device, |
| 252 | ) |
| 253 | |
| 254 | |
| 255 | class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): |
| 256 | """Triton indexer metadata: no SM100 fmha_sm100 plan.""" |
| 257 | |
| 258 | def build( |
| 259 | self, |
| 260 | common_prefix_len: int, |
| 261 | common_attn_metadata: CommonAttentionMetadata, |
| 262 | fast_build: bool = False, |
| 263 | ) -> MiniMaxM3IndexerMetadata: |
| 264 | num_reqs = common_attn_metadata.num_reqs |
| 265 | num_tokens = common_attn_metadata.num_actual_tokens |
| 266 | query_start_loc = common_attn_metadata.query_start_loc |
| 267 | seq_lens = common_attn_metadata.seq_lens |
| 268 | block_table = common_attn_metadata.block_table_tensor |
| 269 | |
| 270 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 271 | split_decodes_and_prefills( |
| 272 | common_attn_metadata, |
| 273 | decode_threshold=self.reorder_batch_threshold, |
| 274 | require_uniform=True, |
| 275 | ) |
| 276 | ) |
| 277 | assert num_decodes + num_prefills == num_reqs |
| 278 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 279 | |
| 280 | # Decode-first batch: context lengths into the stable cudagraph buffer. |
| 281 | context_lens = self.context_len_buffer[:num_reqs] |
| 282 | context_lens.copy_( |
| 283 | common_attn_metadata.compute_num_computed_tokens(), non_blocking=True |
| 284 | ) |
| 285 | |
| 286 | prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None |
| 287 | if num_prefills > 0: |
| 288 | prefill_metadata = MiniMaxM3IndexerPrefillMetadata( |
| 289 | cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to( |
| 290 | torch.int32 |
| 291 | ), |
| 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 | |
| 299 | decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None |
| 300 | if num_decodes > 0: |
| 301 | qsl_cpu = common_attn_metadata.query_start_loc_cpu |
| 302 | query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] |
| 303 | decode_query_len = int(query_lens_cpu[0].item()) |
| 304 | assert decode_query_len > 0 |
| 305 | assert torch.all( |
| 306 | (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0) |
| 307 | ) |
| 308 | assert num_decode_tokens == num_decodes * decode_query_len |
| 309 | decode_metadata = MiniMaxM3IndexerDecodeMetadata( |
| 310 | seq_lens=seq_lens[:num_decodes], |
| 311 | block_table=block_table[:num_decodes], |
| 312 | max_seq_len=common_attn_metadata.max_seq_len, |
| 313 | decode_query_len=decode_query_len, |
| 314 | max_decode_query_len=self.max_decode_query_len, |
| 315 | ) |
| 316 | |
| 317 | return MiniMaxM3IndexerMetadata( |
| 318 | seq_lens=seq_lens, |
| 319 | max_seq_len=common_attn_metadata.max_seq_len, |
| 320 | slot_mapping=common_attn_metadata.slot_mapping, |
| 321 | num_actual_tokens=num_tokens, |
| 322 | num_decodes=num_decodes, |
| 323 | num_decode_tokens=num_decode_tokens, |
| 324 | num_prefills=num_prefills, |
| 325 | num_prefill_tokens=num_prefill_tokens, |
| 326 | prefill=prefill_metadata, |
| 327 | decode=decode_metadata, |
| 328 | ) |
| 329 | |
| 330 | |
| 331 | class MiniMaxM3IndexerImpl(nn.Module): |
| 332 | """Abstract base for the indexer kernel impls. |
| 333 | |
| 334 | Each impl owns its side cache and reports its backend via |
| 335 | ``indexer_backend_cls`` (so each gets its own builder). The Triton and MSA |
| 336 | subclasses each own a full ``forward`` returning ``(decode_topk, |
| 337 | prefill_topk)`` -- no shared forward code. |
| 338 | """ |
| 339 | |
| 340 | # Set by each impl so the side cache reports the matching backend + builder. |
| 341 | indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerBackend |
| 342 | |
| 343 | def __init__( |
| 344 | self, |
| 345 | *, |
| 346 | num_kv_heads: int, |
| 347 | scale: float, |
| 348 | topk_blocks: int, |
| 349 | sparse_block_size: int, |
| 350 | num_index_heads: int, |
| 351 | index_head_dim: int, |
| 352 | prefix: str, |
| 353 | init_blocks: int = 0, |
| 354 | local_blocks: int = 0, |
| 355 | score_type: str = "max", |
| 356 | cache_config: CacheConfig | None = None, |
| 357 | indexer_kv_dtype: IndexerKVDType = "bf16", |
| 358 | topk_indices_buffer: torch.Tensor | None = None, |
| 359 | ) -> None: |
| 360 | super().__init__() |
| 361 | self.num_kv_heads = num_kv_heads |
| 362 | self.scale = scale |
| 363 | self.topk_blocks = topk_blocks |
| 364 | self.block_size = sparse_block_size |
| 365 | self.init_blocks = init_blocks |
| 366 | self.local_blocks = local_blocks |
| 367 | self.score_type = score_type |
| 368 | self.num_index_heads = num_index_heads |
| 369 | self.index_head_dim = index_head_dim |
| 370 | self.indexer_kv_dtype = indexer_kv_dtype |
| 371 | # Shared, stable-address top-k output buffer (set by the model for the |
| 372 | # cudagraph-safe MSA impl); None -> impl allocates fresh (eager). |
| 373 | self.topk_indices_buffer = topk_indices_buffer |
| 374 | topk_completion_counter = None |
| 375 | if current_platform.is_rocm() and topk_indices_buffer is not None: |
| 376 | topk_completion_counter = torch.zeros( |
| 377 | topk_indices_buffer.shape[:2], |
| 378 | dtype=torch.int32, |
| 379 | device=topk_indices_buffer.device, |
| 380 | ) |
| 381 | self.register_buffer( |
| 382 | "topk_completion_counter", |
| 383 | topk_completion_counter, |
| 384 | persistent=False, |
| 385 | ) |
| 386 | # Owns the side cache (registers itself in the static forward context). |
| 387 | self.index_cache = MiniMaxM3IndexerCache( |
| 388 | head_dim=index_head_dim, |
| 389 | prefix=f"{prefix}.index_cache", |
| 390 | cache_config=cache_config, |
| 391 | indexer_kv_dtype=indexer_kv_dtype, |
| 392 | backend_cls=type(self).indexer_backend_cls, |
| 393 | ) |
| 394 | |
| 395 | def forward( |
| 396 | self, |
| 397 | index_query: torch.Tensor, |
| 398 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 399 | """Return ``(decode_topk, prefill_topk)``; implemented per kernel impl.""" |
| 400 | raise NotImplementedError |
| 401 | |
| 402 | |
| 403 | class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): |
| 404 | """Triton indexer score + top-k for both prefill and decode.""" |
| 405 | |
| 406 | def forward( |
| 407 | self, |
| 408 | index_query: torch.Tensor, |
| 409 | *, |
| 410 | attention_block_table: torch.Tensor | None = None, |
| 411 | sparse_block_table_out: torch.Tensor | None = None, |
| 412 | sparse_context_lens_out: torch.Tensor | None = None, |
| 413 | block_page_stride: int | None = None, |
| 414 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 415 | attn_metadata = get_forward_context().attn_metadata |
| 416 | if not isinstance(attn_metadata, dict): |
| 417 | return None, None # profiling run; caches unbound |
| 418 | index_md = attn_metadata[self.index_cache.prefix] |
| 419 | assert isinstance(index_md, MiniMaxM3IndexerMetadata) |
| 420 | num_tokens = index_md.num_actual_tokens |
| 421 | nd = index_md.num_decode_tokens |
| 422 | iq = index_query[:num_tokens].view( |
| 423 | -1, self.num_index_heads, self.index_head_dim |
| 424 | ) |
| 425 | kv = self.index_cache.kv_cache |
| 426 | |
| 427 | # Both sides write into the single shared persistent topk_indices_buffer |
| 428 | # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the |
| 429 | # kernels' out= writes out[:, :total_q]. None -> allocate fresh. |
| 430 | buf = self.topk_indices_buffer |
| 431 | buf_htk = ( |
| 432 | buf if buf is None or current_platform.is_rocm() else buf.transpose(0, 1) |
| 433 | ) |
| 434 | decode_topk: torch.Tensor | None = None |
| 435 | prefill_topk: torch.Tensor | None = None |
| 436 | if index_md.num_decodes > 0: |
| 437 | d = index_md.decode |
| 438 | assert d is not None |
| 439 | fused_sparse_kwargs = {} |
| 440 | if attention_block_table is not None: |
| 441 | fused_sparse_kwargs = { |
| 442 | "attention_block_table": attention_block_table, |
| 443 | "sparse_block_table_out": sparse_block_table_out, |
| 444 | "sparse_context_lens_out": sparse_context_lens_out, |
| 445 | "block_page_stride": block_page_stride, |
| 446 | } |
| 447 | decode_backend_kwargs = {} |
| 448 | if current_platform.is_rocm(): |
| 449 | decode_backend_kwargs["completion_counter"] = ( |
| 450 | self.topk_completion_counter |
| 451 | ) |
| 452 | decode_topk = minimax_m3_index_decode( |
| 453 | iq[:nd], |
| 454 | kv, |
| 455 | d.block_table, |
| 456 | d.seq_lens, |
| 457 | d.max_seq_len, |
| 458 | self.topk_blocks, |
| 459 | self.init_blocks, |
| 460 | self.local_blocks, |
| 461 | self.num_kv_heads, |
| 462 | d.decode_query_len, |
| 463 | d.max_decode_query_len, |
| 464 | out=buf_htk, |
| 465 | **fused_sparse_kwargs, |
| 466 | **decode_backend_kwargs, |
| 467 | ) |
| 468 | if index_md.num_prefills > 0: |
| 469 | p = index_md.prefill |
| 470 | assert p is not None |
| 471 | score = minimax_m3_index_score( |
| 472 | iq[nd:], |
| 473 | kv, |
| 474 | p.block_table, |
| 475 | p.cu_seqlens_q, |
| 476 | p.seq_lens, |
| 477 | p.context_lens, |
| 478 | p.max_query_len, |
| 479 | p.max_seq_len, |
| 480 | self.num_kv_heads, |
| 481 | ) |
| 482 | prefill_topk = minimax_m3_index_topk( |
| 483 | score, |
| 484 | p.cu_seqlens_q, |
| 485 | p.context_lens, |
| 486 | p.max_query_len, |
| 487 | self.topk_blocks, |
| 488 | self.init_blocks, |
| 489 | self.local_blocks, |
| 490 | out=buf_htk[:, nd:, :] if buf_htk is not None else None, |
| 491 | ) |
| 492 | return decode_topk, prefill_topk |
| 493 | |
| 494 | |
| 495 | def select_indexer_impl_cls( |
| 496 | *, |
| 497 | topk_blocks: int, |
| 498 | indexer_kv_dtype: IndexerKVDType = "bf16", |
| 499 | ) -> type[MiniMaxM3IndexerImpl]: |
| 500 | """Pick the indexer impl off the platform, top-k count, and cache dtype. |
| 501 | |
| 502 | On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's |
| 503 | ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is |
| 504 | used for both bf16 and fp8 index caches. Everything else falls back to the |
| 505 | Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8 |
| 506 | support. |
| 507 | """ |
| 508 | if indexer_kv_dtype in ("mxfp4", "nvfp4"): |
| 509 | raise NotImplementedError( |
| 510 | f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) " |
| 511 | "CuteDSL indexer impl." |
| 512 | ) |
| 513 | is_cuda = current_platform.is_cuda() |
| 514 | is_fp8 = indexer_kv_dtype in ("fp8", "fp8_e4m3") |
| 515 | if is_fp8 and not (is_cuda and current_platform.supports_fp8()): |
| 516 | raise NotImplementedError( |
| 517 | f"indexer_kv_dtype={indexer_kv_dtype!r} requires CUDA fp8 support." |
| 518 | ) |
| 519 | is_sm100 = is_cuda and current_platform.is_device_capability_family(100) |
| 520 | use_msa = ( |
| 521 | is_sm100 |
| 522 | and topk_blocks == 16 |
| 523 | and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3") |
| 524 | ) |
| 525 | if use_msa: |
| 526 | # Lazy import so AMD / non-SM100 never import fmha_sm100. |
| 527 | from vllm.models.minimax_m3.nvidia.indexer_msa import ( |
| 528 | MiniMaxM3IndexerMSAImpl, |
| 529 | ) |
| 530 | |
| 531 | logger.info_once( |
| 532 | "MiniMax M3 indexer: selected MSA (fmha_sm100 score + top-k) " |
| 533 | "[topk_blocks=%d, indexer_kv_dtype=%s]", |
| 534 | topk_blocks, |
| 535 | indexer_kv_dtype, |
| 536 | ) |
| 537 | return MiniMaxM3IndexerMSAImpl |
| 538 | if indexer_kv_dtype not in ("bf16", "fp8", "fp8_e4m3"): |
| 539 | raise NotImplementedError( |
| 540 | f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the " |
| 541 | "Triton indexer impl." |
| 542 | ) |
| 543 | logger.info_once( |
| 544 | "MiniMax M3 indexer: selected Triton (no fmha_sm100) " |
| 545 | "[topk_blocks=%d, indexer_kv_dtype=%s, sm100=%s]", |
| 546 | topk_blocks, |
| 547 | indexer_kv_dtype, |
| 548 | is_sm100, |
| 549 | ) |
| 550 | return MiniMaxM3IndexerTritonImpl |
| 551 | |
| 552 | |
| 553 | class MiniMaxM3Indexer(nn.Module): |
| 554 | """Indexer module held by the attention layer (like ``DeepseekV4Indexer``). |
| 555 | |
| 556 | Picks the kernel impl in ``__init__`` (``select_indexer_impl_cls``) and |
| 557 | delegates ``forward``; exposes the impl's side cache via ``index_cache``. |
| 558 | """ |
| 559 | |
| 560 | def __init__( |
| 561 | self, |
| 562 | *, |
| 563 | num_kv_heads: int, |
| 564 | scale: float, |
| 565 | topk_blocks: int, |
| 566 | sparse_block_size: int, |
| 567 | num_index_heads: int, |
| 568 | index_head_dim: int, |
| 569 | prefix: str, |
| 570 | init_blocks: int = 0, |
| 571 | local_blocks: int = 0, |
| 572 | score_type: str = "max", |
| 573 | cache_config: CacheConfig | None = None, |
| 574 | indexer_kv_dtype: IndexerKVDType = "bf16", |
| 575 | topk_indices_buffer: torch.Tensor | None = None, |
| 576 | ) -> None: |
| 577 | super().__init__() |
| 578 | impl_cls = select_indexer_impl_cls( |
| 579 | topk_blocks=topk_blocks, |
| 580 | indexer_kv_dtype=indexer_kv_dtype, |
| 581 | ) |
| 582 | self.impl = impl_cls( |
| 583 | num_kv_heads=num_kv_heads, |
| 584 | scale=scale, |
| 585 | topk_blocks=topk_blocks, |
| 586 | sparse_block_size=sparse_block_size, |
| 587 | num_index_heads=num_index_heads, |
| 588 | index_head_dim=index_head_dim, |
| 589 | prefix=prefix, |
| 590 | init_blocks=init_blocks, |
| 591 | local_blocks=local_blocks, |
| 592 | score_type=score_type, |
| 593 | cache_config=cache_config, |
| 594 | indexer_kv_dtype=indexer_kv_dtype, |
| 595 | topk_indices_buffer=topk_indices_buffer, |
| 596 | ) |
| 597 | |
| 598 | @property |
| 599 | def index_cache(self) -> MiniMaxM3IndexerCache: |
| 600 | return self.impl.index_cache |
| 601 | |
| 602 | @property |
| 603 | def num_index_heads(self) -> int: |
| 604 | return self.impl.num_index_heads |
| 605 | |
| 606 | def forward( |
| 607 | self, |
| 608 | index_query: torch.Tensor, |
| 609 | *, |
| 610 | attention_block_table: torch.Tensor | None = None, |
| 611 | sparse_block_table_out: torch.Tensor | None = None, |
| 612 | sparse_context_lens_out: torch.Tensor | None = None, |
| 613 | block_page_stride: int | None = None, |
| 614 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 615 | if attention_block_table is None and all( |
| 616 | arg is None |
| 617 | for arg in ( |
| 618 | sparse_block_table_out, |
| 619 | sparse_context_lens_out, |
| 620 | block_page_stride, |
| 621 | ) |
| 622 | ): |
| 623 | return self.impl(index_query) |
| 624 | if attention_block_table is None or any( |
| 625 | arg is None |
| 626 | for arg in ( |
| 627 | sparse_block_table_out, |
| 628 | sparse_context_lens_out, |
| 629 | block_page_stride, |
| 630 | ) |
| 631 | ): |
| 632 | raise ValueError( |
| 633 | "MiniMax-M3 fused decode sparse-table arguments must be " |
| 634 | "provided together" |
| 635 | ) |
| 636 | if not current_platform.is_rocm() or not isinstance( |
| 637 | self.impl, MiniMaxM3IndexerTritonImpl |
| 638 | ): |
| 639 | raise ValueError( |
| 640 | "MiniMax-M3 fused decode sparse-table construction is ROCm-only" |
| 641 | ) |
| 642 | return self.impl( |
| 643 | index_query, |
| 644 | attention_block_table=attention_block_table, |
| 645 | sparse_block_table_out=sparse_block_table_out, |
| 646 | sparse_context_lens_out=sparse_context_lens_out, |
| 647 | block_page_stride=block_page_stride, |
| 648 | ) |