Repository navigation
Conversation
…-invariant) test_block_topk_matches_torch_topk compared the top-k selection by index against torch.topk. In "ties" mode the selection boundary is heavily tied, and the kernel breaks ties deterministically by lowest column id (its documented design) while torch.topk on GPU breaks ties by its own order -- not lowest-index, and not deterministic run-to-run. Both pick a valid top-k, but their index sets differ on the tied boundary, so the index comparison fails with a run-to-run-varying mismatch count. Compare the winning score multiset instead (-1 slots counted as -inf), which every valid top-k shares. The kernel is unchanged; the structural assertions (live count, front-packing, ascending order) are kept.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem (ROCm-specific)
tests/models/qwen4_exp/test_qsa_kernels.py::test_block_topk_matches_torch_topk[ties-4-4096-512]fails on ROCm (gfx1201 / R9700) with a mismatch count that varies run to run (e.g. 41/2048, then 80/2048 — so it is also flaky). On CUDA it passes, which is why NV CI never caught it.Measured, torch 2.14 / Triton 3.8, same test:
Root cause
The test compared the top-k selection by index against
torch.topk. Intiesmode the scores arerandint(0, 3), so the top-512-of-4096 boundary is heavily tied.qsa_block_topkbreaks ties deterministically by lowest column id — its documented design (qsa/topk.py: "candidates stay in ascending column order throughout, so ties still go to the lowest column").torch.topk's tie-breaking is implementation-defined (the API does not guarantee which of several equal elements are returned), and the two backends differ:It is the same torch version on both (2.14.0,
+cu130vs+rocm7.14), so this is a backend-kernel difference, not a version issue. Both backends return a valid top-k; the test was relying ontorch.topk's unspecified, backend-specific tie order happening to match the kernel's lowest-column convention, which only holds on CUDA.Fix
Top-k with ties is a set of scores, not of indices: the winning score multiset is unique even when which tied column fills a given slot is not. The test now compares the sorted selected-score multiset (empty
-1slots counted as-inf) instead of the index set — platform-independent, and no longer dependent ontorch.topk's tie order. The kernel is unchanged; the structural assertions (live count, front-packing, ascending order) are kept.Validation
See the table: green on both backends, and deterministic on ROCm where the old index comparison was flaky.