fix(rocke): apply gfx950 dense softmax scale in fp32 by AviralGoelAMD · Pull Request #13122 · ROCm/rocm-libraries

AviralGoelAMD · 2026-10-05T22:50:17Z

ISSUE ID : AICK-2501

What was wrong

The gfx950 dense prefill attention kernel loses BF16 accuracy at a non-default softmax scale. At scale = 0.5, BF16 output exceeded the BF16 tolerance against an fp32 reference. The default scale (1/sqrt(D)) was fine, which is why the existing tests never caught it.

Cause: the kernel multiplied Q by scale * log2(e) in fp32, then rounded the product back to bf16/fp16 before the QK matrix multiply. That rounding error grows with the scale.

What changed

Path	Change
Ordinary kernel	Q feeds the MFMA unscaled. The row max is taken on raw scores and scaled once; the exp2 argument becomes one `fma(s, qk_scale, -m)`. The mask sentinel is now the power of two `-2*99`, so `sentinel qk_scale` is exact.
Persistent kernel	Each fp32 score is multiplied by `qk_scale` after the MFMA.
`run_attention_dense_torch`	Raises `ValueError` for a NaN or infinite scale, and `NotImplementedError` (`NOT_YET_IMPLEMENTED`) for a finite scale outside the supported range `[2-64, 24]`, including `<= 0`.
hipDNN gfx950 dense matcher	Declines graphs whose scale is outside the same range, so another engine serves them. An absent scale still means 1.0.
IR goldens	The 32 gfx950 `attention_dense` cases are re-blessed (sha256 and size only).
`builders/gfx950/attention/ALGORITHM.md` §0, §2, §8	§0 names the raw score, `qk_scale` and the log2-domain score as separate symbols, and τ is now a runtime argument with a default. §8's ABI bullet states the accepted scale range. The §8.1 lever table describes the fused exp2 argument and its bound, with that bound's premise. §8.2 no longer calls the two grids' inner compute byte-identical: it lists where their softmax arithmetic differs.

Why it works

The QK matrix multiply already accumulates in fp32. Applying the scale there keeps full fp32 precision, so the error no longer grows with the scale. The running max stays in log2 units, so the lazy-rescale threshold and the sink handling are unchanged.

Taking the max on raw scores brings three conditions:

Positive scale. The max commutes with the scale only when the scale is positive.
Bounded scale. For the row's maximum element, fma(s, qk_scale, -m) returns the rounding residue of s * qk_scale: at most half an ulp of the product, so it grows with the scale and the score. At the upper bound 2**4, an exhaustive fp32 check over raw scores up to 1e6 in magnitude gives a residue of at most 1. Together with the lazy-rescale threshold, that keeps P within 2**9, inside fp16 range.
Exact mask sentinel. A query row can be fully masked in the first tile it visits (causal sliding window without sinks), and its running max then comes from the mask sentinel. With the old -1e30 sentinel, the residue of sentinel * qk_scale is huge and its sign depends on the scale: where it is positive (for example scale 1/8 and 0.5), exp2 overflows to inf and then NaN. A power-of-two sentinel makes the product exact. Such a row then gets exactly P = 1 when qk_scale <= 1, or exactly P = 0 when qk_scale > 1. Either way, the next tile with a real key rescales it by alpha = 0, as on develop.

How it was validated

On real gfx950 hardware:

New test test_bf16_d128_non_default_scale fails on develop and passes with this change.
New test test_bf16_sliding_window_no_sinks passes. It covers D64 at its default scale, D128 at scale 0.5, 1.0 and 2**4 (the two branches above), and both the ordinary and persistent kernels. With a -1e30 sentinel the ordinary cases below scale ln 2 produce non-finite output. On develop, the D128 scale 0.5 cases exceed the BF16 tolerance.
New test test_d128_scale_range_edges passes at both bounds of the supported range (2**-64 and 2**4), for bf16 and fp16, ordinary and persistent.
The full gfx950 dense numeric file and the bottom-right tests pass.
Strict C++/Python IR parity (ROCKE_CPP_STRICT=1) and the resource-budget test pass; no scratch spill.
A one-off check against an fp32 reference covers four re-blessed shapes that no in-tree GPU test runs (lazy rescale off, ordinary and persistent; varlen; paged sliding window fp16 with a shuffled block table), at the default scale and at 0.5. Seven other re-blessed shapes run on no GPU test: bn128_sq512, noncausal_sq512, ragged_causal_sq500, ragged_full_sq500, ragged_h64_sq500, persist_ragged_sq500 and persist_hkvmaj_sq512.
A scale sweep shows BF16 error no longer grows with the scale; FP16 also improves.

On the host: the installed-golden gate and the new CPU validation test pass. The re-blessed golden changes only the gfx950 attention_dense entries. The new C++ matcher gtest is not built locally (no hipDNN SDK); it runs only in the hipDNN Linux superbuild CI. The matcher guard was checked with the repo's .clang-tidy config, and its decline set was checked in a standalone build.

