[XiaobingSuper]

Part of https://github.com/wuhuikx/LLM-Booster/issues/341.

Summary

Optimize the MiniMax-M3 prefill index pipeline on gfx1250:

• derive score-kernel geometry from query work, context blocks, planes, and runtime CU count;
• use Q64 grouped page stores or Q256 page reuse only where their measured occupancy/re-read trade-off wins;
• split wide score rows into parallel 1024-block local top-k passes when the emitted page table has enough workspace, then merge the exact local lists;
• dispatch heavy prefill rows to aiter's one-pass small-k selector by rows * width / CU;
• preserve main's decode selector crossover, non-gfx1250 behavior, non-FP8 score path, and underfilled fallback;
• preserve M3.1's pages_per_block contract: page-128/MXFP4 output cannot be reused as parallel-selector workspace and falls back safely.

The remaining 28-page multi-plane value is a measured kernel loop tile, not an activation shape window. Activation is derived from query tiles, planes, context work, and CU count; the fixed tile controls work inside an already-selected specialization.

Dependency

The wave32 selector support landed in ROCm/aiter#7878, merge commit 5c1489c242fc50e508ae7e9f5f5b4027a7da0366.

Current-main A/B/A/B

Revisions and environment:

• ATOM base: 74e389b7661fb53ffb0a7daa93b0d73afe4d7476
• measured ATOM head: 58177a3aa716de417820aa4bc6243650dd4b17d8
• review-cleanup head: 2b446fcf0804cb912409acc7b8e08398d5ec5ef2
• aiter dependency: #7878 ce3681929fe19dd2c0e34bc392388fcf331a2496
• MI455X/gfx1250, 256 CUs, e15-09 GPU1
• image: rocm/fw-bringup:gfx1250-minimax-m3-20261110
• complete operator path: score + ordered top-k + sparse-table emit
• five warmups, five batches x 20 iterations; median CUDA-event time per run
• process order: base / head / base / head

Command inside the container:

export PYTHONPATH=/app/ATOM:/app/aiter:$PYTHONPATH
python3 /app/bench_current_vs_main.py
CASESET=acceptance python3 /app/bench_current_vs_main.py

The harness and every raw log are in this reproducible benchmark artifact.

Results:

• 22 production trace shapes:
  • base runs: 1919.612, 1902.898 us
  • head runs: 1542.718, 1551.615 us
  • median aggregate: 1911.255 -> 1547.167 us (1.235x, -19.0%)
• issue refactor: extract CI regression-rerun into models.json + regression_rerun.py #341 complete-path anchor (q=2517, context 105301):
  • 82.693 -> 55.080 us (1.501x, -33.4%)
• DPEP four-head q2K/q4K/q8K/q16K/q32K:
  • 22.5% / 24.1% / 27.9% / 23.5% / 29.5% lower complete-path time
  • aggregate 2927.612 -> 2128.982 us (1.375x, -27.3%)
• representative long trace shapes:
  • q2858 / 3942 blocks: 190.868 -> 166.689 us
  • q4104 / 1895 blocks: 206.085 -> 114.047 us
  • q5459 / 1682 blocks: 228.375 -> 139.802 us
  • q2331 / 2706 blocks: 166.776 -> 102.313 us

The maintainer-requested cleanup was measured separately, A/B/A/B against the
measured head. The complete 22-shape aggregate is
1559.166 -> 1535.917 us (1.015x); issue-anchor + DPEP aggregate is
2131.666 -> 2129.550 us (flat). The raw cleanup logs are in the same artifact.

The short/underfilled shapes retain main's original score-kernel specialization. Their complete-path times remain around the launch floor; the aggregate numbers above include them rather than filtering them out.

Production routing

• Native ATOM and the vLLM plugin pass emit_sparse_block_table=True; they can use the parallel selector when its workload and workspace gates pass.
• SGLang does not request fused emit, so it cannot take the parallel-selector branch; it still receives the score-kernel and heavy aiter-selector improvements.
• Decode keeps main's width-based selector gate instead of the new prefill workload/CU crossover.
• M3.1 page-128/MXFP4 output has pages_per_block=1; the capacity-derived policy keeps it off the aliased parallel workspace while both aiter and Triton emit paths preserve page geometry.

Validation

• gfx1250 suite with the merged #7878 wave32 kernel: 288 passed
• 42,120 scorer launch-policy shapes preserve the resolved launch tuple exactly
• score, ordered topk_idx, sparse_bt, and sparse_ctx parity, including ties and strided queries
• explicit page-128/MXFP4 prefill test: aiter and fallback outputs agree, and the parallel workspace path is refused
• decode dispatch boundary tests preserve main's 1024/2048 non-emit/emit crossovers
• runtime dispatch consistently uses get_gfx_runtime()
• Black, Ruff, and git diff --check: clean

Not validated in this revision: full-model TTFT/throughput and a post-rebase GSM8K run. This PR's performance claim is limited to the complete production index operator; the output regression tests are bit-exact.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2971 --add-label <label>

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[valar-cop]

a few small things..

ci green.

