[ROCm][MiniMax-M3] Cross-layer lightning-indexer top-k sharing - #47269
Merged
Merged
Conversation
Add opt-in cross-layer sharing of the lightning-indexer top-k block
selection for MiniMax-M3 sparse attention. When `use_index_cache` is set,
only 1 of every `index_topk_freq` sparse-attention layers recomputes the
indexer score + top-k; the remaining layers reuse the selection the
preceding compute layer wrote into the shared `topk_indices_buffer` this
same forward pass. Adjacent sparse layers select nearly identical blocks,
so this cuts the indexer cost ~`freq`x with negligible accuracy impact.
`self.skip_index_topk` is computed once per layer at construction (static),
so it is cudagraph-capture-safe. Disabled by default; enable via
`--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}'`.
Validated on amd/MiniMax-M3-MXFP4, MI355X (gfx950), TP=4, 8k1k:
- GSM8K (lm_eval, 5-shot, full 1319Q): 0.9515 flex / 0.9522 strict
(unchanged vs the recompute-every-layer baseline).
- Mean TPOT: ~-10% at conc=1, ~-4% at high concurrency.
The NVIDIA MiniMax-M3 twin (vllm/models/minimax_m3/nvidia/model.py) can
take the same change as a follow-up.
This change was developed with AI assistance (Claude).
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
tjtanaa
enabled auto-merge (squash)
July 1, 2026 08:53
Fangzhou-Ai
added a commit
to SemiAnalysisAI/InferenceX
that referenced
this pull request
Jul 2, 2026
…om entry Addresses Claude review: the new entry was inserted inside the minimaxm3-fp4-mi355x-atom entry (#1967), dropping its config-keys header and producing duplicate mapping keys. Restore the atom entry and append the vllm entry at the end of the file per repo convention. Also pin the image to the latest nightly containing vllm-project/vllm#47269.
philippesic
pushed a commit
to philippesic/vllm-semantic-cache
that referenced
this pull request
Jul 19, 2026
…project#47269) Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds opt-in cross-layer sharing of the lightning-indexer top-k block selection for MiniMax-M3 sparse attention. When
use_index_cacheis set, only 1 of everyindex_topk_freqsparse-attention layers recomputes the indexer score + top-k selection; the rest reuse the selection the preceding compute layer wrote into the sharedtopk_indices_bufferin the same forward pass.Adjacent sparse layers select nearly the same blocks, so this cuts the indexer score/top-k cost ~
freq× with negligible accuracy impact.self.skip_index_topkis computed once per layer at construction (static → cudagraph-capture-safe).Disabled by default. Enable via:
Why this isn't a duplicate
Searched open PRs for MiniMax-M3 index/top-k/index_topk_freq — no PR adds cross-layer indexer top-k sharing. Related M3 ROCm PRs (#46419 AITER MoE, #46474/#46545 shared-expert fusion) touch different paths.
Test plan / results
amd/MiniMax-M3-MXFP4, MI355X (gfx950), TP=4, 8k1k:index_topk_freq=4).python -m py_compileon the changed file passes; no lines >88 chars.Notes
vllm/models/minimax_m3/amd/model.pyonly. The NVIDIA twin (nvidia/model.py) can take the same change as a follow-up.use_index_cacheis enabled).AI assistance (Claude) was used to develop and validate this change; a human reviewed every line.