feat(rocke): optional LSE output on unified tiled 2D attention by AviralGoelAMD · Pull Request #14616 · ROCm/rocm-libraries

AviralGoelAMD · 2026-10-28T21:10:09Z

ISSUE ID : AICK-2685

Builds on #14542 (fully masked leading tiles keep the running max), which is now in develop; this PR's diff is the LSE work only.

What was missing

Unified paged attention could return the softmax statistics (LSE, the log of each row's softmax sum) only on the split-KV path (#13955). backend="auto" sends prefill to the tiled 2D path, so an LSE request for a prefill shape was refused. This is a prerequisite for SDPA with stats output on unified prefill; the hipDNN binding is separate work.

What changed

emit_lse now works on the tiled 2D path on gfx942 and gfx950. Python requests and runs it; the C++ engine builds the same kernels byte-identically, but cannot request LSE yet (known gap 7).

Path	With `emit_lse`
Tiled 2D, gfx942 / gfx950	writes LSE (all three output paths)
gfx942 4-warp GQA cohorts: D256 bf16 prefill, no sliding window, block size 16/32; D128 fp16/bf16 sliding-window prefill, block size 16/32/64	not used. D256 bf16 at block size 16 and the D128 sliding-window shapes take the default 2D builder, which the code documents as slower for them (for D128 sliding window, most at block size 16); not measured, AICK-2833. D256 bf16 at block size 32 exceeds the default 2D builder's LDS budget, so `auto` and dispatch run it on split-KV with LSE (as D256 fp16 at block size 32)
gfx942 2D CUDA-graph fast path	skipped (its cache key has no `emit_lse`), so small LSE prefills pay the normal launch cost
Split-KV	unchanged, except a NaN query row now gets NaN LSE instead of `-inf` (same rule as 2D)
Scalar fallback, sinks, other arches	refused with a reason

Kernels: one FP32 value per query token and head, [total_q, Hq, 1]; rows with no visible key get -inf and O = 0. LSE-off kernels are unchanged.
Host and dispatcher: the 2D gate, spec, cache key, kernel signature, runner and dispatcher candidates carry emit_lse. The 4-warp builder refuses an LSE spec.
C++ engine: mirrors the LSE-on 2D kernels and the split-KV reduce. The engine bindings now pass every shared 2D spec field, and the gfx942 split-KV binding builds the reduce as well.

Why it works

The kernels already hold the final running max and sum per row; the store converts them to natural log. On the striped 16x16 output path the store comes after the O stores: placed before them, LLVM changed how FP16 O was rounded relative to the LSE-off kernel.

How it was validated

gfx942 and gfx950, ROCm 10, against an FP32 paged reference: every 2D output path, fp16/bf16, D64/128/256, ragged batches, empty rows, NaN rows, sliding window, softcap, additive bias (gfx950), the graph bypass and the dispatcher bind path, plus the gfx942 D256 block-size-32 split-KV handover. All pass. O is bit-identical to the same kernel with LSE off; on the gfx942 4-warp cohort shapes the LSE-on kernel differs from the shipped LSE-off one and matches the reference within tolerance.
Mutation checks: removing any LSE store, the NaN guard, the row mapping or the graph bypass fails the tests.
LSE-off: no existing IR golden or differential config changes; check_byte_identity GREEN at llvm20/22/23/24.
Host suites, every library test file run separately against develop at the merge base: the same 39 failures on both sides (environment: no torch, test fakes, clang options; plus three defects already on develop), no test that passes on develop fails here. The platform runner gives an identical failure set too.
No CI lane runs the GPU tests; they were run by hand on the nodes above.
ROCKE_BACKEND=both on both GPUs, with the C++ engine built from this branch's sources: every GPU test passes with the C++ and Python lowerers agreeing on each kernel. The C++ kernel builders are checked on the host by check_byte_identity.

Notes

Admitted but not yet run on silicon with LSE: FP8 KV, ALiBi, QQ-bias, block size 32/64, the hipcc compile backend, gfx950 top-left / right-bound masks, i64 KV addressing, and head ratios that leave a tile partly empty (AICK-2833).
On the gfx942 4-warp cohorts, the shipped LSE-off kernel writes mean(V) (NaN when the KV length is 0) instead of 0 for a row with no visible key; LSE-on requests avoid it. Tracked with the other follow-ups in AICK-2833.
Additive bias with LSE stays Python-only on gfx950 (known gap 5).

Related: additive bias on the gfx950 tiled 2D kernel is Python-only (#13120); #14440 (open) mirrors it into C. Until then, additive bias + LSE is built by Python only.

…ks 8-11)

…e-bless goldens

…2D C++; register the masked-tile tests

…keys

… kernels

The empty-segment guards compared the segment sum and the final denominator with an ordered greater-than, which is false for NaN, so a NaN sum was dropped as an empty segment and a NaN denominator gave O = 0. Treat only an exact zero as empty so NaN propagates as before.

Run every epilogue branch at page sizes 16 and 64, using the 1x-block gfx942 variants at page size 64 where the 2x-block ones exceed LDS, assert each geometry still gives a row a fully masked leading tile, and correct the fp16/bf16 P underflow thresholds in the docstring.

…erage Scalar 3D: write NaN into one segment sum of a row that sees keys between the segment and reduce launches; that row's O must be NaN and every other row must match the reference. An ogt-style empty-segment guard would skip the NaN segment and leave the row finite. Tiled 2D: run the pinned epilogue branches at head sizes 64, 128 and 256 (256 skipped, with the canonicalizer's reason, where the variant exceeds the LDS budget), and add one auto-route D64 bf16 case per arch whose production kernel is the multi-warp default builder.

Build the scalar split-3D reduce in the representative IR golden (scalar_reduce_{fp16,bf16}_d128_seg8) and in the attention_unified C/Python parity emitter pair (configs 6-7), so check_byte_identity compares the reduce between the C++ and Python engines. Goldens gain only the new keys. Also fix a gfx950 comment that named m_new as the state handed to the PV bucket; it is m_carry, the raw running max.

…ry mask The scalar 2D epilogue guarded acc * rcp(l) only for non-default masks, on the claim that the default mask never yields a row without keys. Bottom-right causal with q_len > kv_len does: its leading rows see no key, l is 0 and O was NaN. The C++ twin had no guard at all. Guard every mask, and test l == 0 (fcmp oeq) instead of l > 0 (fcmp ogt) so a NaN sum still reaches O instead of being stored as a keyless 0. Mirror the same ops in the C++ engine and re-bless the representative and differential goldens for the scalar 2D cases. Tests: a default-mask q_len > kv_len case whose keyless rows must be 0, and a NaN query-row probe that must give NaN O under the default and top-left masks.

The tiled 2D epilogue on gfx942 and gfx950 stored 0 for any row whose softmax sum failed 'l > 0'. A NaN sum fails that test too, so a corrupt row came out as a finite 0. Store 0 only for an exact-zero sum (fcmp oeq), the rule the scalar 2D epilogue and the scalar reduce use. Rows with a finite nonzero sum are unchanged. Python builders and C++ twins change in step; the fastKV register-P variant inherits the change. Also correct the scalar 2D and scalar 3D comments about keyless rows, and re-bless the representative and differential goldens of the tiled 2D families.

Add a NaN-query-row case per tiled 2D epilogue branch, state what each mask case of the scalar 2D NaN test pins, and drop the unused segm_max argument of the scalar 3D workspace hook.

…arry Resolve the differential golden conflict with #14501 by regeneration: take develop's file, then re-record the attention_tiled_2d and attention_unified families (ll llvm20/22/23 + ir canonical, python-only). Other families keep develop's entries.

) into unified-2d-lse Plan Task 0.6: the LSE work builds on the raw-max carry and the exact-zero keyless guard. Brings develop up to a3b5d0c as well.

AICK-2685 Task 1. Off by default: the kernel name, spec hash, tuning id and LSE-off IR are unchanged. The kernel name gains an 'lse' token only when on. emit_lse is kept out of the 2D kernel defaults fingerprint (the reduce already excludes it), so stored 2D pins still resolve. New host test file test_attention_unified_2d_lse_ir.py with the LSE-off IR hashes recorded after the #14542 merge (llvm20).

AICK-2685 Task 2. With emit_lse on, the gfx950 build_unified_attention_2d_tiled takes a trailing lse_ptr (after the additive-bias params) and stores one FP32 natural-log LSE per valid row, from one writer lane per row: transposed 32x32 (slot 0, lanes 0-31), non-transposed 32x32 (16 slots, lane_col32 == 0) and striped 16x16 (4 slots, lane_col == 0). row_has_mass is the new _lse_store.lse_row_has_mass, !(l == 0) built as fcmp one OR fcmp uno: the exact inverse of the epilogue's oeq keyless guard, so a keyless row writes -inf and a NaN row keeps NaN. lnot(oeq) is avoided because lower_hip lowers i1 NOT to ~bool, which is always true. LSE off: IR byte-identical (pinned llvm20 hashes), no param, no store.

