MiniMax-M3: share the sparse index top-k across layers and reuse the decode top-k buffer

[zcnrex]

MiniMax-M3: share the sparse index top-k across layers and reuse the decode top-k buffer

MiniMax-M3's sparse attention runs a lightning-indexer pass in every sparse layer
to choose the top-k KV blocks that the main attention then reads. On the
value-disabled sparse layers that indexer has no output of its own — idx_o is
None there, only the block selection is consumed — so its result can be
computed once and shared by a run of consecutive layers. This PR adds
SGLANG_MINIMAX_M3_INDEX_TOPK_FREQ=N (default 2), which groups every N such
layers: the group's first layer computes and stores the reduced top-k, and the
remaining layers reuse it, skipping both the flash-index attention and the
group-reduce. In decode the selection is written straight into a persistent
per-batch-size int32 buffer (topk_out=) rather than a freshly allocated
tensor per layer, which removes an allocation from every sparse decode layer and
keeps the destination address fixed so CUDA-graph capture never allocates. The
buffer is only useful because of the sharing, so the two changes ship together.

Sharing the selection changes which KV blocks the skip layers attend to — see
## Correctness below. It is therefore gated: index_topk_freq is pinned to 1
on every non-ROCm platform and whenever two-batch overlap is enabled, and at
freq 1 no buffer is allocated and every new argument is None/False, so the
default CUDA path is byte-for-byte main's.

Two smaller changes in the same files ride along:

• The layer-invariant prefill seqblock metadata (cu_seqblocks_q,
  max_seqblock_q, all_seqblock_q) is computed once per ForwardBatch in the
  backend and passed down, instead of being recomputed inside each sparse layer.
  Same values. The owning ForwardBatch is part of the cache key because a
  single metadata init can be followed by more than one ForwardBatch reaching
  the layers (two-batch overlap splits into two children with different
  extend_seq_lens), so a hit requires the same object, not just a live cache.
• forward_decode gets the _is_sparse_kv_cached_by_fusion guard that the
  extend path already has on main, so the decode sparse-KV store is not repeated
  when the fused qk-norm-rope kernel already wrote it. Idempotent write either
  way; this only removes the launch.

The decode score kernel's autotune space now drops configs whose token tile is
smaller than a sparse block (BLOCK_SIZE_N < block_size), and on ROCm the sweep
is replaced by one fixed config plus a BLOCK_SIZE_N heuristic, because runtime
autotune can fire under CUDA-graph capture and mistune there.

Performance

Measured on 8x MI350X (gfx950), MiniMax-M3-MXFP8, TP8, 80k input / 600 output.
Two bench reps per boot; the warm (second) rep is reported and the two agreed
within 0.1%. Noise floor is +/-0.4%. Tables are the harness output verbatim,
trimmed to the first five columns. Baseline is upstream/main.
Baseline (upstream/main)

Input lens: [80000]. Output lens: [600]. Cache hit rate: 90.0%.
|   batch size |   input len |   latency (s) |   input throughput (tok/s) |   output throughput (tok/s) |
|--------------|-------------|---------------|----------------------------|-----------------------------|
|            1 |       80000 |         11.26 |                     187984 |                       55.35 |
|            2 |       80000 |         11.88 |                     191222 |                      108.68 |
|            4 |       80000 |         13.39 |                     192965 |                      204.62 |
|            8 |       80000 |         16.27 |                     193687 |                      370.13 |
|           16 |       80000 |         22.2  |                     193188 |                      616.28 |
|           32 |       80000 |         33.78 |                     193245 |                      934.92 |

With this PR

Input lens: [80000]. Output lens: [600]. Cache hit rate: 90.0%.
|   batch size |   input len |   latency (s) |   input throughput (tok/s) |   output throughput (tok/s) |
|--------------|-------------|---------------|----------------------------|-----------------------------|
|            1 |       80000 |         10.44 |                     209588 |                       59.67 |
|            2 |       80000 |         11.14 |                     210417 |                      115.61 |
|            4 |       80000 |         12.35 |                     213259 |                      221.18 |
|            8 |       80000 |         14.41 |                     215085 |                      419.8  |
|           16 |       80000 |         19.37 |                     214509 |                      716.39 |
|           32 |       80000 |         28.94 |                     215010 |                     1127.4  |

Output throughput +6.38%..+20.59%; input throughput +10.04%..+11.49%.

Accuracy

gsm8k, 512 examples, max_tokens 2048, temperature 0, seed 0, against the
same server boot as the perf run.
Baseline

== gsm8k ==
512 examples (single-shot)  |  109.0s  |  1338 tok/s  |  146K tokens

* score           =  97.07%
  stop_rate       =  98.44%
  truncated_rate  =  1.56%  [warn: hitting max_tokens]
  error_rate      =  0.00%

With this PR

== gsm8k ==
512 examples (single-shot)  |  97.7s  |  1463 tok/s  |  143K tokens

* score           =  96.68%
  stop_rate       =  98.44%
  truncated_rate  =  1.56%  [warn: hitting max_tokens]
  error_rate      =  0.00%

Note on the launch command

--moe-runner-backend aiter is not honoured on current main for mxfp8. Arg resolution logs:

mxfp8 quantization supports only cutlass, deep_gemm, flashinfer_trtllm,
flashinfer_trtllm_routed, triton backends. Overriding 'aiter'.

and falls back to triton. Both sides of this A/B therefore ran the Triton MoE runner, so the comparison is apples-to-apples and the deltas stand — but AITER_CONFIG_FMOE is inert under this configuration and the aiter MoE path is not what was measured. The flag is kept in the command above only because it matches the invocation actually used; anyone reproducing will get triton and should see the same numbers.

---

CI States

Latest PR Test (Base): ⏳ Run #34259568986
Latest PR Test (Extra): ❌ Run #34259568651
Latest PR Test (AMD ROCm 7.2): ❌ Run #34259568790