karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Model] M3: route Triton indexer top-k through shared topk_indices_buffer

Browse Source 
Unify the Triton indexer impl onto the same persistent top-k buffer as the MSA
impl, and thread the buffer through the AMD model too.

- minimax_m3_index_decode gains an out= param (writes out[:, :total_q]); the merge
  kernel already writes via strides, so a buffer view works.
- MiniMaxM3IndexerTritonImpl.forward writes decode ([:, :nd]) and prefill ([:, nd:])
  into the shared topk_indices_buffer and returns views into it (no fresh per-step
  top-k allocations), matching the MSA impl.
- amd/model.py: allocate the model-level topk_indices_buffer and thread it
  model -> decoder layer -> sparse attention -> indexer (mirrors nvidia). AMD keeps
  its eager break, so this is purely allocation reuse there.

test_msa_indexer_impl_matches_triton now gives each impl its own buffer and asserts
both decode/prefill outputs are views into it. 42/42 pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Yongye Zhu
2026-06-17 05:32:36 +00:00
co-authored by Claude Opus 4.8
parent d2fbaf73c1
commit 06324dd3df
4 changed files with 52 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_minimax_m3.py
+14 -6
View File
@@ -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(
vllm/models/minimax_m3/amd/model.py
+21
View File
@@ -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",
)
vllm/models/minimax_m3/common/indexer.py
+6
View File
@@ -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
vllm/models/minimax_m3/common/ops/index_topk.py
+11 -5
View File
@@ -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