[ROCm] Fix packed-bag pooling truncation in nbit inference forward - #166
Open
aryaman-gupta wants to merge 4 commits into
Open
[ROCm] Fix packed-bag pooling truncation in nbit inference forward#166aryaman-gupta wants to merge 4 commits into
aryaman-gupta wants to merge 4 commits into
Conversation
In PackedMode the accumulate and store stages map lanes to bags at uint granularity while the load stage uses uint4 granularity, so the two need different entries of Ls[]. The kernel translated between them by shuffling Ls[] in place, inside the L_start loop. That corrupts the load stage, which keeps reading Ls[] on later passes, and corrupts the shuffle source lanes themselves, so from the second pass onwards the longer bag of a packed pair is truncated to its partner's pooling length. Compute the accumulate-stage lengths once into a separate Ls_acc[], leaving Ls[] intact for the load stage. The mapping is loop-invariant, so this also removes OutputRowsPerThread shuffles from every pass of the loop. Also make max_Ls wave-uniform under PackedMode. The L_start loop is wave-collective, so a per-lane bound lets the short-bag lanes exit while the remaining lanes still shuffle against them. Per-row validity checks against Ls[]/Ls_acc[] still bound each lane's own work, and a divergent wave already executes the union of all lanes' iterations, so this adds none. Fixes the INT4 D=160 subtests of test_nbit_forward_nan_zero_fill. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
max_Ls is wave-uniform under PackedMode, so min(InputRowsInFlight, max_Ls - L_start) already evaluates identically on every lane and the shuffle that followed it is the identity. nbit_forward_test.py is unchanged at 12 passed / 8 subtests passed, and the divergent-Ls probes stay clean. Mean of 3 runs on MI350X, INT4 pooled forward: D=160 packed, ragged 65.40 -> 64.97 us D=240 packed, ragged 64.97 -> 64.50 us D=160 packed, uniform 45.73 -> 45.83 us D=1024 not packed 139.80 -> 139.40 us The non-packed config cannot be affected by this change yet moves by a similar amount, so treat ~0.3% as the noise floor and the change as perf-neutral. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
aryaman-gupta
changed the title
…gths test_nbit_forward_nan_zero_fill only covers bag lengths that happen to agree between packed neighbours, so it does not catch a kernel that applies one bag's L to its partner. Add a case with alternating short/long lengths, which is the shape that exposes it. Every row is 1.0 and nothing is pruned, so each bag must sum to its own L. Verified red/green: passes on the fixed kernel, and fails without the wave-uniform max_Ls (max abs diff 60.0, 272/1280 elements). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
avbokovoy
approved these changes
Aug 19, 2026
avbokovoy
left a comment
avbokovoy
left a comment
There was a problem hiding this comment.
Strange thing, but I recall addressing this issue a while ago. Anyway, nice catch and great changes
| self._execute_nan_zero_fill(weights_ty, D, output_dtype, weighted) | ||
|
|
||
| @unittest.skipIf(*gpu_unavailable) | ||
| def test_nbit_forward_packed_bags_uneven_pooling(self) -> None: |
There was a problem hiding this comment.
Should it be guarded with @skipIfNotRocm?
Author
There was a problem hiding this comment.
I previously thought the test could be useful on CUDA also, but you're right that it doesn't make sense to have a test with packed_bags in the name running on CUDA. Done in 7e12fd6
Replace the hand-rolled butterfly with warp_reduce_max() from utils/find_qparams.cuh, which is the same shfl_xor reduction over kWarpSize and is already in scope via embedding_forward_template_helpers.cuh. Guard test_nbit_forward_packed_bags_uneven_pooling with skipIfNotRocm: bag packing is a ROCm-only path, so elsewhere the test would pass without exercising what its name describes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
2 participants
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
On ROCm,
TBE_ROCM_INFERENCE_PACKED_BAGSpacks several bags into a single wave. When those bags have different pooling lengths, the longer one is silently truncated to its partner's length and the embedding output is wrong.It only surfaces when a bag's length exceeds
InputRowsInFlight, and it disappears entirely when packed bags happen to share a length — which is why it went unnoticed.Cause
The accumulate/store stages map lanes to bags differently from the load stage (uint vs uint4 granularity). The kernel translated
Ls[]between the two by shuffling it in place inside theL_startloop. That corrupts the load stage, which keeps readingLs[]on later passes, and also corrupts the shuffle source lanes themselves — so from the second pass onwards every lane broadcasts an already-permuted value.Fix
Ls_acc[]before the loop, leavingLs[]intact for the load stage. Being loop-invariant, this also removes shuffles from every pass.max_Lswave-uniform under PackedMode. The loop is wave-collective (syncwarp,shfl_sync), so a per-lane bound let short-bag lanes exit while the remaining lanes still shuffled against them.Both changes are required — either one alone still fails.
Testing (MI350X / gfx950, ROCm 7.1)
nbit_forward_test.py: 2 failed → 12 passed, 8 subtests passedCUDA and the
nobagpath are unaffected — PackedMode is ROCm-and-pooled-only.