cuda: wave64 (CDNA2 / gfx90a) correctness for turbo KV-cache + ConvRot - #328
Merged
TheTom merged 4 commits intoAug 31, 2026
Merged
Conversation
added 4 commits
August 29, 2026 21:54
The convrot_inverse butterflies and the cooperative-kernel final reduction in mmv-cr.cu omit the width argument to __shfl_sync / __shfl_down_sync. On wave64 (CDNA / gfx90a) the physical wave is 64 lanes, so the two logical 32-lane groups (lane = p & 31, block dim 64) are no longer separate warps: the upper group reads shuffle sources from the lower 32 physical lanes, corrupting the radix-4 inverse rotation, and the reduction inside `if (p < 32)` pulls stale row_sum from lanes that never entered the branch. Pass explicit WARP_SIZE (32) width so each 32-lane subsection is self-contained, matching the convrot.cu / fwht.cu convention. Fixes the q6_cr/q8_cr MUL_MAT failures in test-backend-ops on gfx90a.
k_set_rows_turbo2/3/4 pack the turbo KV cache with warp-cooperative shuffles
and __ballot_sync that assume warpSize == 32. On wave64 (CDNA / gfx90a):
- the qs / nibble gathers use absolute srcLanes from lane = j % 32 but omit
the shuffle width, so threads in the upper 32 physical lanes read from the
lower 32 -> corrupt qs / nibble bytes;
- __ballot_sync returns a 64-bit mask that truncates into uint32_t and the
per-8-lane byte extraction only sees the low 32 lanes -> corrupt sign plane.
Pass explicit WARP_SIZE (32) width on the gathers, and replace the sign ballot
with the same width-32 shuffle-gather the qs packing uses (wave-size agnostic,
no 64-bit-ballot truncation). The intra-warp xor reductions are left as-is:
their offsets are <= 16 so each butterfly stays inside its 32-lane subgroup.
Previously the turbo KV cache was populated with garbage on the MI210.
flash_attn_ext_vec broadcasts each K-position's softmax weight KQ_reg to the V-accumulation lanes with a width-omitted __shfl_sync, but only on the turbo-V path (non-turbo V spills KQ to shared memory). On wave64 (CDNA / gfx90a) the kernel packs two 32-lane logical warps into one 64-lane physical wave (nthreads=128, nwarps=4), so lanes in the upper logical warp read the softmax weight from the WRONG warp's K-positions -> half of every V accumulation is multiplied by the wrong weights -> incoherent output. This is why only turbo KV V-cache (turbo2/3/4) produced garbage on the MI210. Add the explicit WARP_SIZE width to both broadcast sites (V_DOT2 and float paths), matching the KQ_max reduction two lines up. Read-side dequant and the write path were already wave64-correct.
Follow-up to the wave64 correctness fixes. The intra-warp XOR reductions in the turbo set-rows norm/recon kernels and the turbo-wht butterfly omit the shuffle width. They are correct on wave64 today only because their offsets are always < 32 (each aligned 32-lane group self-reduces regardless of physical wave width), but that is latent fragility: a future offset >= 32, or a stricter mask interpretation, would silently break them. Add the explicit WARP_SIZE width so the 32-lane grouping is intent, not accident. No functional change on any current arch; removes the wave64 footgun ahead of UDNA.
Owner
|
Thanks, @jasstrong. I reviewed the shuffle changes and the wave64 failure mode checks out. The explicit width keeps every operation inside the logical 32-lane tile, and replacing the truncated wave64 ballot with an 8-lane shuffle gather is the right fix for the sign bytes. The separate K/V isolation and real MI210 needle test give this useful hardware coverage beyond the backend cases. Merging this. |
TheTom
merged commit Aug 31, 2026
ef68c3c
into
TheTom:feature/turboquant-kv-cache
4 of 21 checks passed
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.
Makes the turbo KV-cache and ConvRot CR kernels correct on wave64 (CDNA / gfx90a / MI210), which previously produced garbage there. Validated on an MI210 — this brings CDNA2 up as a working turbo target.
Root cause
The turbo kernels compute a logical
lane = j % WARP_SIZE(WARP_SIZE is the 32-lane tile width, not the physical wave) then call__shfl_*_syncwith the width argument omitted, so on wave64 the shuffle spans the full 64-lane physical wave and reaches across the 32-lane logical-warp boundary. Fix = the explicit-width idiom already used correctly in these files (convrot.cu,fwht.cu,fattn-vec.cuh:401): passWARP_SIZEas the shuffle width so each 32-lane subsection stays self-contained.Commits
mmv-cr.cu) — theconvrot_inversebutterflies + the cooperative-kernel reduction. Fixes theq6_cr/q8_crMUL_MAT test-backend-ops failures on gfx90a (1694 → 1696/1697; the remaining CR fail is a pre-existing flaky on tiny random shapes, unrelated to wave size).set-rows.cu) — the qs/nibble gathers ink_set_rows_turbo2/3/4(+tails). The__ballot_syncsign-packs are converted to the same width-32 shuffle-gather the qs packing uses (__ballot_syncreturns a 64-bit mask on wave64 and was being truncated into auint32_t).fattn-vec.cuh) — the softmax-weight broadcast to the V-accumulation lanes (turbo-V only; non-turbo V spills KQ to shared memory). This was the root cause of turbo-V garbage.WARP_SIZEwidth on the intra-warp XOR reductions (set-rows norm/recon + the turbo-wht butterfly). Correct today because their offsets are always < 32, so no functional change; this removes the latent wave64 fragility.Validation (MI210 / gfx90a)
test-backend-opsMUL_MAT:q6_cr/q8_crwave64 failures fixed.cache_type_k=turbo3/cache_type_v=turbo4): output went from total garbage to coherent with the needle correctly recalled. Isolation (f16-K/turbo4-V vsturbo3-K/f16-V) confirmed the K and V turbo paths are each independently correct after the fixes.