From 4e5ca89cfe98121642d76b40e32a006f4d0fbf3b Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Wed, 1 Jul 2026 06:50:09 -0400 Subject: [PATCH] [ROCm][MiniMax-M3] Cross-layer lightning-indexer top-k sharing (#47269) Signed-off-by: Fangzhou Ai Co-authored-by: Claude --- vllm/models/minimax_m3/amd/model.py | 42 ++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 7bb8bd722f2..6504ba5d41c 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -135,6 +135,37 @@ def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]: return {i for i, f in enumerate(freq) if f != 0} +def _sparse_attention_layer_ordinals(config: PretrainedConfig) -> dict[int, int]: + """Map each sparse-attention layer id to its ordinal among sparse layers.""" + return { + lid: ordinal + for ordinal, lid in enumerate(sorted(_sparse_attention_layer_ids(config))) + } + + +def _should_skip_index_topk(config: PretrainedConfig, layer_id: int) -> bool: + """ATOM ``index_topk_freq`` (cross-layer index sharing). + + Only 1 of every ``index_topk_freq`` sparse-attention layers recomputes the + lightning-indexer top-k block selection; the rest reuse the selection the + preceding compute layer wrote into the shared ``topk_indices_buffer`` this + same forward pass. This cuts the indexer score + top-k cost ~``freq``x with + negligible accuracy impact (adjacent sparse layers pick nearly the same + blocks; ATOM validated GSM8K with freq=4). Gated by ``use_index_cache``; + enable via ``--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}'``. + """ + if not getattr(config, "use_index_cache", False): + return False + freq = int(getattr(config, "index_topk_freq", 1) or 1) + if freq <= 1: + return False + ordinal = _sparse_attention_layer_ordinals(config).get(layer_id) + if ordinal is None: + return False + offset = int(getattr(config, "index_skip_topk_offset", 0) or 0) + return max(ordinal - offset, 0) % freq != 0 + + def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool: """Whether this layer's MLP is a sparse MoE block (vs a dense MLP).""" moe_layer_freq = getattr(config, "moe_layer_freq", None) @@ -541,6 +572,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 + # Cross-layer index sharing (ATOM index_topk_freq): when True this sparse + # layer reuses the previous compute layer's top-k block selection from the + # shared topk_indices_buffer instead of recomputing it. Static per layer + # -> cudagraph-capture-safe. + self.skip_index_topk = _should_skip_index_topk(config, layer_id) + # Sparse "index" branch dims. index_q has the same head count as the KV # heads (sparse_num_index_heads == num_key_value_heads), so it shards # identically -- including replication when tp_size > num_key_value_heads. @@ -728,7 +765,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): # Single eager break around both: their split-K kernels read per-request # metadata and can't be captured into a cudagraph. The indexer writes its # top-k into the shared ``topk_indices_buffer``; the attend reads it back. - self.indexer(index_query) + # When skip_index_topk is set (ATOM index_topk_freq), reuse the selection + # the preceding compute layer wrote into the shared buffer this forward. + if not self.skip_index_topk: + self.indexer(index_query) return self.impl.forward(self, query, self.kv_cache, output)