perf(hip-kernel-provider): gfx942 4-warp GQA softmax exp2_fast + bf16 swizzle-hoist - #8
Open
AviralGoelAMD wants to merge 1 commit into
Open
AviralGoelAMD wants to merge 1 commit into
AviralGoelAMD wants to merge 1 commit into
Conversation
… swizzle-hoist The gfx942 4-warp GQA attention kernel (build_gfx942_4warp_gqa) ran the online-softmax through the guarded math.exp2 (overflow/underflow range reduction) and hoisted the LDS bank-swizzle columns for fp16 only. - Softmax now emits exp2_fast: both exp2 arguments (m_old-m_new and S*scale-m_new) are <= 0 by the running-max invariant, so the guard is unnecessary and v_exp_f32 flushes the tail correctly. - Enable the swizzle-hoist for bf16 as well as fp16: this cohort is the small-tile BN=32 wide-flash path (HD128_PIPE => BS<=32), whose register pressure leaves headroom for the precomputed columns (the earlier bf16 spill concern was for larger tiles this path never uses). Verified correct on real gfx942 vs an fp32 windowed-SDPA reference for bf16 and fp16 (D128, GQA 32/8, Sq 8192 and 16384). Adds a CPU-only emit test that guards the exp2_fast lowering for both dtypes.
🤖 rocKE PR ReviewPR #8 · 2 files · +96/-7 Branch: A. Byte-identity & two-engine parity — graded 9/9
B. Coverage & silent-fallback — graded 5/5
C. DRY / single-source / blast-radius — graded 8/8
D. Dispatcher & spec integration — graded 10/10
E. Structure & hygiene — graded 10/10
F. Merge-readiness / PR-bot policy — graded 10/10
G. Verification honesty — graded 23/23
H. Kernel input validation & memory safety — graded 5/5This PR modifies only: (1) the H trace (required): No device-side mask or bound is added or modified. Trace is vacuous — all H items are N/A.
I. Reusable design principles — graded 5/5
RollupTotal graded: 82/82 (all items graded) Findings — BLOCKER → HIGH → MEDIUM → LOW
Final Verdict: FAILBlocking items (must fix before merge):
High-priority items (strongly recommended before merge): Non-blocking (clean up in this or a follow-up PR): Token usage (Claude-Sonnet-4.6 via AMD LLM gateway): input 300 · output 67515 · cache 14072491 · total 14140306 · est. $7.128069900000002 |
AviralGoelAMD
force-pushed
the
develop
branch
from
September 6, 2026 16:23
25fc943 to
7b4a377
Compare
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.
Summary
The gfx942 4-warp GQA attention kernel (the paged-prefill path used for D128
sliding-window and D256) did avoidable work in its inner loop: the online-softmax
exp2went through a full overflow/underflow guard on every element even though theexponent is always non-positive, and the LDS bank-swizzle addressing was precomputed
once-per-lane only for fp16 (recomputed on every access for bf16). This change removes
both. The kernel is measurably faster and produces identical results. Confidence is
high: verified correct on real gfx942 against an fp32 reference for both bf16 and fp16
across two sequence lengths, with a new test guarding the change.
What was wrong
exp2(math.exp2, an extrarange-reduction clamp per call) at both exp2 sites, though the exponent is always
<= 0.recomputing
div/mod/xor/mul/addper K/V access) was gated to fp16; bf16 recomputed it.What changed
library/kernels/gfx942/attention_tiled_2d.py(build_gfx942_4warp_gqa):exp2_fast;exp2_fastfor bf16 and fp16.Why it works
exp2_fastisv_exp_f32without the overflow/underflow guard. Both arguments —m_old - m_newandS*scale - m_new— are<= 0by the running-max invariant, so theguard can never fire and the raw op flushes the tail correctly. Safe by construction,
independent of dtype/shape.
(key, col)oncebuf_offis0 mod 16(it is), sothey are loop-invariant and hoistable. bf16 was excluded over a 256-VGPR spill that only
arises at larger tiles; this cohort is the small-tile
BN=32path (HD128_PIPE => BS<=32),which has register headroom (measured; no spill).
How it was validated
D128 GQA 32/8, sliding window, Sq 8192 and 16384 — all within tolerance.
exp2) and passes after.pre-commitclean; deliverable source only (kernel + test), no scratch/bench.Notes
not in this body or the commits.
exp2_fastalso applies to the D256 4-warp path (same shared softmax); its correctnessrides the same
arg <= 0proof, and its IR is covered by the representative golden.attention_d256/gfx942/4warp_gqaentry inrocke_representative_ir_sha256.jsonmust bere-blessed under the CI LLVM flavor. It cannot be blessed from the dev environment used
here (LLVM-flavor mismatch + the C++ engine archive is not built).
with the upstream PR.