[AMD] Add opt-in MiniMax-M3 TP4 indexer context partitioning

[ThomasNing]

Motivation

Under TP4, MiniMax-M3's four index-query heads are split across ranks while the index-K cache is replicated. Each rank scans the full index cache for its head. This PR adds an opt-in decode path that partitions the block reads across those four ranks and scores all heads while each K tile is resident.

The measured benefit is in the indexer chain, including communication: at batch 32 / 128K context, latency falls from 109.84 to 41.63 µs (62.1% lower latency, 2.64× speedup). Short contexts and small batches regress, so the feature defaults off. Full-model serving gains have not been measured.

Modifications

• Add SGLANG_MINIMAX_M3_INDEXER_CP=1, default 0, for gfx950 TP4 ordinary decode with attention CP/DP size 1, four index heads, four main KV heads, dimension 128, 128-token blocks, max scoring, and top-k 16. Unsupported run configurations log a reason and retain the existing path.
• Gather existing post-RoPE queries through the attention TP communicator, then let rank r read blocks r, r+4, r+8, ... for all four heads. The cache remains replicated; this reduces read traffic, not cache capacity.
• Use Triton kernels for shard scoring, local top-16 selection per head, and final merging. Keeping 16 candidates on every shard preserves global top-16 even when all winners belong to one shard. Initial/recent block priorities and native ROCm score/ID ordering are preserved.
• Exchange packed score/ID keys through SGLang's communicator, using AITER custom all-gather when supported and the existing fallback otherwise. Keys use FP32 bitcast views for copying, without numeric conversion. Warm both collective paths before graph capture; merge directly from the received layout.
• Keep the implementation in two new modules with small backend/dispatch changes. Existing top-k reuse is checked first. Prefill, index-value layers, model projections, checkpoint loading, and main sparse attention retain their existing paths.
• Add a manual four-rank correctness/timing harness, reproduction instructions, and all raw timing samples under test/manual/minimax_m3/indexer_cp/.

The initial scope excludes speculation, TBO, HiSparse, FP8 queries, and dense sparse decode. There is no automatic batch/context crossover policy. It does not require a main-attention backend change.

Accuracy Tests

Completed on four MI355X / gfx950 GPUs:

• 8/8 cases passed on all ranks: BF16 and FP8 E4M3 index caches × random scores, ties, winners concentrated on one shard, and mixed lengths.
• Exact selected-ID comparisons against native SGLang and an independent FP32 PyTorch reference, including empty rows and partial blocks.
• Changed sequence lengths after graph capture, including short, partial-block, and empty rows; replay matched both references.
• All 17 timing shapes additionally passed native-versus-CP selected-ID parity in eager execution and graph replay.
• Repository pre-commit hooks passed for all changed files. The scorer, runtime helper, and benchmark script are byte-identical to the files used for the successful GPU run.

Validation still outstanding: runtime feature-gate/backend dispatch through a loaded server, full-model accuracy, serving TPS/TTFT/TPOT, and measurements with the fallback collective. The harness directly invokes the indexer helper. The shape gate also accepts FP16 queries and other FP8 cache variants, which were not covered by this run.

Speed Tests and Profiling

Four MI355X GPUs, BF16 queries, FP8 E4M3 index cache, AITER custom gather available. Both gathers are included; projections and main attention are excluded. Eight indexer calls per HIP graph, 100 replays per round, seven rounds alternating baseline/CP order. Each sample uses the slowest rank; reported values are the median of seven samples.

Inputs/cache are repeatedly reused and warm, without an L2 flush or rotating-layer protocol. These are microbenchmark results and cannot be interpreted as model throughput gains.

Context	Batch	Native TP (µs)	CP (µs)	CP latency change
8K	1	9.40	15.40	+63.8%
8K	8	7.72	18.32	+137.3%
8K	32	12.12	19.58	+61.5%
32K	16	20.80	20.04	-3.7%
32K	32	30.41	24.38	-19.8%
128K	1	12.71	15.93	+25.4%
128K	8	32.85	24.02	-26.9%
128K	16	51.71	32.08	-38.0%
128K	32	109.84	41.63	-62.1%
One 128K row, remaining 1K	16	40.76	27.82	-31.7%
One 128K row, remaining 1K	32	68.89	34.62	-49.7%

Negative latency change is an improvement. All tested 8K shapes and batches 1–2 at 128K regress. The small 32K/batch-16 gain needs independent repetition. The README contains the complete 17-shape table, runtime versions, source hashes, and the reproduction command; results-gfx950.json contains every round.

Checklist

• Format changed files with repository pre-commit hooks.
• Add reproducible four-rank correctness and speed checks (manual GPU harness).
• Document enablement, supported configurations, results, and limitations.
• Provide indexer-level accuracy and speed measurements, including regressions.
• Validate full-model accuracy and serving performance with CP off/on.
• Integrate the multi-GPU checks into registered CI.

---

CI States

Latest PR Test (Base): ✅ Run #36951043442
Latest PR Test (Extra): 🚫 Run #36962666288
Latest PR Test (AMD ROCm 10): ❌ Run #36951042978

[kevin-mii]

/tag-and-rerun-ci

[bingxche]

Hi @ThomasNing , thanks for your contribution!

The test you added test_minimax_indexer_cp did not pass even in this PR checks. https://github.com/sgl-project/sglang/actions/runs/36951042978/job/110670268955#step:6:26861

And it's still failing as of today https://github.com/sgl-project/sglang/actions/runs/37709475489/job/113091796233#step:6:7682. Could you please take a look?

cc @michaelzhang-ai @HaiShaw @hnyls2002