[ROCm][MiniMax-M3] Cross-layer lightning-indexer top-k sharing (#47269)
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Fangzhou Ai committed
Jul 1, 2026 at 06:50 UTC
4e5ca89cfe98121642d76b40e32a006f4d0fbf3b
1 file changed
+41
-1
vllm/models/minimax_m3/amd/model.py
+41
-1
@@ -135,6 +135,37 @@ def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
135
return {i for i, f in enumerate(freq) if f != 0}
136
137
138
+def _sparse_attention_layer_ordinals(config: PretrainedConfig) -> dict[int, int]:
139
+ """Map each sparse-attention layer id to its ordinal among sparse layers."""
140
+ return {
141
+ lid: ordinal
142
+ for ordinal, lid in enumerate(sorted(_sparse_attention_layer_ids(config)))
143
+ }
144
+
145
+
146
+def _should_skip_index_topk(config: PretrainedConfig, layer_id: int) -> bool:
147
+ """ATOM ``index_topk_freq`` (cross-layer index sharing).
148
+
149
+ Only 1 of every ``index_topk_freq`` sparse-attention layers recomputes the
150
+ lightning-indexer top-k block selection; the rest reuse the selection the
151
+ preceding compute layer wrote into the shared ``topk_indices_buffer`` this
152
+ same forward pass. This cuts the indexer score + top-k cost ~``freq``x with
153
+ negligible accuracy impact (adjacent sparse layers pick nearly the same
154
+ blocks; ATOM validated GSM8K with freq=4). Gated by ``use_index_cache``;
155
+ enable via ``--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}'``.
156
+ """
157
+ if not getattr(config, "use_index_cache", False):
158
+ return False
159
+ freq = int(getattr(config, "index_topk_freq", 1) or 1)
160
+ if freq <= 1:
161
+ return False
162
+ ordinal = _sparse_attention_layer_ordinals(config).get(layer_id)
163
+ if ordinal is None:
164
+ return False
165
+ offset = int(getattr(config, "index_skip_topk_offset", 0) or 0)
166
+ return max(ordinal - offset, 0) % freq != 0
167
+
168
+
169
def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
170
"""Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
171
moe_layer_freq = getattr(config, "moe_layer_freq", None)
@@ -541,6 +572,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
572
self.kv_size = self.num_kv_heads * self.head_dim
573
self.scaling = self.head_dim**-0.5
574
575
+ # Cross-layer index sharing (ATOM index_topk_freq): when True this sparse
576
+ # layer reuses the previous compute layer's top-k block selection from the
577
+ # shared topk_indices_buffer instead of recomputing it. Static per layer
578
+ # -> cudagraph-capture-safe.
579
+ self.skip_index_topk = _should_skip_index_topk(config, layer_id)
580
+
581
# Sparse "index" branch dims. index_q has the same head count as the KV
582
# heads (sparse_num_index_heads == num_key_value_heads), so it shards
583
# identically -- including replication when tp_size > num_key_value_heads.
@@ -728,7 +765,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
765
# Single eager break around both: their split-K kernels read per-request
766
# metadata and can't be captured into a cudagraph. The indexer writes its
767
# top-k into the shared ``topk_indices_buffer``; the attend reads it back.
731
- self.indexer(index_query)
768
+ # When skip_index_topk is set (ATOM index_topk_freq), reuse the selection
769
+ # the preceding compute layer wrote into the shared buffer this forward.
770
+ if not self.skip_index_topk:
771
+ self.indexer(index_query)
772
return self.impl.forward(self, query, self.kv_cache, output)
773
774