Notes

gfx942 attention_dense has the same pre-MFMA rounding. It is out of scope here and tracked in AICK-2502.
There is a small measured slowdown, larger on the ordinary (small-problem) path than on the persistent path. It was measured on gfx950 with cold memory, D128, Sq = Skv = 8192, causal. The bf16 GQA 128/8 numbers come from a run with an A/A control. The fp16 and GQA 32/8 numbers come from an earlier cold A/B run in ABBA order without an A/A control. The measured kernels differ from this PR's head only in the ordinary kernel's mask-sentinel constant. Other shapes, head sizes, and architectures were not measured. Numbers and the profiling analysis are on the private results page.
The persistent kernel scales each score after the MFMA, which adds one multiply per score in its hot loop. It does not use the ordinary kernel's fused form (row max on raw scores, scale folded into the exp2 fma). Porting that form would change the persistent masking, and on the ordinary kernel the fused form still measured a larger slowdown than the persistent kernel shows with this approach.

The gfx950 dense prefill kernel multiplied Q by scale * log2(e) in fp32 and rounded the product back to bf16/fp16 before the QK MFMA. The rounding error grows with the softmax scale, so bf16 output exceeded tolerance at scale 0.5. Q now feeds the MFMA unscaled and the scale is applied to the fp32 scores: - ordinary kernel: the row max is taken on raw scores and scaled once, and the exp2 argument becomes one fma(s, qk_scale, -m), so the scale adds no per-score instruction and the running max stays in log2 units; - persistent kernel: each fp32 score is multiplied by qk_scale after the MFMA. The ordinary kernel's max is only correct for scale > 0, so run_attention_dense_torch now rejects scale <= 0 and NaN, and the hipDNN gfx950 dense matcher declines such graphs. gfx942 attention_dense has the same pre-MFMA rounding and is left for a follow-up.

therock-pr-bot · 2026-10-05T22:54:11Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-10-05T22:54:14Z

🎉 All checks passed! This PR is ready for review.

The ordinary gfx950 dense kernel now masks raw scores and scales the row max afterwards. With the -1e30 mask sentinel, a query row that is fully masked in the first visited KV tile (causal sliding window without sinks) starts its running max from fl(-1e30 * qk_scale), and fma(-1e30, qk_scale, -m) returns the rounding residue of that product. For qk_scale < 1 the residue can be a huge positive number, so exp2 overflows and the row becomes inf/NaN. Use the power-of-two sentinel -2**99, whose product with qk_scale is exact, so the exp2 argument is exactly 0 for such rows, as before the scale moved. Bound the accepted softmax scale to [2**-64, 2**24] in run_attention_dense_torch and in the hipDNN gfx950 dense matcher, so the product stays exact, finite, and far below any real score; this also rejects +-inf. Add a GPU test for causal sliding window without sinks (D64 and D128, ordinary and persistent), extend the scale-range CPU test and gtest, correct the bf16 rounding bound in a comment, note that the exp FMAs may be packed, and re-bless the affected gfx950 attention_dense IR goldens.

Copilot

Copilot review overview

🟡 Changes recommended

Masked keys can contribute at small scales, and positive infinity remains accepted despite producing invalid arithmetic.

Review effort: Balanced
Findings: 1

Open (1)

Reject non-finite scales in engine matching · New

What changed in this PR

Moves gfx950 dense attention scaling to fp32 after QK MFMA to improve non-default-scale accuracy.

Changes:

Reworks ordinary and persistent scaling paths.
Adds scale validation and regression tests.
Re-blesses gfx950 dense IR goldens.

File	Description
`TestGfx950AttentionDenseMatchers.cpp`	Tests invalid scale matching.
`Gfx950AttentionDenseNative.cpp`	Validates graph scale values.
`rocke_representative_ir_sha256.json`	Updates generated IR fingerprints.
`test_attention_dense_gfx950_numeric.py`	Tests non-default BF16 scaling.
`test_gfx950_dense_wiring.py`	Tests runner scale validation.
`attention_dense.py`	Applies scaling after QK accumulation.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

tenpercent · 2026-10-05T23:17:22Z

Reviewed head 0e48d2d84fb8105834ee0d7e3276820d818937d2. One correctness issue remains:

P2 — The ordinary kernel can overflow softmax at an explicitly accepted scale.

In attention_dense.py:967, the fused s * qk_scale - m subtracts a maximum calculated with a separately rounded multiplication at line 798. The rounding residual can be large and positive even for the maximum score.

A concrete case is BF16, D128, causal attention without sinks, ordinary/nonpersistent execution, Q=1, K=1.5, V=1, and scale=2**24—the upper bound explicitly accepted by the Python runner and hipDNN matcher. For example, use B=1, Sq=Skv=512, Hq=4, Hkv=1. The raw score is exactly 192. Host fp32 reproduction using libm.fmaf gives:

