diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index ef0958edebc..5d8b9dc0915 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -475,12 +475,16 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch): msa_impl.index_cache.kv_cache = index_cache triton_impl.index_cache.kv_cache = index_cache - # Exercise the shared persistent top-k buffer: the MSA impl must write both - # decode ([:, :nd]) and prefill ([:, nd:]) into it and return views of it. + # Exercise the shared persistent top-k buffer for BOTH impls: each must write + # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views. + # Separate buffers so the two forwards don't clobber each other. nd = sum(q for q in batch.query_lens if q <= 1) msa_impl.topk_indices_buffer = torch.full( (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device ) + triton_impl.topk_indices_buffer = torch.full( + (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device + ) attn_metadata = { msa_impl.index_cache.prefix: msa_builder.build(0, common), @@ -494,10 +498,14 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch): assert msa_prefill is not None and tri_prefill is not None _assert_topk_indices_equal_unordered(msa_decode, tri_decode) _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill) - # decode/prefill outputs are views into the one persistent buffer. - buf = msa_impl.topk_indices_buffer - assert msa_decode.data_ptr() == buf[:, :nd, :].data_ptr() - assert msa_prefill.data_ptr() == buf[:, nd:, :].data_ptr() + # decode/prefill outputs are views into each impl's persistent buffer. + for impl, dec, pre in ( + (msa_impl, msa_decode, msa_prefill), + (triton_impl, tri_decode, tri_prefill), + ): + buf = impl.topk_indices_buffer + assert dec.data_ptr() == buf[:, :nd, :].data_ptr() + assert pre.data_ptr() == buf[:, nd:, :].data_ptr() @pytest.mark.parametrize( diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 4bd92444728..fe40cdd4cf3 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -457,6 +457,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): quant_config: QuantizationConfig | None = None, prefix: str = "", cache_config: CacheConfig | None = None, + topk_indices_buffer: torch.Tensor | None = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -565,6 +566,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): local_blocks=sparse_cfg.get("sparse_local_block", 0), score_type=sparse_cfg.get("sparse_score_type", "max"), cache_config=cache_config, + topk_indices_buffer=topk_indices_buffer, ) # Register the main K/V cache so the KV-cache manager allocates it. @@ -671,6 +673,7 @@ class MiniMaxM3DecoderLayer(nn.Module): quant_config: QuantizationConfig | None = None, force_sparse_attn: bool = False, force_moe: bool = False, + topk_indices_buffer: torch.Tensor | None = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -690,6 +693,7 @@ class MiniMaxM3DecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.self_attn", cache_config=cache_config, + topk_indices_buffer=topk_indices_buffer, ) else: self.self_attn = MiniMaxM3Attention( @@ -771,6 +775,22 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): prefix=f"{prefix}.embed_tokens", ) + # Reserved top-k indices buffer shared by all sparse-attention indexer + # layers (mirrors DeepseekV4); the indexer writes its per-head decode/ + # prefill block selection into it, the attend reads it back. + sparse_cfg = getattr(config, "sparse_attention_config", None) + if sparse_cfg is not None: + tp_size = get_tensor_model_parallel_world_size() + num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size) + self.topk_indices_buffer = torch.empty( + num_index_heads, + vllm_config.scheduler_config.max_num_batched_tokens, + sparse_cfg["sparse_topk_blocks"], + dtype=torch.int32, + ) + else: + self.topk_indices_buffer = None + self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: MiniMaxM3DecoderLayer( @@ -778,6 +798,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): prefix, cache_config=cache_config, quant_config=quant_config, + topk_indices_buffer=self.topk_indices_buffer, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 4e365f856b5..f574ff2d789 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -401,6 +401,10 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): ) kv = self.index_cache.kv_cache + # Both sides write into the single shared persistent topk_indices_buffer + # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the + # kernels' out= writes out[:, :total_q]. None -> allocate fresh. + buf = self.topk_indices_buffer decode_topk: torch.Tensor | None = None prefill_topk: torch.Tensor | None = None if index_md.num_decodes > 0: @@ -418,6 +422,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.num_kv_heads, d.decode_query_len, d.max_decode_query_len, + out=buf, ) if index_md.num_prefills > 0: p = index_md.prefill @@ -441,6 +446,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.topk_blocks, self.init_blocks, self.local_blocks, + out=buf[:, nd:, :] if buf is not None else None, ) return decode_topk, prefill_topk diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index f122d9952bf..79bacd971ad 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -766,10 +766,13 @@ def minimax_m3_index_decode( num_kv_heads: int, decode_query_len: int, max_decode_query_len: int, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Decode index block-score + top-k, both split-K (cudagraph-safe). Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. """ total_q, num_idx_heads, head_dim = idx_q.shape assert num_idx_heads == num_kv_heads, ( @@ -843,11 +846,14 @@ def minimax_m3_index_decode( **score_kwargs, ) - topk_idx = torch.empty( - (num_idx_heads, total_q, topk), - dtype=torch.int32, - device=idx_q.device, - ) + if out is not None: + topk_idx = out[:, :total_q, :] + else: + topk_idx = torch.empty( + (num_idx_heads, total_q, topk), + dtype=torch.int32, + device=idx_q.device, + ) # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts # pow2(num_topk_chunks * pow2(topk)) candidates. TOPK_TARGET_GRID = 64