Repository navigation
perf(rocke): enable exp2_fast for bf16 head_dim=128 in gfx942 dense prefill - #10885
Conversation
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
…refill The bf16-D128 exp2_fast holdout was justified by a pre-lazy-rescale spill over the waves-per-eu=2 cap. rocprofv3 on the current lazy-rescale kernel shows the opposite: VGPR drops and scratch stays zero, output numerically identical to plain exp2. Remove the stale holdout in _use_exp2_fast (return True), update the in-file status/policy comments that referenced the old exception, and re-bless the four bf16-D128 golden cases across both llvm flavors; fp16 and D64 are byte-identical.
d4de8c1 to
5da51c8
Compare
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
There was a problem hiding this comment.
Pull request overview
Enables exp2_fast for gfx942 bf16 D128 dense-prefill kernels and updates their golden IR hashes.
Changes:
- Enables
exp2_fastfor all gfx942 configurations. - Documents updated spill measurements.
- Re-blesses affected LLVM 20/22 golden outputs.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
library/kernels/gfx942/attention_dense.py |
Enables the optimized exponential path. |
library/tests/golden/attention_dense_gfx942_ir_sha256.json |
Updates affected IR sizes and hashes. |
Suppressed comments (1)
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py:538
- This policy change leaves the existing spill-matrix test failing deterministically:
test_exp2_fast_gate_matches_the_spill_measured_matrixstill parameterizes(128, "bf16", False)and asserts the result attests/test_attention_dense_gfx942.py:846-865. Update that case to expectTrueand revise its stale spill rationale so the dense gfx942 test lane passes.
return True
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
aosewski
left a comment
aosewski
left a comment
There was a problem hiding this comment.
Went through this one carefully, including re-running the golden fixture at both flavors and re-measuring the registers with the prefill occupancy probe compiled for gfx942 (ROCm 7.1 / llvm20).
The change itself checks out. The IR delta is exactly the intrinsic swap and nothing else (34 llvm.exp2.f32 -> llvm.amdgcn.exp2.f32, byte-identical otherwise at both llvm20 and llvm22), the golden re-bless is accurate and minimal, and the resource claim holds: bf16-D128 goes 196 -> 174 VGPR / 0 spill on the default grid and 200 -> 180 / 0 spill on the persistent one. bf16-D64 even loses a 12-VGPR spill. No objection to the flip.
Comments below are all prose and test hygiene. The recurring one is that the doc-consistency pass caught the two big docstrings but left about five other places still describing the holdout you just removed.
- Add default-grid numeric coverage and a forced-off/on A/B (torch.equal) that actually distinguishes the two exp2 arms on both bf16-D128 grids; the fp32-SDPA oracle (tol=4e-2) cannot. - Parametrize the softmax-intrinsic test on the tuning, not the policy, so the plain-exp2 codegen arm stays pinned now that the policy is True everywhere; add the bf16-D128 row to the fused-rescale cast-count test. - Rename the policy test to test_exp2_fast_is_enabled_for_every_config and de-spill its docstring/section header. - Correct kernel + README prose: drop the unsupported "lazy-rescale schedule change" cause (the kernel body is unchanged since first commit), restore the tried/rejected/re-adopted lever-record entry, replace "VGPR-neutral" with the measured VGPR reduction, and use "numerically identical for every reachable softmax argument". - Wrap over-length comment lines to the 88-col style limit.
aosewski
left a comment
aosewski
left a comment
There was a problem hiding this comment.
Everything from the last round is addressed — thanks for the thorough pass, particularly reframing the rationale instead of just softening the wording.
I re-verified rather than reading the diff: golden fixture green at both llvm20 and llvm22 (untouched, correctly — only prose and tests moved), CPU lane clean apart from the four pre-existing paged-field failures that also fail on the merge-base, and all three over-length lines are gone.
I paid particular attention to the new test_exp2_fast_matches_plain_exp2, since it's skip-gated and CI can't catch a mistake in it. It holds up: the signatures line up (_spec(..., *, batch, sq) and run_attention_dense_torch(*, spec, q, k, v, out, scale, tuning)), both rows collect, and the premise it rests on is real — forced-off and forced-on genuinely produce two different binaries on both grids:
bf16 D128 default off -> ..._wpe2_e2f0 fast= 0 plain=33
on -> ..._wpe2 fast=33 plain= 0
bf16 D128 persist off -> ..._wpe2_e2f0 fast= 0 plain=33
on -> ..._wpe2 fast=33 plain= 0
That's the gate the "results unchanged" claim was missing, and it's a better answer than the cohort row I asked for.
Approving. Two leftovers below — neither blocks this, but flagging them so they don't quietly become permanent.
The exp2_fast review pass renamed the test section header to "fused rescale" but left three sibling mentions still reading "fused/lazy rescale". On gfx942 the only "lazy" artifact is an inert kernel-name token, so the "/lazy" wording misdescribes the mechanism (a single fused pass: exp2 -> l_local accumulate -> cast -> pack). Align all three to "fused rescale" -- the P2 progress-table header and the rescale comment in attention_dense.py, and the fused-rescale test docstring -- so the file names the mechanism one way. Comment/docstring only; no code, IR, or golden change.
…2 dense prefill) (#11079) **Hot fix** for a short-sequence regression from #10885 (AICK-1991), which enabled `exp2_fast` for bf16 `head_dim=128` unconditionally. On the develop merge base (gfx942, ROCm 7.2.2, same-node A/B) it clearly regresses short-seq dense prefill (~-15 to -21% at Sq <= 1024) while winning at Sq >= 4096 (+8-10%); Sq 2048 sits in the same-node run-to-run noise band. Guard: bf16 `head_dim=128` uses `exp2_fast` only when `seqlen >= 4096`; shorter sequences revert to plain `exp2` (their original, pre-#10885 perf). The conservative cut never regresses a shape (it forgoes at most a noise-level Sq 2048 win). fp16 and D64 are unaffected; gfx950 makes an independent exp2_fast decision and is not covered here. The proper fix is seqlen-aware kernel selection in the rocKE dispatcher; that work is WIP, so this lands the guard in the kernel builder as a hot fix. Jira: AICK-1991
ISSUE ID : AICK-1991
Enable
exp2_fastfor bf16 head_dim=128 in the dense prefill softmax — the last config still on the slower guardedexp2(fp16 and bf16-D64 already used it).The holdout was a stale spill: it was measured on the old pre-lazy-rescale body. The current kernel doesn't spill with
exp2_fast(zero scratch on both grids), and the output is bit-identical to plainexp2(both softmax args are≤ 0, exactly its precondition) — so it's a free softmax speedup on bf16-D128.Diff:
_use_exp2_fast→True, plus a golden re-bless of the four bf16-D128 IR cases (fp16 and D64 unchanged; gfx950 untouched). No C++ mirror underlibrary/kernels/, so byte-identity doesn't apply.