qk_scale                    = 24,204,406
rounded maximum             = 4,647,245,824
fmaf(192, qk_scale, -maximum) = 128
exp2(128)                   = inf

This sends infinite weights into the numerator and denominator; normalization produces NaN instead of the expected output 1. The persistent path rounds each scaled score before subtracting the maximum, giving an exponent argument of zero for this case. Changing the mask sentinel to a power of two fixes cancellation for masked scores, but does not fix this cancellation for real scores.

Please use consistent rounding for scores and maxima, as the persistent path does, or otherwise guarantee bounded exponent arguments. Add a GPU numeric regression at the accepted upper bound; the new matcher test verifies eligibility but not numerical correctness.

Validation on this head: 145 focused host tests and 30 subtests passed; the golden gate passed for LLVM20, LLVM22, and LLVM23. This finding is supported by source/IR inspection and host fp32 fmaf reproduction. I did not run GPU numerics or the C++ matcher tests.

The follow-up commit addresses the earlier tiny-scale masking issue; that is not a remaining finding against this head.

In the ordinary gfx950 dense kernel the exp2 argument fma(s, qk_scale, -m) is, for the row-max element, the rounding residue of s * qk_scale. It grows with the scale, so at the previous 2**24 upper bound ordinary raw scores could push it past what exp2 can represent and produce inf/NaN rows. Lower the accepted range to [2**-64, 2**8] in run_attention_dense_torch and the hipDNN gfx950 dense matcher; there the residue stays far below the lazy-rescale budget. Document both exact outcomes of a row fully masked in its first visited tile (P = 1 when scale * log2(e) < 1, P = 0 otherwise; both cleared by alpha = 0), describe the range as a chosen safety margin, replace stale "softmax argument is always <= 0" comments, and add a scale-1.0 case to the sliding-window no-sinks GPU test so both branches run on hardware.

The ordinary gfx950 dense kernel's exp2 argument fma(s, qk_scale, -m) is, for the row-max element, the rounding residue of s * qk_scale: at most half an ulp of the product. An exhaustive fp32 check over raw scores up to 1e6 in magnitude shows a worst residue of 16 at scale 2**8, enough for P to overflow the fp16 cast before the PV MFMA, and at most 1 at scale 2**4. Lower the accepted range to [2**-64, 2**4] in run_attention_dense_torch and the hipDNN gfx950 dense matcher, and update both tests. Rewrite the matcher guard as an explicit NaN check plus two bound comparisons: the negated conjunction failed clang-tidy readability-simplify-boolean-expr (warnings are errors) in the hipDNN superbuild. Correct comments and the algorithm doc: a row fully masked in its first visited tile gets P = 1 when scale * log2(e) <= 1 and P = 0 above it, and P is bounded by 2**9 rather than 2**8 in the default grid because of the residue.

ALGORITHM.md: name the raw score, qk_scale and the log2-domain score as separate symbols in the notation table; state the premise (scale <= 2**4, |raw score| <= 1e6) behind the fma-residue and P bounds and attribute the residue to rounding m; state the supported scale range in the ABI bullet; and describe how the default and persistent grids' softmax arithmetic differs instead of calling their inner compute byte-identical. run_attention_dense_torch now raises ValueError only for a NaN or infinite scale, and NotImplementedError (NOT_YET_IMPLEMENTED) for a finite scale outside the supported range [2**-64, 2**4]. The sentinel comment no longer claims a non-power-of-two sentinel always overflows: the residue's sign depends on the scale. The kernel IR is unchanged.

Add test_d128_scale_range_edges (scale 2**-64 and 2**4, bf16 and fp16, ordinary and persistent) and a 2**4 case to the sliding-window no-sinks test, so the supported range [2**-64, 2**4] is checked against the fp32 reference at both ends.

Name the C++ mirror of the supported softmax-scale range next to _MIN_SCALE/_MAX_SCALE, derive the out-of-range test cases from the constants, pin the constants to the 0x1p-64F / 0x1p4F matcher literals, and check that both inclusive bounds pass the scale guard.

The matcher guard already declines -inf (it is below 0x1p-64F); the gtest listed +inf but not -inf.

AviralGoelAMD · 2026-10-06T05:09:05Z

@tenpercent thanks for the careful review. I reproduced your host numbers exactly at scale=2**24 (residue 128, exp2(128) = inf).

That case is gone: a later commit (3f37c72) lowered the accepted scale range to [2**-64, 2**4], in both run_attention_dense_torch and the hipDNN matcher. At 2**4 the same inputs (raw score 192) give a residue of about 1.2e-4. A GPU numeric regression at both bounds of the range is in 41d6dc5 (test_d128_scale_range_edges: bf16 and fp16, ordinary and persistent).

