Make Qwen3 device top-k selection exact#787
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughQwen3-14B top-k selection now uses vocabulary groups instead of chunks, handles a vocabulary tail, merges grouped candidates through staged sorting, and derives final indices from selected positions. Test fixtures cover active batches, while the golden reference directly applies ChangesQwen3 grouped top-k selection
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the top-k selection logic in topk_select.py to use a grouped top-k approach with a new helper function _topk_group_pairs, while simplifying the golden reference implementation. The review feedback highlights two key issues: first, the tail group logic should be wrapped in a compile-time check (if TOPK_GROUP_TAIL != 0:) to prevent out-of-bounds slicing when TOPK_GROUP_TAIL is zero; second, an explicit assertion assert VOCAB_CHUNK == 512 should be added to prevent silent correctness bugs since the sorting pipeline currently assumes this specific chunk size.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| tail_start = TOPK_NUM_FULL_GROUPS * TOPK_GROUP_WIDTH | ||
| tail_scores_raw = logits[ | ||
| b : b + 1, tail_start : tail_start + VOCAB_CHUNK | ||
| ] | ||
| tail_scores = pl.fillpad( | ||
| pl.set_validshape(tail_scores_raw, 1, TOPK_GROUP_TAIL), | ||
| pad_value=pl.PadValue.min, | ||
| ) | ||
| tail_indices = pl.arange( | ||
| pl.cast(tail_start, target_type=pl.UINT32), | ||
| [1, VOCAB_CHUNK], | ||
| dtype=pl.UINT32, | ||
| ) | ||
| tail_pairs = pl.sort32(tail_scores, tail_indices) | ||
| tail_pairs = pl.mrgsort(tail_pairs, block_len=64) | ||
| tail_pairs = pl.mrgsort(tail_pairs, block_len=256) | ||
| tail_pairs = tail_pairs[:, 0 : 2 * TOPK] | ||
| tail_vals = pl.gather(tail_pairs, mask_pattern=pl.tile.MaskPattern.P0101) | ||
| tail_ids = pl.gather( | ||
| tail_pairs, | ||
| mask_pattern=pl.tile.MaskPattern.P1010, | ||
| output_dtype=pl.INT32, | ||
| ) | ||
| tail_offset = TOPK_NUM_FULL_GROUPS * TOPK | ||
| for k in pl.range(TOPK): | ||
| pl.write(candidate_vals, [0, tail_offset + k], pl.read(tail_vals, [0, k])) | ||
| pl.write(candidate_ids, [0, tail_offset + k], pl.read(tail_ids, [0, k])) |
There was a problem hiding this comment.
If TOPK_GROUP_TAIL == 0 (e.g., if REAL_VOCAB is a multiple of TOPK_GROUP_WIDTH), the unconditional execution of the tail group logic would slice logits out of bounds (specifically when REAL_VOCAB == VOCAB), causing compilation or runtime failure. Wrapping the tail group logic in a compile-time check if TOPK_GROUP_TAIL != 0: completely avoids this issue.
if TOPK_GROUP_TAIL != 0:
tail_start = TOPK_NUM_FULL_GROUPS * TOPK_GROUP_WIDTH
tail_scores_raw = logits[
b : b + 1, tail_start : tail_start + VOCAB_CHUNK
]
tail_scores = pl.fillpad(
pl.set_validshape(tail_scores_raw, 1, TOPK_GROUP_TAIL),
pad_value=pl.PadValue.min,
)
tail_indices = pl.arange(
pl.cast(tail_start, target_type=pl.UINT32),
[1, VOCAB_CHUNK],
dtype=pl.UINT32,
)
tail_pairs = pl.sort32(tail_scores, tail_indices)
tail_pairs = pl.mrgsort(tail_pairs, block_len=64)
tail_pairs = pl.mrgsort(tail_pairs, block_len=256)
tail_pairs = tail_pairs[:, 0 : 2 * TOPK]
tail_vals = pl.gather(tail_pairs, mask_pattern=pl.tile.MaskPattern.P0101)
tail_ids = pl.gather(
tail_pairs,
mask_pattern=pl.tile.MaskPattern.P1010,
output_dtype=pl.INT32,
)
tail_offset = TOPK_NUM_FULL_GROUPS * TOPK
for k in pl.range(TOPK):
pl.write(candidate_vals, [0, tail_offset + k], pl.read(tail_vals, [0, k]))
pl.write(candidate_ids, [0, tail_offset + k], pl.read(tail_ids, [0, k]))There was a problem hiding this comment.
Addressed in bc73c77: the tail-group path is now guarded by the compile-time condition TOPK_GROUP_TAIL != 0, so zero-tail configurations cannot form an out-of-bounds slice. The current Qwen3 exact Top-32 NPU golden test passes.
| assert VOCAB % VOCAB_CHUNK == 0 | ||
| assert REAL_VOCAB <= VOCAB | ||
| assert TOPK <= VOCAB_CHUNK | ||
| assert CHUNK_TOPK <= TOPK | ||
| assert CHUNK_TOPK <= CHUNK_TOPK_PAD | ||
| assert REAL_NUM_VOCAB_CHUNKS * CHUNK_TOPK <= CANDIDATE_PAD | ||
| assert TOPK_GROUP_WIDTH % VOCAB_CHUNK == 0 | ||
| assert TOPK_GROUP_TAIL == REAL_VOCAB_TAIL | ||
| assert TOPK_NUM_GROUPS * TOPK <= TOPK_CANDIDATE_PAD | ||
| assert NUM_VOCAB_CHUNKS <= GREEDY_CHUNK_PAD |
There was a problem hiding this comment.
The mrgsort pipeline in both the greedy path (lines 96-97) and the tail path (lines 199-200) only goes up to block_len=256, which can only fully sort up to 512 elements. If VOCAB_CHUNK is configured to any other value (e.g., 1024 or 2048) in T.vocab_chunk, the sorting will be incomplete, leading to silent correctness bugs. Adding an explicit assertion ensures that any future configuration changes are caught immediately.
| assert VOCAB % VOCAB_CHUNK == 0 | |
| assert REAL_VOCAB <= VOCAB | |
| assert TOPK <= VOCAB_CHUNK | |
| assert CHUNK_TOPK <= TOPK | |
| assert CHUNK_TOPK <= CHUNK_TOPK_PAD | |
| assert REAL_NUM_VOCAB_CHUNKS * CHUNK_TOPK <= CANDIDATE_PAD | |
| assert TOPK_GROUP_WIDTH % VOCAB_CHUNK == 0 | |
| assert TOPK_GROUP_TAIL == REAL_VOCAB_TAIL | |
| assert TOPK_NUM_GROUPS * TOPK <= TOPK_CANDIDATE_PAD | |
| assert NUM_VOCAB_CHUNKS <= GREEDY_CHUNK_PAD | |
| assert VOCAB % VOCAB_CHUNK == 0 | |
| assert REAL_VOCAB <= VOCAB | |
| assert TOPK <= VOCAB_CHUNK | |
| assert VOCAB_CHUNK == 512 | |
| assert TOPK_GROUP_WIDTH % VOCAB_CHUNK == 0 | |
| assert TOPK_GROUP_TAIL == REAL_VOCAB_TAIL | |
| assert TOPK_NUM_GROUPS * TOPK <= TOPK_CANDIDATE_PAD | |
| assert NUM_VOCAB_CHUNKS <= GREEDY_CHUNK_PAD |
References
- When assuming a single-tile coverage (e.g., T_PAD <= MM_ROW_TILE) in a kernel, make this invariant explicit with a module-level assertion (e.g., assert T_PAD == MM_ROW_TILE) to prevent silent correctness issues if configurations change.
There was a problem hiding this comment.
Addressed in bc73c77: added assert VOCAB_CHUNK == 512 to make the sort-pipeline invariant explicit and fail fast if the tiling configuration changes.
4f6f119 to
f44ddf9
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/qwen3/14b/topk_select.py (1)
264-266: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDistributed fixture doesn't actually cross the half0/half1 merge boundary.
spread_idsuses strideVOCAB_CHUNK(512), so withTOPK=32the max index is ~15,879 — entirely inside groups 0-7 (candidate slots 0-255), which all land inhalf0(slots 0-2047). This fixture never places a top-32 value inhalf1(groups 64+, tail) or near the 2048-slot boundary, so it can't exercise the cross-half merge that is central to this PR's exactness guarantee.Consider spreading across the full
REAL_VOCABrange instead, e.g.:🧪 Proposed fix to actually exercise the half0/half1 boundary
- spread_ids = torch.arange(TOPK, dtype=torch.long) * VOCAB_CHUNK + 7 + spread_ids = torch.arange(TOPK, dtype=torch.long) * (REAL_VOCAB // TOPK) + 7 logits[1, spread_ids] = torch.arange(TOPK, 0, -1, dtype=torch.float32) + 1000.0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/qwen3/14b/topk_select.py` around lines 264 - 266, Update the BATCH > 1 fixture in the distributed top-k test around spread_ids so selected logits span the full REAL_VOCAB range, including indices near and beyond the 2048-slot half0/half1 boundary. Preserve the descending score ordering and TOPK cardinality while ensuring the generated candidates exercise the cross-half merge path.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@models/qwen3/14b/topk_select.py`:
- Around line 264-266: Update the BATCH > 1 fixture in the distributed top-k
test around spread_ids so selected logits span the full REAL_VOCAB range,
including indices near and beyond the 2048-slot half0/half1 boundary. Preserve
the descending score ordering and TOPK cardinality while ensuring the generated
candidates exercise the cross-half merge path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 77dc0714-7f8b-487a-a464-0852fba7c65c
📒 Files selected for processing (1)
models/qwen3/14b/topk_select.py
ed45657 to
bc73c77
Compare
|
Addressed the latest review feedback in bc73c77. In addition to the tail guard and 512-chunk invariant, the distributed Top-32 fixture now spans the full real vocabulary and places candidates in both halves of the final merge. The updated fixture passes exact NPU comparison against |
8c6d271 to
9fdc34a
Compare
2 participants
Summary
Follow up on #769 by replacing the Qwen3-14B per-chunk Top-4 approximation with a provably exact full-vocabulary Top-32 selector.
What Changed
sort32andmrgsort.75 x 32 = 2400candidates in a 4096-entry padded buffer and select the exact global Top-32.torch.topk(logits[:, :REAL_VOCAB], 32).If an entry is not in its group's Top-32, at least 32 entries in that group rank ahead of it, so it cannot belong to the global Top-32. Therefore the union of exact group Top-32 sets contains the exact global Top-32.
Validation
selection_k=1regression: passed.Paired serving PR: hw-native-sys/pypto-serving#77