The GPU LSE tests only use the LLVM path. Lower a tiny kernel through lower_kernel_to_hip (no prologue), compile it with the host C++ compiler and check l = 0, -0, 1, -3.5, NaN, inf -> 0, 0, 1, 1, 1, 1. Swapping in lnot(oeq) (lower_hip emits ~bool) or l > 0 (hides NaN) fails it.

AICK-2685 Task 3. Three new representative cases, attention/gfx950/2d_lse_{transposed32,mfma32,striped16}_bf16_d128, one per LSE store site, recorded for llvm20/22/23. No existing case hash changes; only the per-flavor summary counts grow by 3. check_ir_validity and check_arch_domain stay GREEN with the new cases.

AICK-2685 Task 4. Same contract as gfx950: kw_only bool, checked in __post_init__, 'lse' kernel-name token only when on. The kernel name, spec hash, tuning id (the defaults fingerprint already drops emit_lse) and LSE-off IR are unchanged; the gfx942 LSE-off hashes are pinned (llvm20).

AICK-2685 Task 5. With emit_lse on, the gfx942 build_unified_attention_2d_tiled takes a trailing lse_ptr (after qq_bias_stride_0; gfx942 has no additive bias params) and stores one FP32 natural-log LSE per valid row from one writer lane per row: transposed 32x32x8 (slot 0, lanes 0-31), non-transposed 32x32x8 (16 slots, lane_col32 == 0, before the Acc_lds barriers) and striped 16x16 (REGS_PER_LANE slots, lane_col == 0). row_has_mass is lse_row_has_mass (!(l == 0), NaN kept), like gfx950. LSE off: IR byte-identical (pinned llvm20 hashes). New IR tests cover the one-slot transposed state, the q-major grid and two atoms per warp, and pin that production gfx942 specs land on branches A and C.

AICK-2685 Task 6. Three new representative cases, attention/gfx942/2d_lse_{transposed32,mfma32,striped16}_bf16_d128, one per LSE store site, recorded for llvm20/22/23. No existing case hash changes (no gfx942 2D or 4-warp drift); only the per-flavor summary counts grow by 3. check_ir_validity and check_arch_domain stay GREEN with the new cases validated.

AICK-2685. The C++ twins of the gfx942 and gfx950 unified split-KV reduce now mirror Python's emit_lse (merged in #13955 as Python-only): an emit_lse spec field (default off), the 'lse' kernel-name token, a trailing lse_ptr after seq_lens_ptr, and the lane-0 LSE store after the combine pass. The binding reads emit_lse into the derived reduce spec. New shared header instance_lse_store_internal.h: C twins of _lse_store.lse_store / lse_row_has_mass and natural_lse_from_log2_stats, one builder call per statement in Python order. Parity emitters gain LSE-on reduce configs (gfx942 18-19, gfx950 16-17). check_byte_identity GREEN at llvm20/22/23; dropping or reordering a C LSE op turns exactly those configs DRIFT. LSE off emits the same IR as before.

AICK-2685. The C++ twin of the gfx942 build_unified_attention_2d_tiled now mirrors emit_lse: the shared rocke_attention_tiled_2d_spec_t gains emit_lse (default off; read and written by the bindings), the 'lse' name token, the trailing lse_ptr after qq_bias_stride_0, m_final, and the three epilogue store sites with the same writer lanes as Python. The C fastkv_regp wrapper builds gfx950 kernels through this gfx942 C epilogue, so branch A's writer guard is now land(op_mask_t, lane_half32 == 0) in both Python builders (the gfx950 order). Only the attention/gfx942/2d_lse_transposed32 representative golden moves. Parity emitter configs 11-15 cover transposed32 (16 and 1 state slot), mfma32, striped16 and striped16 two-atom + q-major grid. check_byte_identity GREEN at llvm20/22/23 for gfx942_attention_tiled_2d and _stress; reordering a C LSE op turns the matching configs DRIFT.