You are right that the fused argument is not bounded for every score magnitude. At scale = 2**4 an exhaustive fp32 check gives a residue of at most 1 for |raw score| <= 1e6; beyond roughly 1.2e7 (fp16 P) or 9.3e7 (bf16) it can still overflow. That limit is documented next to _MAX_SCALE in attention_dense.py and in ALGORITHM.md §8.1. I also tried a per-row correction (rounding the max up by one ulp), but for fp16 P it trades the overflow for an underflow, so it is not in this PR. A full guarantee needs per-score rounding as the persistent path does, which this PR does not take because of its cost on the ordinary path.

On your inline question: the matcher guard already declined -inf (it is below 0x1p-64F), and 6956a78 adds it to DeclinesAttentionScaleOutsideSupportedRange.

…ke/gfx950-dense-fp32-softmax-scale # Conflicts: # dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json

Port of #13122 (AICK-2501) onto #12719's dispatcher update. Q feeds the QK MFMA unscaled; the scale is applied in fp32 as fma(s, qk_scale, -m) with the running max in log2 units. Accepted scale range [2**-64, 2**4], mirrored by the hipDNN matcher. Conflict resolutions against #12719: keep its knobs (lazy_rescale_threshold, use_exp2_fast, PV_STEPS/EXP_PER/SCHED_GROUP) with the fix's fp32 math; keep the matcher's required attn_scale_value (DeclinesMissingAttentionScale) and add the fix's range check; regenerate the 32 changed gfx950 dense IR goldens.

Copilot

Copilot review overview

🟡 Changes recommended

The fused exponent argument can still overflow for valid finite inputs with large raw scores.

Review effort: Balanced
Findings: 1 · 2

Open (3)

Unbounded FMA residue can overflow exp2_fast · New
Remove measured speedup from documentation · New
Remove measured speedup from module docstring · New

Resolved since last review (1)

Reject non-finite scales in engine matching

…edups Record the ordinary grid's known limitation next to _MAX_SCALE: nothing bounds the raw score, so the fused exp2 argument grows with the row max and overflows at pre-softmax logits around 1e8 to 1e9, far beyond real inputs. Note the fix to apply if it is ever hit (scale each fp32 score after the MFMA), and point to it from the exp2 sites and ALGORITHM.md. Replace the measured speedups in the two ALGORITHM.md lever rows this branch edits with qualitative wording.

…cstring The module docstring line this branch rewrote kept a measured speedup; keep it qualitative.

…ke/gfx950-dense-fp32-softmax-scale # Conflicts: # dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py # dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx950_numeric.py

Port of #13122 (AICK-2501). Q feeds the QK MFMA unscaled; the scale is applied in fp32 as fma(s, qk_scale, -m) with the running max in log2 units. Accepted scale range [2**-64, 2**4], mirrored by the hipDNN matcher, which keeps develop's absent attn_scale_value = 1.0 and logs a "scale" decline for out-of-range values. The 32 gfx950 dense IR goldens are re-blessed for each llvm flavor.

Copilot

🔵 Needs a closer look

The numerical GPU-kernel rewrite requires final human review, and one persistent-kernel docstring remains inconsistent with the implementation.

1 open finding

Update persistent attention docstring for post-MFMA score scaling · New

3 resolved since last review

Unbounded FMA residue can overflow exp2_fast
Remove measured speedup from module docstring
Remove measured speedup from documentation

🧠 Review effort: Balanced

tenpercent

minor nits, thanks!

