Repository navigation
perf(rocke): enable exp2_fast for bf16 head_dim=128 in gfx942 dense prefill - #10832
AviralGoelAMD wants to merge 3 commits into
Conversation
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
✅ 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. |
AviralGoelAMD
changed the title
There was a problem hiding this comment.
Pull request overview
This PR updates the gfx942 attention_dense softmax policy to always use the exp2_fast path (removing the prior bf16 + head_dim=128 exception) and refreshes the associated golden IR hash artifacts to match the newly emitted kernels.
Changes:
- Enable
exp2_fastunconditionally inattention_densesoftmax policy for gfx942 (includingbf16 D128). - Update golden IR SHA256/size records for affected
bf16 D128cases under bothllvm20andllvm22.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py | Removes the bf16 D128 policy exception by making _use_exp2_fast() always return True. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json | Updates sizes and SHA256 hashes for the impacted bf16 D128 gfx942 attention_dense golden artifacts. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…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.
8499f41 to
d4de8c1
Compare
_use_exp2_fast now returns True for all configs (the bf16 D128 spill holdout is stale on the lazy-rescale schedule), so the pinned matrix in test_exp2_fast_gate_matches_the_spill_measured_matrix must expect True for (128, bf16). Flip the case and refresh the docstring; the parametrized case fails on the pre-change gate and passes after, guarding the flip.
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
…no-spill - black reformats the flipped (128, bf16, True) parametrize case - _use_exp2_fast docstring now cites the llvm-readelf --notes code-object read (.private_segment_fixed_size / .vgpr_count) for the bf16-D128 no-spill claim, matching the PR body: the lazy-rescale build reports 0 scratch on both the default and persistent grids (VGPR reduced, not increased, vs plain exp2).
| correctness one. | ||
| """ | ||
| return dtype == "fp16" or head_size != 128 | ||
| return True |
There was a problem hiding this comment.
The reason I keep this function to just return true is becasuse we can reuse this later in a new scenario where some datatype would need to be excluded.
|
Recreated as #10885 on |
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.