[Godmook]

pr description is generated by AI tool

Motivation

minimax_decode_topk has three selection regimes. Two of them already break exact score ties by block id (the lower id wins):

• the small path (num_blocks <= 128): is_greater(...), "lower block id wins"
• the ROCm path: a packed (score, ~id) total order

The CUDA radix paths (num_blocks > 128, i.e. contexts above 16k tokens at block size 128) place the final-round ties in atomicAdd order instead. When tied blocks straddle the top-k boundary, the selected set changes from run to run. ROCm indexer CP already depends on a fixed tie order ("requires ... the ROCm radix top-k tie ordering"); the CUDA paths gave no such guarantee.

Exact ties are common for indexers whose keys carry no RoPE: the block score is the max over units, so repeated text produces equal scores. On such a model, serving the same prompts twice from a cold cache with the radix cache on gave 43/45 identical greedy outputs; with this change, 45/45.

Modifications

• Final radix round: rank the tied candidates by block id (warp ballot + prefix sum) and keep the lowest ids, matching the small and ROCm paths. Applied to both the single-element path and the multi-element (register bucket) path.
• Fast path: when every tie is selected (equal_count <= topk_remain), the order does not matter (the output is sorted ascending), so the ties are written with a plain atomic as before and the ranking is skipped.
• Removes the now unused counter_final.
• Adds test_decode_topk_ties_pick_lowest_ids, which checks the selected set against (score desc, id asc) for 300 to 9,000 blocks with heavy ties.

ROCm is unchanged.

Accuracy Tests

MiniMax-M3 (bf16, TP8, H200, radix cache on), main vs this PR:

• Greedy outputs on 40 MMLU-Pro prompts plus needle-in-text prompts up to 60k tokens (radix path): 45/45 identical.
• GPQA-Diamond (n=1, max_tokens 32k): 89.4 (main) vs 90.4 (this PR). (range in noise)

Speed Tests and Profiling

MiniMax-M3, TP8, H200, one request at a time, radix cache flushed before each request:

input tokens	prefill TTFT, main / PR	decode TPOT, main / PR
4k	0.395 / 0.396 s	13.42 / 13.11 ms
16k	1.496 / 1.499 s	13.62 / 13.26 ms
32k	3.042 / 3.051 s	13.68 / 13.46 ms
48k	4.639 / 4.643 s	13.76 / 13.54 ms
60k	5.936 / 5.951 s	13.81 / 13.59 ms

Kernel microbenchmark (2 heads, 8 to 64 rows, 2k to 8k blocks): within ±0.2 us of main for distinct scores; equal or faster with ties.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (No API or behavior change to document.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #41007060631
Latest PR Test (Extra): ❌ Run #41007060650
Latest PR Test (AMD ROCm 10): ❌ Run #41007060555

[Godmook]

@DarkSharpness @BBuf Hi. Could you kindly review this PR? Thanks!

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py

Targeted validation for a9209eed68b5.

[github-actions]

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py

[DarkSharpness]

Why do we need deterministic here? This may introduce non-trivial overhead. The result is already correct.

[Godmook]

Results are valid either way, but not stable : tied blocks get picked in atomicAdd order, so the same prompt can select different blocks across runs. The small and ROCm paths already break ties by lower id, so this just makes CUDA radix match.

[Godmook]

@DarkSharpness FYI, The ranking only runs when ties straddle the top-k boundary, so overhead is negligible. Microbench is within ±0.2 us of main, and E2E shows no regression. In practice, serving the same prompts twice from a cold cache gave 43/45 identical greedy outputs on main and 45/45 with this PR.

Feel free to give feedback!!

[BBuf]

LGTM once related CI passes; non-blocking: add a constructed cross-slot tie case.

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py

Targeted validation for a9209eed68b5.

[github-actions]

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py