ISSUE ID : AICK-2502 ## What was wrong The gfx942 dense prefill attention kernel loses BF16 accuracy at a non-default softmax scale. At `scale = 0.5` and `1.0`, BF16 output exceeds the BF16 tolerance against an fp32 reference, on both the default and the persistent grid. The default scale (`1/sqrt(D)`) is fine, which is why the existing tests never caught it. Cause: the kernel multiplied Q by `scale * log2(e)` in fp32, then rounded the product back to bf16/fp16 before the QK matrix multiply. That rounding error grows with the scale. This is the same defect fixed for gfx950 in #13122 (AICK-2501). ## What changed | Path | Change | |---|---| | `kernels/gfx942/attention_dense.py` | Q feeds the MFMA unscaled. Each fp32 score is multiplied by `qk_scale` right after the MFMA, in the one work-item body both grids share. The `exp2_fast` guard's docstring now says its sweep predates this change, and the VGPR range note is refreshed. | | `tests/test_attention_dense_gfx942_numeric.py` | New GPU tests: D128 at scale 0.5 and 1.0 (bf16 and fp16, default and persistent), and an fp16 case with large equal scores at scale 16. The file's inline fp32 SDPA oracles now share one helper. | | `tests/golden/attention_dense_gfx942_ir_sha256.json` | The gfx942 `attention_dense` IR golden is re-blessed (sha256 and size only). | ## Why it works The QK matrix multiply already accumulates in fp32. Scaling the fp32 scores keeps full fp32 precision, so the scale step no longer adds error. Masking, the row max and `exp2(s - m)` all use the same rounded score, so the row max's exp2 argument is exactly 0 at any scale or score magnitude. No scale restriction is needed, and every scale the kernel accepted before still works. ## How it was validated On an MI325X (gfx942), ROCm 10 / LLVM 23: - New tests fail on `develop` for BF16 and pass with this change: `test_dense_non_default_scale` (D128 at scale 0.5 and 1.0, D64 at 1.0, both grids) and two scale-1.0 sliding-window rows (D128 default grid, D64 persistent grid). The FP16 rows and the large-equal-score case pass on both; they guard against regressions. - The full gfx942 dense numeric file passes. - A scale sweep (BF16 and FP16, both grids) shows the scale-step error is gone: at scale 1.0 the BF16 error drops about 7x, to the level of the no-rounding control (scale = ln 2, where `scale * log2(e)` is exactly 1). The small growth that remains across scales is BF16 output rounding. FP16 also improves. - Register use and scratch were compared against `develop` with the same compiler: no kernel spills. The `develop` FP16 D128 default-grid kernel spilled to scratch; the fixed kernel does not. On the host: the gfx942 CPU tests, the gfx942 IR golden test and pre-commit pass. gfx942 dense has no C++ twin and no strict C++ parity or resource-budget test, so those gates do not apply here. ## Notes - Performance against `develop` (9db0477), on an MI325X (gfx942), ROCm 10 / LLVM 23: batch 1, D128, Sq = Skv = 8192, causal, GQA 128/8 and 32/8, BF16 and FP16, both grids. Cold memory (buffers rotated through at least 1 GiB per round), an A/A control, and 60 paired rounds. In the same run, every timed arm passed an fp32 numeric check at the timed shape. - BF16: slightly slower on the default grid; on the persistent grid slightly slower at GQA 128/8 and unchanged within noise at 32/8. - FP16: faster on the default grid; unchanged within noise on the persistent grid. - Other shapes, head sizes and devices were not measured. Numbers are on the private results page. - The `exp2_fast` default-grid guard for BF16 D128 was tuned before this change and was not re-measured here; its docstring now says so.

