[XiaobingSuper]

Summary

Allow the existing topk_per_row_small_k narrowing selector to build on wave32. ROCm ballots still return an i64 mask on wave32 with the high half clear, so the existing low/high mask storage, lane shifts, threshold search, and tie ordering work unchanged.

This enables the selector on gfx1250 without changing gfx942/gfx950 behavior.

MiniMax-M3 results

MI455X e15-09 GPU0, exact ordered top-16 with forced local blocks, paired against ATOM Triton selectors:

• width 958 / 2,420 rows: 21.82 -> 10.24 us selector-only (2.13x)
• width 1,895 / 4,104 rows: 34.16 -> 11.86 us (2.88x)
• width 3,942 / 2,858 rows: 54.16 -> 11.08 us (4.89x)
• DPEP width 640 / 65,536 rows: 173.95 -> 81.18 us (2.14x)

With ATOM score + sparse-table emit included, a workload/CU-gated dispatch improves the 22 production prefill shapes by 1.045x aggregate; long shapes improve 11.8-12.9%, and DPEP q2K-q32K improves 9.2-15.6%.

Validation

• gfx1250 gauss, recency, clustered/tied, ascending worst-case distributions
• widths 256/640/1024/2048/4096; rows 16/512/4096
• init/local forced blocks and ragged row lengths
• every selected id exact against the stable torch reference
• Black and Ruff clean

ATOM integration is a separate stacked change.

Made with Cursor

PR Self-Review

Reviewed revision

ce3681929fe19dd2c0e34bc392388fcf331a2496

Behavior and intent

The existing small-k row selector rejected every wave32 build even though its i64 ballot representation is valid there: ROCm leaves the upper 32 ballot bits clear, and the kernel already stores/reconstructs the low and high halves separately. This change permits wave32 while preserving the wave64 code path, threshold search, survivor bound, forced-block pins, and larger-column tie rule. The purpose is to let gfx1250 use the existing narrowing algorithm instead of sorting a wide score row tile by tile.

Validation and unverified scope

Validated on one MI455X/gfx1250 GPU with FP32 score rows and top-k 16:

• widths 256, 640, 1024, 2048, and 4096; rows 16, 512, and 4096
• gaussian, recency-biased, clustered/tie-heavy, and strictly ascending worst-case distributions
• init/local forced blocks and ragged row lengths, including rows shorter than k
• every selected id matched the stable Torch reference exactly
• MiniMax-M3 production-shape parity included up to 131,072 selector rows; q32K was repeated 20 times with zero mismatches
• Black, Ruff, repository checks, wheel builds, MI35X OPUS/standard tests, and the aiter test gate pass

The PR does not change selector dispatch policy in aiter and does not claim performance for arbitrary wave32 GPUs or k values other than the measured MiniMax-M3 top-16 use. gfx1250 integration and workload/CU gating live in ROCm/ATOM#2971.

PR-specific question

Before this PR, build_topk_per_row_small_k_module(..., wave_size=32) failed immediately with ValueError, so no gfx1250 correctness or performance test could reach the kernel. Existing gfx942/gfx950 coverage exercises only wave64 and therefore cannot detect a broken wave32 ballot/mask path. The added gfx1250 entry makes the canonical distribution and ragged matrices execute that previously rejected path.

Changed op tests

op_tests/test_flydsl_topk_per_row_small_k.py adds gfx1250 to the supported architecture list. This is necessary and non-redundant because the file is the canonical oracle/performance matrix for this selector: it covers adversarial ascending rows, FP8-like tied distributions, forced blocks, ragged lengths, and exact tie ordering. No other op test previously instantiated this kernel with wave32.

[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.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter standard tests on MI35X
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:mi300x	Run Aiter standard tests on MI300X in PRs; main branch always runs them
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the Radeon R9700 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7878 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Reviews:
A reviewer is usually assigned within one working day after the PR is ready for review (see Code Review Process). If nobody has responded for about a week, mention the area contact listed in CONTACTS.md. Want to help review? See REVIEWERS.md.

[aiter-bot]

WAITING ON AUTHOR

PR Self-Review

Revision ce3681929fe1 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

1. What behavior does this PR change, and why is the existing behavior insufficient?
2. What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
3. What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
4. For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

• FAIL — validation answer must name a command, result, or evidence link
• FAIL — validation answer must explicitly state what remains unverified (or say none with a reason)

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[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.