[valar-cop]

these magic numbers lost their derivation, please leave a comment so we know where 22/3 comes from..

[XiaobingSuper]

Fixed in e82c45d. The shape ranges and fixed chunk/stage values are gone; the launch is now derived from runtime CU count, query planes, and context blocks, with comments for the occupancy thresholds.

[valar-cop]

the unmasked load lost its justification, please restore the comment or add mask=?

[XiaobingSuper]

Fixed in e82c45d. Restored the justification immediately above the load: each page id names a fully allocated 128-row cache page, and the later causal mask removes positions beyond the live sequence.

[valar-cop]

the chunk window has three owners now, can we keep one table keyed on (arch, is_fp8, max_query_len, ...) instead of re-deriving it?

[XiaobingSuper]

Fixed in e82c45d. _prefill_score_config is now the single gfx1250 FP8 policy owner; it returns the complete tile/chunk/warp/stage/store configuration from (arch, dtype class, q, blocks, batch, heads, CUs). The original generic chunk policy is only the unchanged fallback.

[valar-cop]

this would still pass if the tuned launch stopped firing, please assert the selected policy and the new chunk=24 branch..

[XiaobingSuper]

Fixed in e82c45d. The policy test now asserts every selected launch tuple, including fallback boundaries and the 16-page cadence adjustment; GPU tests compare scores and all emitted outputs against the original path. Final runs: CI simulation 215 passed / 1 skipped, gfx1250 191 passed / 1 skipped.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[XiaobingSuper]

GPU0 follow-up at b43da855 is complete.

• e15-09 gfx1250, previous head f3a7918 vs candidate, production n_valid_column_per_row, 5 alternating A/B/A rounds.
• 22 trace shapes complete path: 1728.614 -> 1601.123 us (1.080x); score-only: 1146.220 -> 979.835 us (1.170x).
• DPEP 4-head q2K/q4K/q8K complete path improves 16.5% / 6.4% / 7.7%; q16K/q32K and the original q2517 anchor are flat.
• 8 boundary cases: 1399.933 -> 1371.737 us aggregate; all exact.
• GPU regression: 268 passed, 1 optional AITER-selector skip. Black, Ruff, py_compile, and diff check pass.

The policy is now derived from query tiles, batch/heads, context-block waves, and runtime CU count; no model name or trace-shape table. An experimental direct port of the MI355 FlyDSL scorer was rejected: 83.397 us vs 69.282 us Triton on the anchor, with non-bit-exact scores and changed ordered top-k.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[XiaobingSuper]

Wave32 selector follow-up pushed in 44efc8b4, stacked on ROCm/aiter#7878.

This is the selector hotspot found by kernel profiling, not another scorer tile change. The gfx1250 policy uses rows * width / runtime_CU and keeps short calls on the fused Triton selector.

• 22 production shapes complete path: 1594.627 -> 1526.337 us (1.045x)
• long trace shapes: 10.6-11.4% lower
• DPEP q2K-q32K: 8.5-13.5% lower
• selector-only: 2.13-4.89x on representative shapes
• ordered top-k and sparse outputs exact; q32K repeated 20 times with zero mismatches
• ATOM GPU tests: 276 passed; aiter distribution/ragged matrix, Black and Ruff clean

The direct SGLang-style fused score+streaming-topk prototype was also measured and rejected (85-153 us on the anchor versus ~52 us current complete path).

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[valarLip]

⚠️ Action required before merge (5)