Brings the branch up to origin/develop dd76b08. Conflicted files: - dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json Generated file; not hand-merged. The conflict was only in the per-flavor summary counts (develop added cases in #13122 / #13249, this branch adds the 2 nontemporal cases). Re-blessed with the harness bless path: PYTHONPATH=platform/python:library python3 \ platform/tests/instances/rocke_ir_parity_harness.py \ --write platform/tests/golden/rocke_representative_ir_sha256.json Result: llvm20 172 ok / 4 xfail, llvm22 172 ok / 4 xfail, llvm23 176 ok. Verified the regenerated per-case hashes equal exactly develop's hashes plus this branch's nontemporal/* hashes in every flavor (no drift); only the summary counts differ (develop + 2: gfx942 +1, gfx950 +1).

assistant-librarian · 2026-10-09T01:20:44Z

TheRock Submodule Bump Activity

Newest first

2026-10-09 — Included in TheRock via ROCm/TheRock#8860.

The temporary override always built rocjitsu at a pinned rocm-systems commit. Make it retire itself instead: read the rocm-systems commit that therock_ref pins, and build ROCJITSU_MIN_REF only when that pin does not contain it (or the comparison fails). Once TheRock picks up #13122 the job uses TheRock's rocjitsu with no further change here. Build with the same options as TheRock (no LTO or -march=x86-64-v3), so the switch does not change speed. In a local per-test comparison of rocjitsu ad69963, LTO with x86-64-v3 was within 2% of a plain build; the newer commit itself was 1.21x faster than TheRock's d081905. Co-Authored-By: Claude <noreply@anthropic.com>

…tion (#13140) ISSUE ID : AICK-2508 Rebased on current `develop`, which includes #12719 (dispatcher update), #13122 (gfx950 FP32 softmax scaling), and #13138 (gfx942 FP32 softmax scaling). The PR diff contains the LSE feature and its tests/docs; it does not reintroduce those merged scale fixes. Hackathon B8 scope: covers DN-01 (gfx950 dense) and DN-03 (gfx942 dense), including the Python dispatcher path. It does not cover U-02, DK-01, SH-01, gfx1151, hipDNN / ST10 stats-tensor wiring, or ST0 shared cases. ## What was wrong The gfx950 and gfx942 dense prefill attention kernels return only the output `O`. Some callers also need the softmax statistics: LSE, the log of each row's softmax sum. Examples are backward passes and merging partial attention results. The merged #13122/#13138 FP32 scaling fixes let this PR validate LSE at a non-default scale without reintroducing those fixes. ## What changed | Commit | Change | |---|---| | `feat(rocke)` gfx950 | New `emit_lse` opt-in on `Gfx950AttentionDenseSpec`. With it on, the kernel also writes FP32 natural-log LSE into a buffer the caller owns. With it off, the kernel is byte-identical. | | `feat(rocke)` gfx942 | The same opt-in on `Gfx942AttentionDenseSpec`. | | `test(rocke)` | LSE tests run at scale 0.5 for BF16 too, now that both fixes are in. | | `test(rocke)` | Packed varlen (the one mode where LSE changes the `O` store guard) now checks `O` bit for bit against LSE off. ALGORITHM.md names the log2 running max `m_2` and states `LSE = ln2 (m_2 + log2 l)`. | | `feat(rocke)` dispatcher | `AttentionRequest.emit_lse` asks for the LSE output and reports the `lse` feature. Only the dense candidates declare it, so an unpinned LSE request is refused rather than routed to a kernel that can't write LSE. Both dense `_base_spec` builders carry the flag, and `bind_torch` forwards a caller-owned `tensors["lse"]`. | | `test(rocke)` goldens | LSE-on kernels are pinned by IR hashes: 3 gfx950 cases in the representative golden and 3 gfx942 cases (one built through the dispatcher), across LLVM 20/22/23. The gfx942 fixture also re-blesses the 20 LLVM 23 non-LSE hashes that merged scale fix #13138 updated for LLVM 20/22 but left stale; every existing LLVM 20/22 hash is unchanged. | ## Why it works - LSE comes from the running max and sum the kernel already keeps. No extra pass over K or V. - Rows with no allowed key write `O = 0` and `LSE = -inf`, or the sink value. The zeroing is an integer AND on the packed output bits. A float select let LLVM fuse multiply and convert, which changed nonempty FP16 `O` by one bit when LSE was on. - `emit_lse` is an output field (`DENSE_OUTPUT_FIELDS`), so the tuning space never sweeps it. ## How validated | Check | gfx950 (MI355X) | gfx942 (MI325X) | |---|---|---| | Dense numeric file, incl. LSE tests at scale 0.5 for FP16 and BF16 | 222 passed, 0 skipped (incl. dispatcher LSE on all 3 dense candidates) | 98 passed, 0 skipped (incl. dispatcher LSE) | | Strict C++ lowering, builds, dispatch (`ROCKE_CPP_STRICT=1`) | 472 passed, 0 skipped | 744 passed, 1 skipped of 745 on the pre-rebase head; current post-rebase CI pending | | `check_byte_identity.py` (C++ engine built fresh, Python vs C++ `.ll`) | GREEN, 73/73 families | same run | | Host suites | 682 passed (GPU-only and C++-engine skips) | same run | | IR goldens | Post-rebase representative golden is scoped to 3 new LSE cases per flavor; no existing case changed | Post-rebase fixture adds 3 LSE cases per flavor and completes #13138's missing 20-case LLVM 23 re-bless; focused golden test passes locally | | hipDNN matcher + gtest (`TestGfx950AttentionDenseMatchers`) | Built on Linux and Windows by this PR's hipDNN superbuild CI (`HIPDNN_ENABLE_KERNEL_INGESTOR=ON`); `hip_kernel_provider_tests` passed on both | n/a (no gfx942 matcher) | ## Notes - Conflict choices: keep #12719's knobs (`lazy_rescale_threshold`, `use_exp2_fast`, PV step / exp-per-step / sched-group) with the FP32 scale math. In the hipDNN matcher, keep `develop`'s absent scale = 1.0 (#12926) and add the range check with a logged decline. - GPU runs on the develop-based head: ROCm 10 image, LLVM 23 flavor; MI355X job 68043340, MI325X job 68043342. - **Inherited gfx950 ordinary-grid limitation:** adversarial pre-softmax logits around `1.9e8` or larger can make the separate scaled-max rounding overflow the FP16 probability conversion. The persistent grid avoids this by scaling each FP32 score before taking the maximum. Matching that arithmetic in the ordinary grid adds per-score work and may reduce performance, so this LSE PR preserves `develop`'s faster path and its documented `KNOWN LIMITATION`; a correctness/performance change belongs in a separate measured follow-up. - hipDNN does not request LSE yet: its gfx950 matcher still declines graphs that ask for SDPA stats.

…13125) ## Summary Closes the gfx950 forward-attention top-left causal gap in rocKE: the unified kernels silently served top-left as bottom-right (and were unconditionally causal), their dispatch candidates over-claimed a mask they did not implement, and dense mishandled cross-length / ragged shapes. The mask is now the cuDNN/hipDNN band (left = `sliding_window`, right = `right_bound`, alignment = `diagonal_alignment`) computed exactly or rejected with a reason. Part of epic [AICK-2500](https://amd-hub.atlassian.net/browse/AICK-2500). JIRA ID : AICK-2505 ## Risk Assessment High risk. On gfx950 the default request mask (`NO_MASK`) now means full attention in the unified family instead of being served by a causal kernel, and `AttentionRequest` gains fields, a `band` capability feature and new rejections; the kernel changes are large. Default-off IR is pinned byte-identical by the existing platform representative golden (re-run after merging develop), and the new variants' lowered IR is pinned by 26 added cases in that golden. gfx950 hardware parity is green, but no gfx942 / gfx1250 hardware run has been done for the dispatch changes that apply to those arches, and the new bodies are not yet mirrored in the C++ engine (see Known gaps). ## ASIC Coverage Kernel changes are gfx950-only, so gfx950 hardware must pass (done locally, see below). The dispatch-layer changes are arch-wide: `SLIDING_WINDOW` ordinal handling, the new request fields and `band` feature, and gating of the gfx942 dense-pipe and gfx942/gfx1250 tiled paths. gfx942 and gfx1250 keep their legacy `NO_MASK` (causal) body by design, and non-causal bands are rejected there. That behaviour is covered by CPU tests, but a gfx942 GPU-lane run is required before merge to confirm no routing regression. ## Testing Summary - gfx950 hardware parity (MI355X): unified 2D / 3D / scalar kernels against an explicit PyTorch-style reference for top-left, full attention, non-causal and two-sided windows, lookahead, decode, sinks, softcap, D64/D128/D256, GQA, the tiled-lever matrix, and exactly-zero empty rows; dense band matrix through both the spec and the dispatch path, including persistent and ragged cross-attention. - Default-off IR byte-identity at `llvm20`, `llvm22` and `llvm23` via the existing platform harness golden (`rocke_representative_ir_sha256.json`: 32 `attention_dense/gfx950` and 10 `attention/gfx950` cases). Scalar and tiled window / QQ-bias / D256 variants that golden does not cover were compared once against the base commit (31 cases x 3 flavors, 0 differences); no separate library golden is added. - Dispatch wiring tests: distinct kernel names / cache keys for every unified candidate, band normalization, capability features, off-gfx950 gates. - Empty-row behaviour cross-checked against PyTorch 2.10 SDPA (fully masked rows return 0). ## Testing Checklist - [x] rocKE library CPU lane - `python library/tests/run_all.py` - Status: Passed (2314 passed; 3 failures, all pre-existing and identical on clean develop: `test_tiled_2d_shipped_geometries_compile_and_fit_budget` and the gfx942 dense `llvm23` golden) - [x] rocKE library GPU lane - `python library/tests/run_all.py --gpu-only` - ASICs: gfx950 - Status: Passed (536 passed, on the branch merged with develop) - [x] rocKE installed golden gate - `python platform/tests/rocke_installed_golden_test.py` - Status: Passed (llvm20, llvm22, llvm23), re-run after merging develop (including the fp32 softmax-scale and dispatcher-restructure changes) - [x] Pre-commit hooks - `git commit` - Status: Passed - [x] Review fixes validated - GPU lane 544 passed (gfx950); CPU lane 2317 passed with the same 3 failures as clean develop; platform golden gate PASS at llvm20, llvm22, llvm23 with the 26 added cases - [ ] rocKE library GPU lane on gfx942 - `python library/tests/run_all.py --gpu-only` - ASICs: gfx942 - Status: Not run - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - Review follow-up: top-left is no longer folded into bottom-right on the unified family when the request's `seqlen_q == seqlen_k` (batch maxima there; the kernels derive the diagonal per sequence); dense keeps that shortcut. Dense q-blocks entirely past the keys no longer read other batches' tokens (start tile clamped when rows can be empty; the dense band matrix now runs at `batch=2`). The new variants' IR is pinned in the platform golden, and a GPU case covers request -> dispatch -> launch on a mixed-length cohort. - Merged `develop` (#12719 dispatcher restructure, #13122 fp32 softmax scale): the dense band moved into `gfx950_dense.py` (`_dense_band`, `_is_ragged`), the unified AttentionSpec fields into `gfx950_unified.py`, and `right_bound` is classified as a dense problem field. Under #13122's finite mask sentinel a row with no visible key is no longer NaN but averages V, so the dense kernels now zero such a row at the epilogue when its running max is still at the sentinel. - Unified 2D tiled, 3D split-KV and scalar kernels: new `causal_top_left` and `right_bound` problem/spec fields. The window anchor stays `d - k < W` and the upper compare becomes `k <= d + R` (skipped when unbounded); tile bounds, the sliding-window tile end, the peel split and the `mask_limit` threshold all follow `R`. `R = 0` lowers byte-identically. - Dense kernel (default and persistent grids): `right_bound`, non-causal windows, lookahead, ragged cross-attention, and a zero-output guard for rows with no visible key (previously stale data from `exp2(-inf - -inf)`), emitted only when such rows can exist. - Dispatch: `AttentionRequest.right_bound` / `diagonal_alignment`, a `band` capability feature, `SLIDING_WINDOW` read as the band trio (rejected when it names no bound), distinct kernel names (`_tl`, `_rbu` / `_rb<R>`, `_sw<W>`) and cache keys; unsupported combinations (e.g. ALiBi/QQ-bias with a right bound, dense bottom-right window at `sq != sk`) are rejected with a reason. - Parity harness: `ref_paged_attn` band reference and `--set topleft|nomask|band`; gfx950 benchmarks that relied on the default mask now pin bottom-right causal. - Known gaps (follow-ups): - **C++ mirror (tracked follow-up):** the `causal_top_left` / `right_bound` bodies of the tiled 2D / 3D and unified (scalar) kernels are not yet mirrored in `platform/cpp/instances/`, and there are no `library/tests/parity/*_emit` pairs for them (e.g. `attention_tiled_2d_ctx_prologue.cpp` still hard-codes `context_len = seq_len - cur_batch_q_len`). The default mask is unchanged, so Python/C++ byte-identity still holds; until the mirror lands the engines diverge on the new variants, whose Python IR is pinned only by the platform golden cases. Dense has no C++ mirror. - On arches with no top-left kernel (gfx942 / gfx1250) the legacy equal-length shortcut remains: top-left is served by the bottom-right body, correct only when every sequence has `q_len == kv_len`. gfx942 / gfx1250 tiled kernels remain causal-only; dense window/lookahead at bottom-right alignment with `sq != sk`; pre-existing 2D tiled ALiBi + sliding-window NaN at bottom-right. 🤖 Generated with [Claude Code](https://claude.com/claude-code) [AICK-2500]: https://amd-hub.atlassian.net/browse/AICK-2500?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: Claude <noreply@anthropic.com>

AviralGoelAMD requested review from a team as code owners October 5, 2026 22:50

github-actions Bot added project: hip-kernel-provider rocke-lib: attention Touches rocKE library attention/FMHA kernels rocke: platform Touches rocKE platform labels Oct 5, 2026

AviralGoelAMD marked this pull request as draft October 5, 2026 23:01

tenpercent reviewed Oct 5, 2026

View reviewed changes

Comment thread ...provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionDenseMatchers.cpp Resolved

tenpercent requested a balanced review from Copilot October 5, 2026 23:02

Copilot started reviewing on behalf of tenpercent October 5, 2026 23:03

View session

Copilot AI reviewed Oct 5, 2026

View reviewed changes

Comment thread .../hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseNative.cpp Resolved

assistant-librarian Bot added the organization: ROCm label Oct 5, 2026

AviralGoelAMD added 2 commits October 5, 2026 19:26

github-actions Bot added the documentation label Oct 5, 2026

AviralGoelAMD added 3 commits October 6, 2026 00:38

AviralGoelAMD marked this pull request as ready for review October 6, 2026 04:46

test(hipdnn): cover -inf in the gfx950 dense scale decline test

6956a78

The matcher guard already declines -inf (it is below 0x1p-64F); the gtest listed +inf but not -inf.

Merge remote-tracking branch 'origin/develop' into users/avirgoel/roc…

549b905

…ke/gfx950-dense-fp32-softmax-scale # Conflicts: # dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json

AviralGoelAMD mentioned this pull request Oct 6, 2026

fix(rocke): apply gfx942 dense softmax scale in fp32 #13138

Merged

AviralGoelAMD mentioned this pull request Oct 6, 2026

feat(rocke): add optional LSE output to gfx950 and gfx942 dense attention #13140

Merged

tenpercent requested a balanced review from Copilot October 6, 2026 16:02

Copilot started reviewing on behalf of tenpercent October 6, 2026 16:03

View session

Copilot AI reviewed Oct 6, 2026

View reviewed changes

AviralGoelAMD added 3 commits October 6, 2026 13:16

docs(rocke): drop measured speedup from the gfx950 dense exp2_fast do…

9cb34cd

…cstring The module docstring line this branch rewrote kept a measured speedup; keep it qualitative.

vidyasagar-amd approved these changes Oct 6, 2026

View reviewed changes

AviralGoelAMD enabled auto-merge (squash) October 6, 2026 19:27

AviralGoelAMD requested review from a team and tenpercent October 7, 2026 15:43

tenpercent requested a balanced review from Copilot October 7, 2026 15:56

Copilot started reviewing on behalf of tenpercent October 7, 2026 15:57

View session

Copilot AI reviewed Oct 7, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py Resolved

tenpercent reviewed Oct 7, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md Resolved

tenpercent approved these changes Oct 7, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py Resolved

AviralGoelAMD merged commit 89cc3af into ROCm:develop Oct 7, 2026
96 checks passed

AnaghaRaoAMD mentioned this pull request Oct 8, 2026

feat(rocke): top-left causal and attention band coverage on gfx950 #13125

Merged

5 of 7 tasks

Conversation

AviralGoelAMD commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What was wrong

What changed

Why it works

How it was validated

Notes

Uh oh!

therock-pr-bot Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

tenpercent commented Oct 5, 2026

Uh oh!

AviralGoelAMD commented Oct 6, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Uh oh!

Uh oh!

tenpercent left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

assistant-librarian Bot commented Oct 9, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

AviralGoelAMD commented Oct 5, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 5, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 5, 2026 •

edited

Loading