AICK-2685. The C++ twin of the gfx950 build_unified_attention_2d_tiled now mirrors emit_lse: trailing lse_ptr after qq_bias_stride_0, the 'lse' name token (also in the fastkv_regp wrapper's own name list), m_final, and the three epilogue store sites. A new hoist_qh_r ctx array keeps the per-slot head index (the C hoist_q_head slot is reused for the ALiBi slope); filling it adds no IR. Parity emitter configs 56-59 cover transposed32 with scalar state, mfma32, striped16 GQA-16 and striped16 + ALiBi; fastkv_regp config 6 runs the wrapper with LSE on. check_byte_identity GREEN at llvm20/22/23; reordering a C LSE op, or dropping the wrapper's 'lse' token, turns the matching configs DRIFT. The C twin has no additive bias, so bias + LSE stays Python-only.

…figs Regenerated with run_diff.py --python-only --record-golden --only attention_tiled_2d,attention_tiled_3d (ll llvm20/22/23 + ir canonical). Adds the 14 new configs per kind (fastkv_regp 6, gfx942 2D 11-15, 3D reduce 16-19, gfx950 2D 56-59); no existing entry changes. --check-golden: GOLDEN OK.

cgmillette · 2026-10-31T00:11:29Z

Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max

PR #14616 — optional LSE output on unified tiled 2D attention (AviralGoelAMD, AICK-2685) — 🟡 Follow-up: clean; land order pending

Prior item	Status
Author (body, 10-28): untested-on-silicon LSE paths → AICK-2833; no C request (gap 7)	🟡 DEFERRED • body unchanged; unsupported paths refuse with a reason (read) • acceptable for an opt-in output; CI covers none of it (N1)
CongMa13 (10-29) APPROVED at `254c35e8`	🟡 Stale • predates 10 own commits (`34c793f9`…`13619c6c`) + 7 merges • still counts
Copilot (10-29) ×3: platform imports `library`; gfx942 reduce drops `waves_per_eu`; tracker IDs	🟢 FIXED `34c793f9`, `3b1c6aea`, `41ac07e6`, `84a86725` • hold after the merges: no `library` import, wpe copied (`rocke_engine_attention.cpp:1204`), no tracker IDs (read) • wpe parity cfg is now 22; byte identity GREEN (ran)
Copilot (10-29) ×3: `emit_lse` grows 3 public spec structs, no ABI bump	🟢 FIXED `6683a604`, renumbered `d8b32561` • append-only: `emit_lse` last in all 3 structs (ran: diff vs develop) • develop is 5; takes develop + 1 when it lands (order: `rk abi`) • our nit stands: line 6 says they grow; x86-64 sizes don't (ran: probe)
AnaghaRaoAMD (10-29) `attention_unified.py`: gfx942 D256 bs32 LSE has no path	🟢 FIXED `dd30ab54` • `d8b32561` ports it to develop's `select_path(arch)` (read) • routing tests pass: merge (ran), CI at head (read)
AnaghaRaoAMD (10-29) ×3 tests: no CTest lane; domain golden; hash pins	🟢 FIXED `439874a0`, `774d19e9`, `f9cd0c41` • both files ran on both gfx94X shards: 53 ✓ each (read: CI logs) • JSON golden and `_OFF_SHA256` still gone (read)
AnaghaRaoAMD (10-29) APPROVED at `13619c6c`	🟡 Stale • merges `53d7efdd`, `d8b32561` re-resolve the patch onto develop • golden shards, ABI renumber, parity re-index; still counts
Our 10-29 N1: LSE GPU numerics never run in CI	🟡 OPEN • file unchanged; skipped on both gfx94X shards, no torch (read: CI logs) • author acts; not blocking
Our 10-29: land after #14542, keep this PR's side	🟡 OPEN • #14542: approved, auto-merge on, CI running • then merge develop: 7 files conflict; this PR's side = today's tested tree (ran: merge-tree)
Bots / CI	🟡 Running • fails: gfx94X hipkernelprovider 2/2 — `test_ir_cases_match_golden_sha256` timeout, develop's #14733; passed on 1/2 → re-run • fails: HipDNN Windows Superbuild — develop's FlyDSL census break; fix #14801; not required • running: Windows gfx110X hipkernelprovider ×2; 37 ✓ • merge: blocked; auto-merge off (author)
Our checks	🟢 Clean • merge vs `1e9db890`: clean, 3 behind (ran) • ctest 35/35; byte identity GREEN ×3; parity 2 ✓ ×3 (ran) • goldens PASS ×3; differential OK, 0 missing (ran) • PR tests 447 ✓ / 84 skip; compile, rules 1–3 clean (ran)

Net-new: none.

Utility: 🟢 Positive — unified prefill can return LSE; LSE-off kernels unchanged; GPU numerics author-run only.

Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted

Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with develop (build + ctest, byte identity, parity, goldens, PR tests, compile), ABI header diff + sizeof probe, #14542 land-order merge-tree, CI logs at head. Flags whether a closer look is needed; not an approval or rejection.

Develop now carries the base PR as a squash commit whose rocKE tree equals the base-PR head this branch already merged. Keep this branch's side of every conflict: the merged tree is identical to the previous head.

cgmillette · 2026-10-31T02:42:22Z

Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max

PR #14616 — optional LSE output on unified tiled 2D attention (AviralGoelAMD, AICK-2685) — 🟡 Follow-up: develop merge clean; CI running

Prior item	Status
Our 10-29: land after #14542, keep this PR's side	🟢 FIXED `eb0d43ed` • develop merged after 14542 landed; 0 behind (ran) • same tree we last checked; no 14542 hunk lost (ran) • LSE reads the raw running max, ln 2 once (ran: IR, 16 kernels)
Our 10-29 N1: LSE GPU numerics never run in CI	🟡 OPEN • file unchanged; CI has no torch, so it skips (read) • author acts; not blocking
Author (body, 10-28): untested-on-silicon LSE paths → AICK-2833; no C request (gap 7)	🟡 DEFERRED • Notes unchanged; unsupported paths refuse with a reason (read) • body's "diff is the LSE work only" holds (ran)
CongMa13 (10-29) APPROVED at `254c35e8`	🟢 APPROVED • 10 own commits + 9 merges since • still counts on develop
Copilot (10-29) ×3: platform imports `library`; gfx942 reduce drops `waves_per_eu`; tracker IDs	🟢 FIXED `34c793f9`, `3b1c6aea`, `41ac07e6`, `84a86725` • unchanged since our last check (ran: tree compare) • byte identity GREEN ×3 (ran)
Copilot (10-29) ×3: `emit_lse` grows 3 public spec structs, no ABI bump	🟢 FIXED `6683a604`, renumbered `d8b32561` • append-only: `emit_lse` last in all 3 (ran: diff vs develop) • develop is 5; takes develop + 1 when it lands (order: `rk abi`) • our nit stands: history says sizes change; x86-64's don't (read)
AnaghaRaoAMD (10-29) `attention_unified.py`: gfx942 D256 bs32 LSE has no path	🟢 FIXED `dd30ab54` • routing tests pass at head (ran)
AnaghaRaoAMD (10-29) ×3 tests: no CTest lane; domain golden; hash pins	🟢 FIXED `439874a0`, `774d19e9`, `f9cd0c41` • LSE IR file passes at head: 48 ✓ (ran) • its gfx94X CI shards still running
AnaghaRaoAMD (10-29) APPROVED at `13619c6c`	🟢 APPROVED • merges `53d7efdd`, `d8b32561`, `aceb1506`, `eb0d43ed` since • still counts on develop
Bots / CI	🟡 Running • fails: HipDNN Windows Superbuild + Summary — develop's FlyDSL census break; fix #14801; not required (read: log) • running: gfx94X hipkernelprovider ×2, Windows gfx110X build; 30 ✓ • merge: blocked until Multi-Arch summary; auto-merge on
Our checks	🟢 Clean • merge = head, 0 behind `809c17de`; ctest 35/35 (ran) • byte identity GREEN ×3; parity, goldens pass; differential 0 changed (ran) • tests 400 ✓ / 2 skip; compile nothing new; rules 1–3 clean (ran)

Net-new: none.

Utility: 🟢 Positive — unified prefill can return LSE; LSE-off kernels unchanged; GPU numerics author-run only.

Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted

Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with develop (build + ctest, byte identity, parity, goldens, PR tests, compile), 14542 patch survival and tree compare, LSE-on IR probe (16 kernels), ABI header diff, CI logs at head. Flags whether a closer look is needed; not an approval or rejection.

## Summary Fix forward the rocKE installed golden failure introduced when the LLVM 24 representative goldens landed without the attention LSE cases subsequently merged into `develop`. This re-records only the LLVM 24 attention entries. Reverting the LLVM 24 qualification would remove intended test coverage, so updating the stale golden is the safer fix-forward. ## Risk Assessment Risk level 1. This changes generated regression-test metadata only; it does not change kernel emission or product behavior. LLVM 20, LLVM 22, and LLVM 23 entries are unchanged. ## Related - Work tracking: #14815 - Related PRs: #14692, #14616 Fixes #14815 ## Device / Architecture Coverage No GPU validation is required because this is a host-only representative-IR golden update. The affected cases cover `gfx942` and `gfx950`, and the installed golden gate evaluates every recorded LLVM flavor without requiring local hardware. ## Testing Summary - Reproduced by the failure recorded in #14815. - Confirmed the refreshed installed golden passes for LLVM 20, LLVM 22, LLVM 23, and LLVM 24. - Confirmed the diff contains only six new LLVM 24 LSE cases and two refreshed LLVM 24 `reduce_lse` entries. ## Testing Checklist - [x] Installed representative-IR golden gate — `cd platform && PYTHONPATH="$PWD/python:$PWD/../library" python tests/rocke_installed_golden_test.py` — Status: Passed - [x] Golden flavor-diff audit — LLVM 20–23 unchanged; LLVM 24 has six additions and two expected updates - [ ] PR CI — GitHub PR checks — Status: Pending ## Flags / Guardrails None. This does not alter runtime behavior. ## Adjacent Tests Considered The byte-identity and GPU numeric lanes are not required because emitted IR and kernel behavior are unchanged. The installed representative-IR golden gate is the lowest-level test that directly reproduces and validates this regression. ## Technical Changes - Add the six missing LLVM 24 attention LSE representative cases for `gfx942` and `gfx950`. - Refresh the two stale LLVM 24 `reduce_lse` hashes for `gfx942` and `gfx950`. - Preserve all LLVM 20, LLVM 22, and LLVM 23 golden entries unchanged.

#14616 landed with ROCKE_ABI_VERSION 6, so this PR's rocke_gemm_trait_spec_t layout change takes 7: abi.h keeps develop's history through 6 and adds ours as 7, and BINARY_ABI in portable_ir/src/abi.py is 7 to match. The differential golden takes develop's (with #14616's gfx950_attention_tiled_2d re-record), then re-records moe_fused_mega_fp8[0..5] and gfx950_attention_tiled_2d[27], [31] at llvm20/22/23/24 for this PR's LDS alignment change: the same 32 hashes as before, nothing else. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Conflicts: platform/CMakeLists.txt (both sides extended the attention CPU pytest comment; combined) and known_gaps.md (develop's #14616 rewrote the LSE gap table after landing its own C LSE store; took develop's). #14616 added rocke/instance_lse_store_internal.h, a C twin of _lse_store.py's lse_store and natural_lse_from_log2_stats, now used by the unified 2D and split-KV reduce builders. This branch had a second copy of the same two functions (rocke_attention_dense_lse_store and helpers/attention_fwd_ext.cpp). Dropped this branch's copies; the dense LSE parity emitter calls rocke_lse_store instead. dense_row_has_key and zero_keyless_row, which develop does not have, stay in instance_attention_dense_lse_internal.h. The attention_dense_lse pair stays byte-identical (8 configs GREEN).

Merge develop 0e707b6 (AICK-2683). Conflicts: - library/tests/parity/gfx942_attention_tiled_3d_emit.{py,c}: develop (#14616) added configs 20-22 (reduce emit_lse, waves_per_eu decode). Our strided decode configs move from 20-29 to 23-32; the C strided threshold is idx >= 23; make_reduce_spec takes idx and sets emit_lse for 20-21; the Python reduce is reduce_spec_for_segment plus emit_lse for _REDUCE_LSE. Both emitters keep a single build_segment / make_reduce_spec. - platform/CMakeLists.txt: union of both CPU test lists. - golden/llvm_gfx_all.json: develop's file plus our 10 configs recorded as 23-32 in ir_canonical and ll llvm20/22/23/24 (python-only record into a throwaway copy; the hashes equal our old 20-29 ones). - integration-tests/SUPPORT_MATRIX_OVERVIEW.md: regenerated with feature_matrix.py from the merged claims and catalogue. Silent trap fixed: develop's prove_parity_binding requires attention_tiled_parity_configs.GFX942_ATTENTION_TILED_3D to cover every emitter config; the 10 strided configs are appended as rows 23-32.

assistant-librarian · 2026-11-03T11:07:44Z

TheRock Submodule Bump Activity

Newest first

2026-11-03 — Included in TheRock via ROCm/TheRock#9698.

…se MMA instruction (#14328) ## Motivation rocKE should know every production target its MMA catalog names, and should be able to issue every dense MMA instruction each of those targets has, in both engines. Today some targets are unknown, several MMA lowerings print a fixed operand type around any value, and integer WMMA signedness is hard-coded to signed. This PR makes the ISA backends complete for the dense instructions: - every op is described as data; - every target compiles for itself where its LLVM flavor accepts the name; - operand signedness comes from the operand dtype. ISSUE ID : AICK-2571 ## Technical Details The 18 feature commits below, then the review fixes listed under **Changes after review** (and merges of develop). Each builds and passes its own tests on its own, and the Python and C++ engines emit byte-identical LLVM IR at every commit. 1. **Delete the duplicate f16 32x32x16 MFMA handler.** 2. **Resolve each LLVM flavor's intrinsic declarations in one place**, in both engines. Each declaration generation (LLVM 22, 23, 24) holds only its own changes and inherits the older ones. 3. **Describe matrix intrinsic calls as data.** `MATRIX_CALLS` lists each op's intrinsic and operand forms. An op lowered through it is refused on a target whose catalog does not list it. 4. **Let a target row take its facts from another row.** Every member of a generic target gets a pointer row. 5. **Add the missing gfx targets.** - gfx908 gets a full fact row; its occupancy is measured with `llc`. - gfx1100-1103 and gfx1150-1154 point to gfx11-generic; gfx1200 points to gfx12-generic. - Every new target's arch-domain cells are measured on each flavor's compiler. 6. **Compile each target for itself where the compiler that runs accepts it.** Otherwise it compiles for the generic its row points at; if neither works, the compile fails loudly. The choice follows the LLVM major of the compiler that will run (see Changes after review). 7. **Add CDNA f32, i8, f64, bf16 and mixed fp8 MFMA ops.** gfx942's f64 ops keep the lowering #13794 added; gfx90a and gfx950 list the f64 16x16x4 atom. 8. **Add gfx12 fp8, bf8, i8 and 16-bit-output WMMA ops.** 9. **Add gfx12 int4 WMMA ops, and derive integer WMMA signedness from the operand dtypes.** The unsigned rows are named `_u8_i8`, `_i8_u8` and `_u8_u8` (and the u4 equivalents). The gfx11 integer handler moves into `MATRIX_CALLS`. 10. **Add gfx11 16-bit-output WMMA ops.** 11. **Describe MMA instructions whose result type differs from the accumulator.** - A row may carry a result dtype `dst`, the name agreed on #11642. `MmaOp.dst_dtype` defaults to the accumulator dtype, and catalog queries take `dst_dtype`. - The C++ op struct grows, so the binary ABI goes from 6 (develop) to 7 in `abi.h` and in Python's `BINARY_ABI`. - Adds the gfx1250 bf16-result op. 12. **Add gfx1250 16-bit-output, fp8-to-f16 and i8 WMMA ops, and wire f32 16x16x4.** 13. **Add the gfx1250 f8f6f4 and f4 32x16x128 WMMA ops.** 14. **Add the gfx950 f8f6f4 MFMA ops with a 6-bit or 4-bit operand.** Covers every format pair, at both shapes. 15. **Add the larger-K variants of existing MMA shapes.** 16. **Accept logical and packed operands on the gfx1250 scaled WMMA ops.** A and B can be logical fp8 vectors or packed 6/4-bit words, as well as the padded `<16 x i32>` form. The gfx950 scaled MFMA ops keep the packed-word form: their only caller already loads those words. 17. **Run the fp64 GEMM numerics on every target that admits them.** They built for gfx942 only; they now build for the GPU they run on and run on gfx90a, gfx942 and gfx950 (the hardware-verified targets, marked `arch_specific`). An admitted target that has not run them yet skips, or fails when the lane requires it. 18. **Move the fp64 GEMM numerics to their own file and share the GPU checks.** `tests/numeric/test_f64_gemm_numeric.py` holds them; the gfx942 f16/bf16 swizzle regression keeps its own gfx942 file. A CPU-only test checks that both engines refuse fp64 `direct_to_lds` on targets without DRAM-to-LDS DMA. Shared test helpers replace copies in the numeric tests: `current_device_arch_in` (the GPU's arch if it is in a list, else skip, and fail under `ROCKE_REQUIRE_<ARCH>=1`), `available_engines` / `all_engines_backend`, and `universal_gemm` (lower and launch a universal GEMM spec with a required target arch). **Changes after review.** These commits follow the 18 above and change behavior: - **Compile-target routing follows the compiler that runs.** `compile_kernel` and `lower_kernel_for_comgr` route by the LLVM major of the COMGR compiler they load, and the hipcc helpers by hipcc's clang major, through `compiler_target_for_llvm_major` in `rocke.core.arch`. `ROCKE_LLVM_FLAVOR` still picks the emitted IR's conventions only. A major in a qualified lane uses that lane's measured processor answers; a newer or unknown major keeps the processor's own name, and a COMGR build refuses a compiler it cannot place. The hipcc helpers now run `hipcc --version` and raise if it prints no `clang version N.` line. - **Dispatch candidates list each target with its own facts once.** `arches_with_own_facts()` in `rocke.core.arch` gives the known targets whose rows hold their own facts; norm and the unified attention paths declare it, and the fp16/bf16 RCR WMMA GEMM candidates share `RDNA_WMMA_GEMM_ARCHES`. A processor is admitted through the generic its row points at (`listed_arch_for`), so example-parity builds each distinct target once. - **Generic members resolve through their generic's backend row, in both engines.** The backend registries hold rows only for targets with facts of their own; a member gets its generic's backend under its own name. Both engines refuse an unknown target with the same message, built from the registry. - **The C++ engine's compiler processor rejections are generated** from the arch-domain columns by `tools/gen_arch_domain_cpp.py`, through the Python reader, instead of a hand-kept list. - **New refusals, in both engines where both lower:** an f8f6f4 operand that is short of its packed words (for example one fp4 word where eight are needed) instead of zero-padding it; an MMA result declared in another number format than the instruction returns instead of bitcasting it; an arch-domain processor answer that is not a JSON boolean. - **Universal GEMM test launches** pack the arguments the builder declares (strides for every batched kernel) and default C to NaN in the kernel's storage dtype. - **Tests and docs.** The test-helper and gfx908 occupancy tests moved under `platform/tests/support` and `platform/tests/benchmark`, where `rocke_pytest` collects them; the shared llc helpers live in `rocke.tests.llc`. The unknown-target verdict test skips on a host that cannot judge any module, such as the installed test tree. The backend-registry and arch-data ctests are in the provider's quick tiers, and a test fails when a C++ test is in no tier. The fp64 GEMM docs name gfx90a, gfx942 and gfx950 everywhere, with a test that checks every such sentence, and the dispatch docs say that a listed generic target covers its members. Every lane map comes from an arbiter: hardware, a public calculator, public ISA tables, or CK where it is the named reference. None is invented. The docs list every new op and target: `dtype_support.md`, `mfma_atom_catalog.md` and `mma_metadata.md`. **Built on develop with #13976 and #13723:** - #13976 removed the two gfx950 ops that do not exist on the hardware, which was this PR's first commit, so that commit is gone. - The gfx950 f8f6f4 16x16x128 rows call the declaration #13976 made shared (`mfma.scale.f32.16x16x128.f8f6f4.v8i32.v8i32`), so a kernel that mixes them with #13976's builder ops emits one `declare`. - Develop's binary ABI is 6 (#14706, #14616); this PR is 7. - Develop's LLVM 24 lane (#14692): the named-barrier declarations are LLVM 24's own overrides in the chain, and the LLVM 24 arch-domain column and goldens cover this PR's targets and ops. - The LLVM 24 arch-domain column is regenerated with the llvm24 compiler CI runs (AMD clang 24.0.0git, llvm-project 2792f8ceeb5a), and its header now records that compiler. It answers 38 of develop's cells differently from the build develop's column recorded: on the five targets without gfx12 barriers the s.barrier.init/join/signal/signal.var/wait and s.cluster.barrier declarations crash the compiler (`toolchain_crash`) instead of being refused, the s.barrier.leave refusal has new wording, and gfx1250 cannot select the bf16 fmuladd, maxnum and minnum. That compiler accepts gfx1154 and gfx1250 by name, so at llvm24 both compile for themselves. - #14156 (gfx1250 LDS 320 KiB) is independent of this PR. ## Test Plan - One test file per op group (Python, both engines in strict mode) and one ctest per group (C++ builders and catalog lookups): - the IR for every new op at llvm20, llvm22, llvm23 and llvm24, and both engines equal; - the refusals, with their exact texts; - selection over every target and dtype; - lane-map checks against each arbiter; - compile and link per target. - Numeric tests on hardware where a GPU is available: the gfx90a MFMA f64, f32, i8 and bf16 ops, and the fp64 GEMM numerics on gfx90a, gfx942 and gfx950. - The representative-IR golden gains the 10 new gemm cases (`gemm.json` shard) at all four LLVM flavors, in commit 5. The other entries are unchanged. ## Test Result - Byte identity green at llvm20, llvm22, llvm23 and llvm24. - ctest 46 / 46 at the head (the arch-domain reproducibility test skips without its compiler). - Full platform suite at the head (Python engine): 6015 passed, 0 failed. - The portable-IR tests against a freshly built shared engine pass: 196. - The touched tests pass with both engines in strict mode, apart from `test_the_cpp_fallback_does_not_warn_twice`, which fails in strict mode on develop too. - The other platform test files and the library smoke: no new failure against develop. - pre-commit passes. - Per-commit check: every commit builds, passes ctest and passes its own tests (strict mode; the same develop test as above is the only exception). - The arch-domain columns equal a full regeneration at the head on llvm20, llvm22, llvm23 and llvm24, cell for cell; the llvm24 column at every commit is that commit's full regeneration with CI's llvm24 compiler. - At llvm24 with CI's llvm24 compiler: the compile-module and target identity tests pass, and so do the gfx90a numeric MFMA tests on a gfx90a. - gfx90a numeric tests pass. - fp64 GEMM numerics, both engines in strict mode, against an fp64 reference (exact on integer inputs, within the runbook tolerance on random inputs): every admitted config passes on gfx90a, gfx942 and gfx950, with partial tiles and a problem smaller than one tile; the Python and C++ IR are identical. gfx90a refuses the four `direct_to_lds` configs (no DRAM-to-LDS DMA), as the test expects. - The other targets are compile-verified. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: John Shumway <jshumway@amd.com>

…ng (#14836) ## Motivation Five tests fail on develop whenever the `rocke_engine` binding is built. Since #14616 the gfx942 split-KV attention binding (`gfx942_attention_tiled_3d_lower_llvm`) builds the reduce kernel after the segment kernel, the same as the gfx950 binding and both engines' parity emitters. The tests still expected the segment alone from the gfx942 binding: - `platform/tests/test_binding_compiler_selection.py`: its probe subprocess stopped with an `AssertionError` on the module count. - `library/tests/test_strided_kv_decode.py`, `test_attention_tiled_3d_kv_lds_swizzle.py`, `test_attention_tiled_3d_wave_page_id_load.py` and `test_attention_tiled_3d_running_max.py`: 22 gfx942 native-builder cases failed, the binding's output carrying a second module (the reduce) after the expected segment. The binding is correct; the test expectations are out of date. Fixes #14807 ## Technical Details Test-only changes. `platform/tests/test_binding_compiler_selection.py`: - Every tiled 3D binding (gfx942 and gfx950, paged and strided KV) is now expected to return two modules: the segment, then the reduce. - The reduce-stage resolver-failure check ran only on gfx950. It now runs on gfx942 as well, since gfx942 also has a second stage: a compiler-lookup failure in the reduce stage must propagate, and the resolver must have been called exactly twice. `library/tests`: - New helper `_tiled_3d_binding_ir.py` (`lower_like_tiled_3d_binding`) gives the binding's expected output one home: the Python-lowered segment kernel, then the reduce derived from the segment spec the way the binding derives it (head size, head counts, dtype, segment count and `waves_per_eu`). - Every tiled 3D native-builder test, gfx942 and gfx950, now compares the binding against that helper instead of its own copy of the reduce spec. This fixes the four stale gfx942 expectations and removes the per-test copies in `test_attention_tiled_3d_kv_cache_policy.py`. No engine code, emitted IR, or golden changes. ## Test Plan - Reproduce on develop: build `librocke_core.a` and the `rocke_engine` binding, run the five tests. - Run the updated tests with the binding built, in strict both-engine mode (`ROCKE_BACKEND=both ROCKE_CPP_STRICT=1`). - Build, ctest, byte identity at llvm20, llvm22, llvm23 and llvm24. - Full platform pytest suite with `rocke_engine` importable, failures compared by name against develop. ## Test Result - Before the fix: `test_python_bindings_use_the_compiling_library` fails on the module-count assertion; the four library tests fail 22 gfx942 cases (6 strided KV decode, 6 KV LDS swizzle, 6 wave page-id load, 4 running max). - After the fix: `test_binding_compiler_selection.py` 3 passed (Valgrind present, leak case ran). In the five library files, all 47 native-builder cases pass (22 failed before), in strict both-engine mode as well. The only other failures in those files are the 21 `test_strided_kv_decode.py` cases tracked in #14652 (`AttributeError: 'object' object has no attribute 'device'`), which this change does not touch. - Build OK; ctest 35/35 passed (the arch-domain reproducibility test was skipped). - Byte identity GREEN at llvm20, llvm22, llvm23 and llvm24. - Platform suite (`pytest tests` from `platform/`, `rocke_engine` importable): 6092 passed, 900 skipped, 0 failed; the only failure-name difference against develop is `test_binding_compiler_selection.py::test_python_bindings_use_the_compiling_library` (fails on develop, passes here). ISSUE ID : AICK-2899 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com>

Conflicts came from develop's attention work landing in the same code as the runtime head-count, window and batch changes (#14530, #14727): - #14616 2D LSE: lse_ptr is declared after the runtime kernargs in both engines, and _tiled_2d_signature takes emit_lse. The 2D and 3D LSE stores index the buffer by the runtime head count rather than a baked NUM_QH, in Python and the C++ twins. - #14367: VALID_ROWS guards combine with the runtime head bound, and the head_size 192 K/V path uses the runtime token stride. - Emit configs: develop keeps its indices; ours move to gfx950_attention_tiled_2d 68-70 and fastkv_regp 7-9. - Goldens: develop's representative shards with the attention, attention_d128_swa and attention_d256 shards re-blessed; the differential golden re-blessed for the gfx942/gfx950 tiled and strided-KV families. - The LSE dispatch test keeps develop's pinned configs under their attention ABI v2 tuning ids; every spec hash is develop's. Co-Authored-By: Claude <noreply@anthropic.com>

…r engines (#14242) ## Summary Packages rocKE's gfx950 unified attention as two hipDNN kernel-ingestor engines: `hipkernel:Gfx950AttentionUnified2D` (prefill and the decode shapes 3D does not cover) and `hipkernel:Gfx950AttentionUnified3D` (split-KV decode, a segment and a reduce launch per kernel through #14158's multi-launch `sources`). Both ship AOT, so an SDPA-forward graph on gfx950 can be served by unified attention with no JIT. Both catalogs are reduced to what a correctness-checked MI355X sweep measured as needed: Unified2D ships 222 of its 792-kernel full space and Unified3D 123 of 348, with no graph losing coverage or more than the measured noise (the final reduced-vs-full sweep is below). JIRA ID : AICK-2542 JIRA ID : AICK-2682 ## Risk Assessment Risk 2 of 5. The PR adds two engines and their descriptors. No existing engine's code, catalog or ranking changes: an unpinned SdpaFwd graph still ranks `hipkernel:Gfx950AttentionDense` first (device test `DenseStillRanksAheadWhenUnpinned`). The shared pieces it touches: - generator tooling: `registry_enumerator` full-space and multi-launch sections, the regeneration gates, provenance stamping, and two fixes in `variant_reduce`. Only `provenance.generated_from` is opt-in. The expanded registry output now always writes per-kernel `provenance.rocke_spec_id` / `rocke_tuning_id`, including Dense's 840 catalog kernels (its manifest pins a `spec_id`). The compact `variants:` selections drop it, so no shipped descriptor changes, but any other consumer of the expanded output sees the new field; - one rocKE builder, `build_unified_attention_3d_tiled_module`; - the ingestor save census's list of dispatch symbols. These are covered by their own tests and by the existing dense regeneration gates, which still pass. 3D is the first production user of `launchNames()` (#14158). ## ASIC Coverage gfx950 only: both engines ship only for gfx950, and other architectures' shards carry neither engine. PR CI has no gfx950 device lane (`hkp-probe-gfx950` packs and asserts but launches nothing), so the device evidence comes from MI355X (and, before the rebase, MI350X) runs. ## Testing Summary - Host: unit and census tests for both packs, the arch-neutral geometry, and the generator suite under `hipdnn_strict_skips`, including regeneration gates g1–g7 for all three gfx950 manifests; static descriptor gates on the installed shards (variant sets, coverage rungs 1–2, validator, launch surface). - Device, MI355X (gfx950), on the shipped (reduced) catalogs at this head's code: - engine-pinned SdpaFwd quick and standard suites against the CPU and GPU references, both engines plus a Dense regression run; - the knobs suite: every shipped kernel class forced and checked against the CPU reference, plus declines, precedence and a NaN-prefilled workspace; - a timing comparison of every engine that serves each corpus graph, and PyTorch; - RUNBOOK stage 7's final sweep: the full catalog against the reduced one in one session. - How the catalogs were reduced: a 3-round dnn-benchmark sweep with correctness on, over 188 prefill graphs (2D) and 88 decode graphs (3D), then `variant_reduce`. ## Packaging and performance statistics All numbers below come from tool output (packaging logs, device suite runs, dnn-benchmark sweeps). **Packaging** (gfx950 shard, shipped catalogs; pack time is the build's own `hkp_pack` command per engine, median of 3, relative to Dense's on the same job) | Engine | Kernels (full space) | AOT budget | Code objects, zstd | KDP | Pack time vs Dense, build default workers | Serial time per kernel vs Dense | |---|---:|---:|---:|---:|---:|---:| | Unified2D | 222 (792) | 1000 | 1.55 MiB | 0.29 MB | 0.11× | 1.58× | | Unified3D | 123 (348), 246 launches | 400 | 0.64 MiB | 0.26 MB | 0.41× | 1.06× | | Dense (catalog unchanged, for scale) | 840 (52 distinct binaries since #14059) | 1000 | 6.16 MiB | 1.26 MB | 1 | 1 | - The gfx950 kpack is 8.49 MiB at this head; the two engines are 2.19 MiB, 26% of it. Since #14059 (runtime head counts for Dense), Dense's 840 catalog kernels compile to 52 distinct binaries, but the kpack still stores one frame per kernel, so Dense grew from 5.79 to 6.16 MiB. Collapsing those entries is #14059's stated follow-up. The pack-time ratios and the full-space figure (14.45 MiB) were measured against Dense before #14059. - Provider load (`hipdnn_list_engines`, median of 5) is 1.24× its time without these descriptors. **Correctness and coverage** (MI355X, shipped catalogs) | Suite | Unified2D | Unified3D | Dense | |---|---|---|---| | SdpaFwd quick + standard, CPU and GPU reference (served / declined / failed) | 43 / 6 / 0 | 4 / 45 / 0 | 45 / 4 / 0 | | Knobs suite (forced kernel classes, declines, precedence) | 30 / 30 (both engines) | | 62 / 62 | | Corpus graphs served and validated against PyTorch | 188 / 188 | 82 / 88 | | Every decline is intentional: Unified2D leaves decode to Unified3D and declines Hkv > Hq; Unified3D serves decode only, and its 6 corpus declines are batch-256 decode graphs whose K/V is over the 32-bit addressing cap (ASM SDPA serves those). **Kernel time** (MI355X, autotuned, median of 3 rounds; ratio = Unified / other, below 1 = Unified faster; every engine's round-to-round drift ≤ 0.4%). The Unified2D rows were re-measured at this head, on the kernels develop's tiled-2D changes recompiled (see the checklist). | | Graphs both serve | Geomean ratio | Time-weighted (other / Unified) | Unified faster on | |---|---:|---:|---:|---:| | Unified3D vs ASM SDPA (decode) | 60 | 0.200 | 8.17 | 60 / 60 | | Unified3D vs PyTorch | 82 | 0.062 | 34.2 | 82 / 82 | | Unified2D vs Dense (prefill) | 179 | 2.37 | 0.38 | 0 / 179 | | Unified2D vs PyTorch | 188 | 1.38 | 0.55 | 39 / 188 | - Unified3D serves 22 decode graphs no existing hipDNN engine serves. - Unified2D is slower than Dense wherever Dense serves, and Dense keeps ranking first there; it adds coverage Dense lacks (the 9 bottom-right causal graphs in these corpora). **Variant reduction** | Engine | Kernels | Chosen eps | Measured drift | Final sweep, reduced / full kernel time (geomean, rounds 1–3) | |---|---|---:|---:|---| | Unified2D | 792 → 222 | 0.02 | 0.62% | 1.000, 0.999, 1.000 (u2d_attn); 1.000, 0.999, 1.000 (u2d_rocke) | | Unified3D | 348 → 123 | 0.03 | 1.08% | 1.004, 1.002, 0.999 (u3d_attn); 1.002, 1.007, 1.004 (u3d_rocke) | Both arms of the final sweep served the same graphs and passed correctness on every corpus. ## Testing Checklist - [x] Unit (host, no GPU) - `HIP_VISIBLE_DEVICES=-1 hip_kernel_provider_tests` - 1454 passed, including `TestKernelIngestorEngine.EveryHandlerReportsSupportsSavingExactlyWhenItSavesAndRestores`. Our build host also fails 11 Dense signature/dispatch tests: it has a `/nonexistent` directory, which those tests assume is absent. They pass in CI - Status: Passed - [x] Census - `ctest -R 'hkp-census|rocke_unified'` - 36/36 (the census inventories are the reduced catalogs') - Status: Passed - [x] Generator suite (strict skips) - `pytest -p hipdnn_strict_skips IngestorGenerator/tests` - 1646 passed, 0 skipped - Status: Passed - [x] Regeneration gates on the reduced selections - `tests/test_registry_regeneration.py` (g1–g7) 39 passed; AOT budget tests 12 passed - Status: Passed - [x] rocKE attention CPU tests and packaging guards - 2D/3D mask_type, 3D module, decode split, KV cache policy, spec-predicate parity: 145 passed, plus 6 develop KV-cache tests that need the compiled `rocke_engine`; `test_hkp_pack_producer_guards.py` and `test_gfx950_attention_unified_3d_bundle_coverage.py`: 26 passed - Status: Passed - [x] Golden IR, before the reduction - `ctest -R rocke_golden_static` - the reduction changes no rocKE code - Status: Passed - [x] Launch surface - `launch_surface.py configs/gfx950_attention_unified_{2d,3d}.engine.yaml --check` - Status: Passed - [x] Variant sets - `verify_variant_sets.py --mode structural` and `--mode full --arch gfx950 --manifest …` - 2D 222/222 distinct binaries, 3D 123/123 - Status: Passed - [x] Coverage gate rungs 1–2 - `coverage_gate.py --mode full --expect-engine …` - Status: Passed - [x] Validator - `hipdnn_validate_descriptors <shard> --expect-engine … --json` - success, no warnings - Status: Passed - [x] Packaging probe - gfx942 configure with `HIPKERNELPROVIDER_ENABLE_PACKAGING_PROBES=ON`, `ctest -R '^hkp-probe-gfx950'` - 12/12; `hkp-probe-gfx950` packs 51 UKDs of the current descriptors in a gfx942-only build - Status: Passed - [x] Device suites, MI355X - engine-pinned SdpaFwd quick + standard with GPU reference (`hip_kernel_provider_gfx950_attention_*_gpu_ref_integration_tests`) and CPU reference (`*-external-integration_{quick,standard}_suite`) - 0 failures: Unified2D 43 served / 6 declined, Unified3D 4 / 45, Dense 45 / 4 - Status: Passed - [x] Device knobs suite, MI355X - `hip_kernel_provider_integration_tests --gtest_filter='*Gfx950AttentionUnified*'` - 30/30; Dense knobs 62/62 - Status: Passed - [x] Fault re-check, MI355X - the whole `u2d_rocke` corpus in one autotuned process, twice and then with serialized launches - 93/93 each, 0 errors - Status: Passed - [x] Reduction sweep + `variant_reduce.py` - 3 rounds, correctness on; 2D `SWEEP_DONE`, 3D `SWEEP_DONE`; 2D 792 → 222, 3D 348 → 123 - Status: Passed - [x] Final sweep, reduced against full (RUNBOOK §7) - same coverage, correctness passed on both arms, ratios above - Status: Passed - [x] Unified3D catalog regenerated for #14367's LLC-aware K/V cache policy - all 123 kernels now `kv_cache_policy: "all"` (was `"stream"`), same kernel classes; MI355X re-run: suites and knobs unchanged, 246/246 decode rows correct, kernel time against the `"stream"` catalog geomean 0.998, time-weighted 0.999 - Status: Passed - [x] Device evidence at this head (rebased onto develop 2026-11-03, after the review-round refactors) - Dense's and Unified3D's code objects are byte-identical to the previous device runs; develop's #13652 recompiled 201 of the 222 Unified2D code objects (same kernels and ids), and the Unified2D/3D matchers moved onto a shared header, so both engines' device runs were repeated at this head: suites Unified2D 43 / 6 / 0, Unified3D 4 / 45 / 0, Dense 45 / 4 / 0 (passed / skipped / failed), knobs 30/30 and 62/62; compare 188/188 Unified2D graphs (564/564 rows) and 82 Unified3D graphs (246/246 rows) correct against PyTorch in all 3 rounds. Kernel time vs the previous runs: Unified2D geomean 0.985 (none of 188 graphs >2% slower), Unified3D 1.003 (ASM_SDPA control 1.006). A final rebase picked up #14059, which recompiled Dense's 840 code objects only (Unified2D/3D byte-identical to the run above); the suites and knobs were re-run at this head with the same results, `DenseStillRanksAheadWhenUnpinned` included - Status: Passed - [x] DVC - both reduced KDPs on the `ingestor` remote; `dvc status -c` in sync - Status: Passed - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - **Unified2D pack** (`Gfx950AttentionUnified2DNative.cpp`): matches BSHD SDPA forward on the identity paged-KV layout and builds block_tables, seq_lens and query_start_len in the workspace. - It declines: batch > 511 (the baked 9-iteration binary search), Skv not a multiple of 16, i32 KV over 0x7FFF0000, any window, and the decode cells Unified3D covers (Sq = 1 with the GQA group dividing 16). - Engine knobs: tile_size, num_warps, block_m_per_warp, use_mfma_32x32 and use_transposed_qk_32x32. - **Unified3D pack** (`Gfx950AttentionUnified3DNative.cpp`): decode only. - `launchNames()` is `{segment, reduce}`, and both definitions are resolved through `launchKernelDefinition` with a signature check before any load. - The segment launch writes `[token, Hq, seg, D]` scratch and the reduce always runs, so the workspace needs no zeroing. - Knob: num_segments. No batch limit. - **Arch-neutral geometry** (`UnifiedAttentionGeometry.hpp`): the launch geometry and workspace layouts take a `UnifiedAttentionArchTraits`, and the gfx950 headers are thin wrappers, so a gfx942 pack (AICK-2544) only supplies traits. - **rocKE**: `build_unified_attention_3d_tiled_module` packs the segment and its derived reduce (`reduce_spec_for`) as one `KernelModule`. The module predicate gates both halves. The 3D pack compiles develop's running-max fix for split-KV (#14276). - **Generator**: - `enumerate.also` sections may use `space: full` and carry manifest `launches`. The enumerator refuses a registry that pairs a reduce when the manifest declares no launches, and keeps one default per cell. - The gates are launch-aware: g1 keys on launches, g3 uses `--subset` for full spaces, g4 compares launch sets, and g6 walks every launch and checks module roles. - g2 and g3 run per engine with the shipped ids reused (#14109), comparing every descriptor byte for byte. - `generate.py --stamp-provenance` stamps `provenance.generated_from` (config, manifest, generator and rocKE commits). It is opt-in, so a default regeneration stays byte-identical. - `variant_reduce` fixes: - it reads the format-2 winner lines the runtime has written since #13028; - it keys a graph on its content with `Graph.id` (minted fresh at each finalization) zeroed, so sweep rounds match. - **Catalogs** (the manifests enumerate the full spaces; the selections are `variant_reduce`'s): - Unified2D: 222 kernels over 198 cells (33 head configs × bf16/fp16 × 3 masks): the narrow nw4/mw16/t64 geometry in every cell, transposed 32x32 nw4/mw32/t64 in 21 and narrow nw2/mw16/t32 in 3. The wide 32x32 geometry won no graph and is not shipped. - Unified3D: 123 kernels: num_segments 16 in all 116 decode cells, 32 or 64 in 7 bf16 bottom-right D128 variants. The K/V cache policy follows #14367's LLC-aware rule, which picks `"all"` for every shipped decode cell. - **Descriptors**: both are generated with develop's generator (post-#13028 UED/UHD, UKD `sources` for 3D), every kernel id reused from the full catalog. A regeneration with the shipped ids reused reproduces them byte for byte. The KDPs live in DVC (`ingestor` remote). - **Shared code** (one home per fact): - `packs/UnifiedAttentionGraphMatch.hpp`: the graph-match helpers both unified packs use (operand/BSHD checks, SDPA node, mask classification, dtype, metadata readers); the tests' graph builder is `UnifiedAttentionTestGraphs.hpp`. - rocKE `kernels/common/attention_mask_type.py` (mask ordinals and band -> kind, used by both tiled specs and re-exported by dispatch) and `attention_split_kv_reduce.py` (the segment -> reduce derivation the dispatcher and the module export share). - CMake registers Dense, Unified2D and Unified3D by one loop per lane (external integration, gpu_ref, census). - **Tests**: - The device knobs suite (`IntegrationGpuGfx950AttentionUnifiedKnobs.cpp`) adds a harness `prepareWorkspace` hook. It forces the default geometry/segment count across the head configurations and each other shipped class on a cell that carries it, and `CoversEveryCatalogKernelClass` holds those tables to the catalog. - The engine-pinned `gpu_ref` targets and external-integration lanes cover both engines. - The ingestor save census names all seven pack dispatch symbols. - The Unified2D/3D integration TOMLs are reserved with the gfx950 census shard (`hkp_reserve_census_shard(... FILES ...)`, as #14573 requires of every shard entry). **Constraints and follow-ups** - #13372 and #13889 are both in develop; this branch no longer carries a copy of #13889. - Heads are catalog cells (baked), window = 0, 2D batch ≤ 511, i32 KV ≤ 0x7FFF0000. D256 is deferred for both engines. The kernarg work is tracked in AICK-2675. - There is no 2D/3D `select_path` crossover replay: the split is the static decode rule above. - Runtimes older than UKD 1.1 skip the 3D descriptors. Strided-KV 3D is a follow-up. - Follow-up: execution-plan save/restore (#13315) for both handlers. Today they report `supportsSaving() == false`, as the RDNA handler does. - One earlier MI355X sweep hit a single GPU fault in Unified2D on `u2d_rocke/rocke_sdpa_035` (bf16, D 128, Hq 64 / Hkv 8, causal, S 512). It has not reproduced since: not in the forced geometries on that cell, 15 single-graph runs, 4 full-corpus runs on MI350X, or the MI355X full-corpus re-check above. - Unified2D kernel time moved with develop's tiled-2D changes (#14542, #14616), not with this PR: at this head it is 1.035× (geomean) its pre-rebase time on the same 188 graphs, up to 1.36× on long causal bf16 D128 prefill (S 8192), while Dense in the same runs is flat (0.997). Dense still ranks first wherever it serves; the table above is at this head. - Measurement note: the pinned dnn-benchmarking (v1 result schema, which `sweep.py` reads) has no bottom-right causal PyTorch reference. The sweeps used upstream dnn-benchmarking #84's reference, with the 9 bottom-right corpus graphs written as `causal_mask: false, right_bound: 0, diagonal_alignment: BOTTOM_RIGHT` (the same band under hipDNN's `resolveDiagonalBand`): #84 reads `causal_mask: true` as top-left regardless of the alignment. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>

…lpers (#14244) ## Summary C++ mirror of the shared dense attention spec (`library/kernels/common/attention_dense_spec.py`) and of the dense-attention LSE helpers in `_lse_store.py`. First step of [AICK-2711](https://amd-hub.atlassian.net/browse/AICK-2711) (gfx950 dense attention C++ mirror); the gfx950 spec layer and both builders follow in separate PRs. Nothing in the shipped provider calls the new C code yet; the gfx950 and gfx942 builder ports will. ## What's ported | Python | C | |---|---| | `AttentionDenseSpec` fields and defaults | `rocke_attention_dense_spec_t`, `rocke_attention_dense_spec_default()` (`include/rocke/instance_attention_dense_spec.h`) | | `DENSE_TILE_GEOMETRIES` | `ROCKE_DENSE_TILE_GEOMETRIES` | | `__post_init__` | `rocke_attention_dense_spec_validate()` | | `resolved_persist_decode` / `resolved_nonpersist_decode` | `rocke_attention_dense_spec_resolved_*_decode()` | | `kernel_name()` and its `_*_name_parts` hooks | `rocke_attention_dense_kernel_name()`, `rocke_attention_dense_name_parts_t`, `rocke_attention_dense_name_parts_base()` | | `check_dense_spec_preflight` | `rocke_attention_dense_spec_preflight()` | | `PERSIST_DECODES` / `NONPERSIST_DECODES` (name, tag, check) | one `{name, tag, check}` table per grid in `instances/common/attention_dense_spec.cpp` | | `dense_row_has_key`, `zero_keyless_row` | `rocke_dense_row_has_key`, `rocke_zero_keyless_row` in `include/rocke/instance_lse_store_internal.h`, next to develop's `rocke_lse_store` (#14616), which the dense path reuses | Shared helpers this PR adds to or reuses: `rocke_spec_set_reason` (now marks a cut reason with `...`) and `rocke_strbuf_append_py_repr` (Python `repr(str)`, now also used by the IR printer). **Not ported here:** - `attention_dense_cache_key` (Python launcher-cache key; no C caller). - `validate_dense_lse` / `validate_unified_lse` (host-side torch tensor checks; no C caller). - `kernels/common/attention_dense_kv_split.py` (#14367: `check_kv_segments`, `segment_tile_range`, `emit_segment_store`, `declare_segment_params`, `MAX_KV_SEGMENTS`) and the `kv_segments` field / `kvs{N}` token. These belong to the gfx950 (AICK-2711) and gfx942 (AICK-2710) builder PRs, whose parity pairs must cover `kv_segments > 1`. Recorded on both tickets. - The runtime-shape builder surface from #14059 (`runtime_shape`, `fastdiv_divisors`, `runtime_kernarg_fields` on the spec; `make_fastdiv` / `fastdiv_kernargs` in the decode module). It feeds kernel params and the persistent fast division, so it goes with the builder ports. ## Design Subclass hooks are passed in explicitly: - per-class answers (`supports_*`, both decode sets, type name) through `rocke_attention_dense_traits_t`; - the per-spec `_aligned_causal` override (gfx950: `right_bound != 0`) through a `lookahead` argument; - name parts through `rocke_attention_dense_name_parts_t`. Preflight runs the base checks on the spec it is given, then an optional arch validator (`rocke_attention_dense_validate_fn`) for the arch's own fields, as both Python subclasses do with `super().__post_init__()`. A decode a traits list names but the C table doesn't implement is rejected, and `name_parts_base` returns an error for it rather than the default order's empty tag. `kernel_name` returns `ROCKE_ERR_VALUE` instead of a truncated name. The accessors, `name_parts_base` and `kernel_name` require a spec that passed preflight; on one that didn't, they return 0 / NULL / `ROCKE_ERR_VALUE` instead of dividing by zero. Every `ValueError` reason matches the Python text byte for byte. The two paged geometries where Python raises `ZeroDivisionError` (`block_m < 32`, or more waves than `block_n` rows) return a fixed `"division by zero"`; CPython's text varies by version, so those are compared by verdict only. ## How it's pinned **Spec layer: the `rocke_attention_dense_spec` ctest.** The reasons, block orders and names don't all reach emitted IR, so byte identity can't see them. Both engines are pinned to `platform/tests/golden/attention_dense_spec_vectors.json`. `library/tests/test_attention_dense_spec_vectors.py` recomputes it with the current Python (also in `rocke_library_attention_cpu_pytest`) and owns `--regenerate`; the ctest embeds the same file and requires the same answers from C. The file holds: - every field default and the tile table; - 129 hand-picked base cases and 400 cases drawn with a fixed seed (2711) around the auto block-order thresholds. `test_cases_reach_every_reject_site` traces the Python and requires every reject the C port mirrors to be reached by some case, and each of the two paged divisions to raise `ZeroDivisionError`; - a `Gfx950AttentionDenseSpec` hook table (16 cases) and a `Gfx942AttentionDenseSpec` hook table (9 cases). `test_hook_tables_exercise_every_hook` reverts each hook a class overrides and requires a recorded answer to change. Changing one token of `attention_dense_spec.cpp` at a time (thresholds, decode preconditions, ignoring a traits list or `lookahead`, name tags, defaults, the tile table, preflight's base and arch halves, accessor guards) fails the ctest in every case tried. **LSE helpers: byte identity.** The `attention_dense_lse` parity pair (`library/tests/parity/attention_dense_lse_emit.{py,c}`, 8 configs) is GREEN in `check_byte_identity.py`, as are all 97 families. ## Testing - All 36 host ctests pass (build tree); `rocke_attention_dense_spec` and `rocke_library_attention_cpu_pytest` pass from the installed tree. - `test_attention_dense_spec_vectors.py`, `test_attention_dense_lse.py`, `test_lse_store_ir.py` pass. - `check_byte_identity.py`: GREEN, all families. - pre-commit clean. ## Land order #14059 landed first; the vectors are regenerated on top of it (only gfx950 hook-case name parts moved). Any later change to the shared spec or a subclass's hooks needs `python test_attention_dense_spec_vectors.py --regenerate`. ## Known gaps (pre-existing in Python, tracked separately) - Preflight checks Q/O in elements, but gfx950's ragged path addresses Q in bytes: [AICK-2876](https://amd-hub.atlassian.net/browse/AICK-2876). This PR mirrors the current Python check; the fix goes to Python first. - gfx950's `right_bound` has no upper bound, so `q + right_bound` can overflow i32: [AICK-2915](https://amd-hub.atlassian.net/browse/AICK-2915). gfx950's validator isn't in this PR; the fix lands with the gfx950 spec layer PR, Python first. [AICK-2711]: https://amd-hub.atlassian.net/browse/AICK-2711?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ [AICK-2876]: https://amd-hub.atlassian.net/browse/AICK-2876?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: Chris Millette <chris.millette@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

AviralGoelAMD added 30 commits October 27, 2026 11:54

test(rocke): host contract tests for tiled 2D LSE (red; AICK-2685 Tas…

d5d09f5

…ks 8-11)

test(rocke): tiled 2D rows with a fully masked first KV tile

c7dbb64

fix(rocke): carry the unclamped running max in tiled 2D attention

120a0da

fix(rocke): mirror the tiled 2D running-max carry in the C++ twins; r…

e38908f

…e-bless goldens

fix(rocke): drop dead running-max copies and stale comments in tiled …

5b2a01d

…2D C++; register the masked-tile tests

test(rocke): scalar unified attention rows with fully masked leading …

d584aed

…keys

fix(rocke): carry the raw running max in the scalar unified attention…

a4d3c22

… kernels

fix(rocke): re-bless scalar attention goldens

5f20fa2

fix(rocke): describe the scalar reduce guards in the C++ phase comments

0feeefa

test(rocke): state the masked-tile underflow points consistently

3cd8c2a

test(rocke): pin NaN propagation through each tiled 2D epilogue branch

cdf5221

Add a NaN-query-row case per tiled 2D epilogue branch, state what each mask case of the scalar 2D NaN test pins, and drop the unused segm_max argument of the scalar 3D workspace hook.

Merge users/avirgoel/rocke/tiled-2d-running-max-carry (AICK-2829, #14542

0fd5414

) into unified-2d-lse Plan Task 0.6: the LSE work builds on the raw-max carry and the exact-zero keyless guard. Brings develop up to a3b5d0c as well.

AviralGoelAMD disabled auto-merge October 30, 2026 22:01

Merge branch 'develop' into users/avirgoel/rocke/unified-2d-lse

aceb150

AviralGoelAMD enabled auto-merge (squash) October 31, 2026 01:22

Merge origin/develop into users/avirgoel/rocke/unified-2d-lse

eb0d43e

Develop now carries the base PR as a squash commit whose rocKE tree equals the base-PR head this branch already merged. Keep this branch's side of every conflict: the merged tree is identical to the previous head.

AviralGoelAMD disabled auto-merge October 31, 2026 01:24

AviralGoelAMD enabled auto-merge (squash) October 31, 2026 01:24

AviralGoelAMD merged commit 7cd54ff into develop Oct 31, 2026
93 of 95 checks passed

AviralGoelAMD deleted the users/avirgoel/rocke/unified-2d-lse branch October 31, 2026 03:41

shumway mentioned this pull request Oct 31, 2026

fix(rocke): refresh LLVM 24 attention goldens #14816

Merged

2 of 3 tasks

ThruptiRajLakshmanaGowda mentioned this pull request Nov 1, 2026

feat(hip-kernel-provider): add gfx950 unified attention 2D/3D ingestor engines #14242

Merged

19 of 20 tasks

cgmillette mentioned this pull request Nov 1, 2026

test(rocke): expect segment and reduce from the gfx942 tiled 3D binding #14836

Merged

aosewski mentioned this pull request Nov 2, 2026

feat(hip-kernel-provider): mirror additive attention bias in the C++ engine and reach it from AttentionRequest #14440

Open

10 tasks done

JohnUzoka mentioned this pull request Nov 3, 2026

feat(rocke): C++ mirror of the shared dense attention spec and LSE helpers #14244

Merged

dpeabody-amd mentioned this pull request Nov 3, 2026

feat(hip-kernel-provider): paged/ragged SDPA engine for gfx942/gfx950 #14247

Open

Navigation Menu

Repository navigation

feat(rocke): optional LSE output on unified tiled 2D attention - #14616

feat(rocke): optional LSE output on unified tiled 2D attention#14616
AviralGoelAMD merged 76 commits into
developfrom
users/avirgoel/rocke/unified-2d-lse

AviralGoelAMD commented Oct 28, 2026 •

edited

Loading

Uh oh!

cgmillette commented Oct 31, 2026

Uh oh!

cgmillette commented Oct 31, 2026

Uh oh!

Uh oh!

assistant-librarian Bot commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

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

Uh oh!

What was missing

What changed

Why it works

How it was validated

Notes

Uh oh!

cgmillette commented Oct 31, 2026

PR #14616 — optional LSE output on unified tiled 2D attention (AviralGoelAMD, AICK-2685) — 🟡 Follow-up: clean; land order pending

Uh oh!

cgmillette commented Oct 31, 2026

PR #14616 — optional LSE output on unified tiled 2D attention (AviralGoelAMD, AICK-2685) — 🟡 Follow-up: develop merge clean; CI running

Uh oh!

Uh oh!

assistant-librarian Bot commented Nov 3, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

AviralGoelAMD commented Oct 28, 2026 •

edited

Loading