1. Drop @functools.cache from _aiter_selector_wins. index_topk.py:894 — the key now includes rows = heads * total_q, so every distinct prefill size adds a permanent entry on every arch (gfx950 too, where rows doesn't even affect the result). CPU probe: 40k distinct steps -> 40k entries, maxsize=None. The body is two integer compares; just remove the cache (bounding it at 128 would only churn) and make rows/compute_units required. lru_cache on parallel_prefill_selector_wins (prefill_policy.py:62) buys nothing either.
2. Make prefill_score_config the single owner of the scorer launch (valar-cop's "three owners" ask is half done). Warps/stages/chunk still come from three places merged with or (index_topk.py:1717,1728-1729: _prefill_score_launch_config + _score_chunk_blocks + the policy). Have the policy return a complete config (no None sentinels), fold the gfx1250 FP8 (4, 1) default into it, and delete the or fallbacks.
3. Remove the duplicated fallback page body. index_topk.py:533-563 re-inlines exactly what _index_block_score_page(..., POST_REDUCE_SCALE=False) computes; Triton inlines @triton.jit helpers, so call the helper (or show the ISA/timing difference that justifies keeping two copies).
4. Delete the try: _small_k_selector() except ImportError: skip in test_gfx1250_wave32_selector_matches_prefill_fallback. aiter#7878 is merged and CI installs aiter main; the import is mandatory (index_topk.py:62), so the skip only hides a broken selector (same failure mode as #3162).
5. Add a multi-request parity case that actually reaches the parallel selector. None of the GPU cases with batch > 1 is wide enough ([1703,717] ~958 blocks, [640]*4 ~257) to pass parallel_prefill_selector_wins, so the partial/merge kernels' pid_b/row_starts indexing is only covered for batch 1. E.g. [2858, 2331] with ~3900-block prefixes.

Smaller: the prefill-warps rationale is now orphaned above DECODE_TOPK_* (index_topk.py:758-764); the "short/underfilled" comment at :534 is wrong (that branch is every non-gfx1250 / non-FP8 launch); "GPU0 complete-path A/B" numbers at :884 lost their log pointer; rows means heads*total_q in one gate and total_q in the other — rename. Perf numbers in the description are fine; please run ci:atom/ci:vllm MiniMax-M3 accuracy before merge (description says GSM8K wasn't rerun).

---

_{Reviewed at 58177a3. valar-cop items already fixed: unmasked-load justification restored (:441, :541), policy test now asserts the selected launch. Not stale vs main (no minimax_m3 changes since base 74e389b, clean merge); aiter#7878 merged.}

[valar-cop]

ci triage for 2b446fc:

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43698338926

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43698338887

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43698338963

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[XiaobingSuper]

Addressed all five review items in 2b446fcf:

1. Removed both selector-policy caches; _aiter_selector_wins now requires explicit arch / flattened selector rows / CU count.
2. prefill_score_config now returns the complete launch (Q tile, chunk, warps, stages and kernel specialization); the launcher has no None sentinels or or fallbacks, and the separate scorer launch/chunk owners are gone.
3. Removed the duplicated fallback page body; every STORE_GROUP=1 launch now calls _index_block_score_page.
4. Removed the optional aiter import skip now that aiter#7878 is merged.
5. Added [2858, 2331] / ~3942-block batch-2 coverage and assert that it actually launches the partial/merge parallel selector.

Also moved the prefill warp rationale to its policy constant, restored an accessible raw-log pointer, and distinguished flattened selector rows from query rows.

Validation on e15-09 GPU0 (gfx1250): 288 passed; Black/Ruff clean. A 42,120-shape host probe preserves every resolved scorer launch tuple. Review-cleanup A/B/A/B is flat-to-better: 22-shape aggregate 1559.166 -> 1535.917 us; issue anchor + DPEP 2131.666 -> 2129.550 us. Raw logs are in the PR's linked benchmark artifact. ci:atom and ci:vllm labels are set for the requested accuracy runs.

@valarLip ready for re-review.

[valarLip]

⚠️ Action required before merge (3)

All five items from the last round are fixed (details below). What's left is small and doesn't block:

1. Fix the decode-selector comment that now contradicts itself. index_topk.py:708-709 (new) says "the best single pair over the deployed shapes is (512, 8)", but :700-701 says the best single pair is (256, 4) and (512, 8) is only 0.6% behind it. Change it to "(512, 8) is the chosen pair; the warps are needed because decode's grid is 50-200 CTAs".
2. Assert that the batch-2 case actually takes the parallel selector. tests/test_minimax_m3_index_topk.py:957-968 checks bool(parallel_launches) is expect_parallel, and expect_parallel comes from the gate under test. If that gate regresses, the [2858, 2331] case silently stops covering the partial/merge pid_b/row_starts indexing. A CPU probe shows it reaches the path at 128/256/304 CUs with heads=1 (it sits exactly on the workspace bound, 4x16 = 64 keys). Add assert parallel_launches for qlens == [2858, 2331] and heads == 1.
3. Use block units in the policy comments. prefill_policy.py:25-27,209 say "page chunk" and "28-page chunks", but chunk_blocks counts 128-token blocks, and in this module a "page" is a 16-token physical page (8 per block). Say "block chunk" and "28-block chunks".

Optional: prefill_score_config's lru_cache(maxsize=128) (prefill_policy.py:120) is bounded now, but it is keyed on per-step max_query_len/max_block, so eager prefill mostly misses. A CPU probe with 5000 distinct q gave 0 hits. The body is integer arithmetic, so dropping the cache costs nothing.

CI at head: ATOM Test passed, including MiniMax-M3-MXFP4 accuracy. ATOM SGLang Test passed. The red vLLM/ATOM gate entries are runs that were cancelled and superseded. In ATOM vLLM Test (run 43680334394), MiniMax-M3-MXFP4 EAGLE3 TP4 passed and 3 non-MiniMax jobs are still pending. Merge once those are green.

---

_{Re-reviewed at 2b446fc. Already fixed: _aiter_selector_wins cache removed and its kwargs required (probe re-run: no cache left on either selector gate); prefill_score_config returns the complete launch with no None/or fallbacks, and the gfx1250 FP8 (4,1) default is folded into _default_prefill_score_config; the duplicated fallback body is gone (every STORE_GROUP=1 launch calls _index_block_score_page, same op order as main); the ImportError skip is removed; the [2858, 2331] parallel-selector case is added; the warp rationale is moved, the "short/underfilled" comment dropped, the gist log pointer restored, and selector_rows/query_rows renamed. Cleanup A/B is in the description. Merges cleanly with main incl. #3104 (its _launch_select call is decode-only and compatible).}