[thanhhao98]

Motivation

The MiniMax-M3 score-only prefill indexer (_flash_attn_fwd_with_block_score_kernel) runs one CTA per (query tile, index head) over the whole prefix, so a short turn over a long context runs on a few CTAs while running decodes wait. In an nsys capture on GB300, a 305-token turn over 80,512 cached tokens ran on 3 CTAs and spent 148 of the 169 ms prefill window in this kernel.

Modifications

• A third grid axis splits the score-only causal K range into contiguous K-tile ranges; no softmax state spans K tiles, so nothing is merged and each score is written once with one split's arithmetic.
• Split count: flash_prefill_with_topk_index aims for about 8 CTAs per SM, capped at the K tiles; the new num_kv_splits argument overrides it.
• Unchanged: the callers, the index-value path and the gfx950 score-only kernel (MiniMax-M3: run the sparse prefill main attention through AITER Gluon paged attention #36546); gfx942 runs the split without AMD CI coverage.
• Prior art: AMD draft M3 perf #39466 (same split for gfx950, never landed).

Accuracy Tests

• New test_minimax_sparse_prefill_topk_idx.py (10 pass on GB300): scores and top-k ids bit-identical to one split (short turns, mixed batch, 16K chunk, fp8 K, lse, 2/7/4096 splits); one case also checks the scores against an fp32 reference.
• The existing MiniMax sparse kernel tests (54) pass.
• GSM8K: MiniMax-M3-NVFP4, TP4, both servers back to back on one GB300 node, sgl-eval 0.1.3 cookbook settings + --max-tokens 32768.

	main	this PR
GSM8K 1319 q	97.19 %	96.82 %
GSM8K 100 q, one request at a time	100 %	100 %

The 5-question gap is inside one binomial SE of the difference between two temperature-1.0 runs (±0.66 points); the block scores are bit-identical to one split.

Speed Tests and Profiling

Score + top-k, bf16, one layer, one TP4 rank, one request, GB300, cold L2 (marker.do_bench, like the new bench_minimax_prefill_index_score.py):

prefix × new	one split ms	split-KV ms	speedup
80,512 × 305	2.172	0.090	24.1×
900,000 × 1,000	24.086	1.927	12.5×
200,000 × 16,384	8.652	7.555	1.15×

End to end: InferenceX AgentX c24 on one 4× GB300 tray, the two arms back to back, both with EAGLE3 #43804 at synthetic acceptance length 2.78, #45451 and fp8 KV. The with-PR arm is this PR's first revision, at 4 CTAs per SM; 8 matches it on short turns and is faster on long chunks.

c24	tok/s/user (1000 / p90 ITL)	total tok/s/GPU	p90 TTFT
without this PR	78.2	22,770	2.01 s
with this PR	94.6	24,134	1.83 s

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests.
• Update documentation (no user-facing change).
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #40295111261
Latest PR Test (Extra): ❌ Run #40295111103
Latest PR Test (AMD ROCm 10): ❌ Run #40295111065

[nvpohanh]

deleted. too nit-picking