Repository navigation
perf(rocke/attention): recover fp16 gfx942 D128 prefill via depth-2 sliced-K ring - #9662
Merged
Merged
Conversation
AnaghaRaoAMD
marked this pull request as ready for review
July 21, 2026 21:27
AnaghaRaoAMD
force-pushed
the
users/anarao/s34-ring-perf-perf
branch
from
July 21, 2026 21:32
1cffdfe to
2ef1678
Compare
Base automatically changed from
users/anarao/s34-fix-bf16-d128-ring-correctness
to
develop
July 22, 2026 21:32
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
…liced-K ring
Stacked on the D128 ring correctness fix. That fix put fp16 D128 on the non-ring
flash path (correct but ~0.55-0.99x AOTriton flash). This restores the ring for
fp16 D128 using a depth-2 pipeline that is both correct and faster.
Root cause of the original bug was the depth-3 ring's kg%3 slot map: at k_groups=4
(D128) slice 3 reused slot 0 and the reusing DMA was unfenced, corrupting the QK
accumulation. Two fixes here:
1. Drain-on-reuse fence (now unconditional in the ring schedule): before a
prefetch DMA that reuses the LDS slot slice(kg-1) was just read from, drain
the LDS reads + barrier. Previously this drain only fired under LDSSEQ.
2. ring_depth spec field. Depth-2 (k%2, 2 slots) keeps the live set {kg, kg+1}
with no slot reuse in the k_groups=4 case, and uses less LDS -> lower
occupancy pressure. It is the measured best fp16 D128 prefill geometry.
Routing (attention_spec_builder / _select_gfx942_flash_ring_depth):
* fp16 D128 -> depth-2 ring. Correct at magnitude (max_abs ~0.002 vs fp32
oracle) and 1.18-1.56x faster than the non-ring path (0.65-1.55x flash;
beats flash at S4096). Verified on MI300X, GQA 32/8,128/8 + MHA 16/16,
S1024-4096.
* bf16 D128 -> stays non-ring. At the production block_size=64 the correct
rings are all slower than non-ring T=64 for bf16 (the faster T=32 ring needs
block_size=32, which production never sends for D128 -- latent).
* D64 (both dtypes) -> depth-3 ring, unchanged.
Python-only, consistent with how the sliced-K ring already ships: the ring
schedule lives only in the Python engine (the C++ engine has ring primitives but
no schedule), and attention is not in the byte-identity gate. #9198 enabled the
D128 ring with no C++ change for the same reason.
Guardrail test updated for per-dtype routing (bf16 non-ring, fp16 depth-2 ring)
and the on-GPU numeric check now gates both paths. 177 attention tests pass.
Co-Authored-By: Claude <noreply@anthropic.com>
AnaghaRaoAMD
force-pushed
the
users/anarao/s34-ring-perf-perf
branch
from
July 22, 2026 21:41
2ef1678 to
d0cdf1c
Compare
AviralGoelAMD
approved these changes
Jul 22, 2026
AviralGoelAMD
left a comment
AviralGoelAMD
left a comment
Contributor
There was a problem hiding this comment.
LGTM with minor suggestion.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR restores the gfx942 sliced-K ring path for fp16 D128 prefill by introducing a depth-2 ring schedule and routing logic, while keeping bf16 D128 on the non-ring flash path. It also updates the D128 routing guardrail tests to validate the per-dtype policy and numerics on-GPU.
Changes:
- Add
ring_depthto the gfx942 2D tiled attention spec and tag it into the kernel name for cache separation. - Implement a depth-2 sliced-K ring schedule (and unconditional drain-on-reuse fence) in the gfx942 tiled kernel builder.
- Update gfx942 ring enablement/selection to route fp16 D128 to depth-2 ring and keep bf16 D128 non-ring; update regression tests accordingly.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py | Adds ring_depth, tags kernel name, and adjusts sliced-K ring scheduling to support depth-2 plus unconditional drain-on-reuse fencing. |
| dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py | Updates gfx942 ring routing: bf16 D128 stays non-ring; fp16 D128 can use ring with depth selected via new helper. |
| dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py | Wires ring_depth from selector into the spec construction paths. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py | Updates regression coverage to assert per-dtype routing and gate bf16 non-ring + fp16 depth-2 ring numerics on-GPU. |
…lback
Addresses review: ring_depth was silently sanitized to 3 when not in (2, 3),
masking misconfigurations (e.g. a typo ring_depth=23) and letting the compiled
kernel diverge from the spec. Validate ring_depth in {2, 3} in __post_init__ when
use_k_sliced_ring is set (raise early), and use spec.ring_depth directly in the
builder. Non-ring specs are unaffected (ring_depth unused when the ring is off).
Co-Authored-By: Claude <noreply@anthropic.com>
|
🎉 All checks passed! This PR is ready for review. |
Addresses review: depth-2 (fp16 D128) and depth-3 rings are distinct schedules (different slot map + LDS footprint), so _tiled_cache_key must distinguish them or two shapes sharing the geometry but differing on depth would collide on the same cached launcher. Add _select_gfx942_flash_ring_depth to the key (only when the ring is active; None otherwise). Matches the rd2 kernel_name tag. Verified the key differs for fp16 D128 (depth-2) vs D64 (depth-3). Co-Authored-By: Claude <noreply@anthropic.com>
github-actions
Bot
added
the
rocke-lib: attention
Addresses review (Copilot): the CK LdsSeq maps ((1,2,0,1) / (1,2)) are 3-slot
layouts -- they reference slot 2. With ring_depth=2 (fp16 D128) K_lds is allocated
with only 2 slots, so returning slot 2 indexes past the allocation -> out-of-bounds
LDS -> corruption. Reachable via HIPDNN_GFX942_K_LDSSEQ=1 on fp16 D128.
Three layers:
* _kslot: gate the LdsSeq maps on RING_DEPTH == 3; depth-2 uses the modulo map
(respects K_SLICE_SLOTS).
* __post_init__: raise if use_k_sliced_ldsseq is set with ring_depth != 3.
* selector: _enable_gfx942_flash_k_sliced_ldsseq returns False at depth-2 even
when the env asks, so the env can't build a raising spec.
Verified: ldsseq+depth2 raises; fp16 D128 + K_LDSSEQ=1 now falls back to ldsseq=off
and passes on-GPU (max_abs 0.002); D64 depth-3 ldsseq still enables. 63 tests pass.
Co-Authored-By: Claude <noreply@anthropic.com>
1 task
AviralGoelAMD
added a commit
that referenced
this pull request
Jul 29, 2026
…st) across gfx950/gfx1250 (#10126) JIRA ID: AICK-1492 ## Summary Restores **`supports_tiled_2d` signature parity across arches** and renames the flag per review. The shared feasibility check `supports_native_unified_attention_tiled` (in `kernels/common/attention_unified.py`) builds one kwarg list and passes the D256-fast flag to whichever arch's `supports_tiled_2d` applies. That flag was declared only on **gfx942** (added with the D256 4-warp fast path), so **gfx950 and gfx1250** — keyword-only signatures with no such param and no `**kwargs` — raised: ``` TypeError: supports_tiled_2d() got an unexpected keyword argument ... ``` for any problem routed through the check. It stayed hidden because production (`run_unified_attention_torch`) skips this check; only the bench/parity `path="auto"` path hits it. ## Changes - **Rename** the kwarg `use_d256_gfx942_fast` → **`use_d256_fast`** everywhere it is declared/passed (per review): the shared caller and the gfx942/gfx950 signatures. The predicate `_d256_gfx942_fast(problem)` that supplies the value is unchanged. - **gfx950**: add `use_d256_fast: bool = False` (accept-and-ignore) — the original parity fix. - **gfx1250**: add `use_d256_fast: bool = False` (accept-and-ignore, plus the `del` discard) — it had the **same latent gap** as gfx950, so this completes the parity rather than leaving gfx1250 to hit the identical error on its next bench/parity run. ## Scope / impact - **Production is unaffected**: `run_unified_attention_torch` does not call this feasibility check; it is bench/parity-only. No kernel/emitter change; no C++ references (rename is Python-only, no byte-identity impact). - Surfaced while GPU-verifying gfx950 D128 sliding-window prefill via the parity harness (the harness could not run any gfx950 shape until this was fixed). - Note: an unrelated, pre-existing test `test_dense_pipe_wiring.py::test_bf16_ring_enabled_for_d128` fails on develop (it still asserts the bf16-D128 ring that #9647/#9662 moved to non-ring); it is not touched by this PR. Changes are behind the feature flag HIPKERNELPROVIDER_ENABLE_ROCKE which is defaulted to OFF.
6 tasks done
dpeabody-amd
added a commit
that referenced
this pull request
Aug 6, 2026
…ion_dense (#10102) ## Motivation `attention_dense` selects its LDS pitch from the DMA row-packing factor: ```python ROWS_PER_INSTR = 128 // D # 1 at D=128, 2 at D=64 LDROW = D + PAD if ROWS_PER_INSTR == 1 else D VROW = D + _LDS_PAD_V if ROWS_PER_INSTR == 1 else D ``` At D=64 both pads collapse to zero, so neither LDS bank-conflict lever is active on that head size. The module docstring already recorded this as known — "its reads take some LDS bank conflict (V-pad is a future D=64 perf lever)" — because D=64 was added for coverage under the constraint that it cost the D=128 schedule nothing. Why it matters, in bank terms. On gfx950 bank = `(byte_addr / 4) % 64` over 64 banks. A wave64 `ds_read_b128` moves 1024 B while LDS delivers 64 banks x 4 B = 256 B/cycle, so **4 lanes per bank is the conflict-free floor**, not merely "better". Aggregating all 64 lanes of one QK read (`n=8` bf16 = 4 dwords each, `krow = nsub*32 + lane%32`, `col = ks*16 + (lane//32)*8`): | K LDS layout | banks touched | lanes per bank | |---|---|---| | D=128 unpadded `[BN, 128]` | 8 / 64 | 32 | | D=128 shipped `[BN, 136]` | 64 / 64 | 4 (floor) | | **D=64 shipped `[BN, 64]`** | **16 / 64** | **16** | | D=64 group-padded `[BN/2, 128+8]` | 64 / 64 | 4 (floor) | So D=64 today sits far closer to the unpadded D=128 case than to the shipped one. ## Technical Details A per-row pad is not possible at D=64: one `async_buffer_load_lds` moves 64 lanes x 2 bf16 = 128 contiguous elements, which is exactly one D=128 row but *two* D=64 rows, so a padded row pitch would place the second row at the wrong offset. The pad therefore goes **between DMA row-groups** — the group stays contiguous for the DMA, while the group pitch restores the bank spread. - New `lds_k_group_pad` spec field (default 8; 0 reproduces the old unpadded layout for A/B). `__post_init__` rejects a non-multiple-of-8 pad, because `smem_load_vN` stamps `align 16` on the `n=8` read unconditionally, so an 8-byte-aligned pitch would keep the `ds_read_b128` while breaking its alignment contract silently. It also joins `kernel_name()` on the packed path — it changes the emitted layout, so it has to be part of the kernel identity or two layouts share a symbol name and a launcher-cache entry. The `head_size=128` name is unchanged. - K LDS becomes `[NBUF, BN // ROWS_PER_INSTR, ROWS_PER_INSTR * D + spec.lds_k_group_pad]` when `ROWS_PER_INSTR > 1`; unchanged `[NBUF, BN, D + PAD]` otherwise. - `_async_load` strides its destination by *group* rather than by row. The sixth parameter is renamed `ldrow_bytes` -> `group_bytes`; the two are equal when `ROWS_PER_INSTR == 1`. V passes the plain unpadded `ROWS_PER_INSTR * D` pitch, so its layout is bit-identically addressed (`V_GROUP_BYTES * grp` reduces to the previous `VROW * 2 * row0`). - `do_qk` maps `krow` to `(krow // K_GROUP, (krow % K_GROUP) * D + col)`. This is exact because `nsub*32` is a multiple of `K_GROUP`, so the group/sub-row split commutes with the `nsub` offset. Both terms are hoisted above the KV loop. - Mirrored in the default and the persistent builder, with an assert that the group split divides evenly so a future head size fails loudly. - **Why 8 and not 16.** The whole-wave model above cannot separate 8, 16 and 24 — all three reach the 4-deep floor. A 16-lane phase (16 x 4 dwords = one full 64-bank sweep) does separate them: pad 8 gives 16 distinct start banks, pad 16 repeats each twice. A sweep over {0, 8, 16, 24, 32} confirms 8 and 16 are indistinguishable in practice while 0 and 32 are clearly worse (32 falls back to 32 banks at 8 deep). 8 is therefore chosen as the model-optimal and cheaper of the two equivalents. - New `library/tests/test_attention_dense_d64_lds.py` (CPU-only, 63 cases): the pad is inert at D=128 (identical IR *and* identical LDS, both grid variants), live at D=64, grows LDS by exactly one pad per K row-group across `block_n` 32/64/128 on both builders, and a misaligned pad is rejected. Every invariant runs on both allowed dtypes: fp16 takes the same packed group-pad path as bf16, and fp16 at D=64 is already a recorded golden case, so bf16 alone would not establish it. - Also fixed, in the same file: `_DENSE_LAUNCHER_CACHE` is keyed on `kernel_name()`, which omits `batch` — but `batch` is baked into the kernel (buffer extents, persistent work-item count). Two specs differing only in batch collided, so the second silently reused the first binary: B=1 then B=2 produced wrong output with only `1/B` of the tensor written, and B=2 then B=1 raised `hipError(700) illegal memory access`. The key is now `(kernel_name(), batch)`; keying here rather than widening `kernel_name()` leaves the emitted IR untouched. Pre-existing, reproduces on unpatched develop. - Exposed on the dense prefill benchmark (`library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py`) as `--lds-k-group-pad`, forwarded at both spec sites, so the pad can be swept directly rather than being reachable only from a unit test. Each sweep row now records the tuning knobs it was measured at, so the JSON/CSV report is self-describing instead of relying on `kernel_name()` to carry them. - The same cache-key bug exists in that script, and is fixed here too. Its private `_LAUNCHER_CACHE` keyed on `kernel_name()` alone, so it collided on `batch` *and* `waves_per_eu`. That one fires on the shipped default run rather than only in principle: the varlen sweep drives `[2048] * 4` (B=4) and `[256, 512, 768, 1024, 1536, 2048]` (B=6) at the same `max_s`, which produce an identical `kernel_name()` but different lowered IR, so on `--mode all` the B=6 config silently reused the B=4 binary. The key is now `(kernel_name(), batch, waves_per_eu)`. Pre-existing, reproduces on unpatched develop. V is deliberately left alone. The same 128-element DMA granularity blocks a per-row V pad, and a group-padded V would need a pad-aware `pv32_v_load_paired` — SSA-emitting `platform/` helper code, which pulls in the C++ mirror and a byte-identity re-bless. `ROCKE_DENSE_VPAD` is currently inert at D=64 (`VROW` is 64 at both 0 and 32), so conflict-free transposed-V remains an untapped D=64 lever, tracked as follow-up rather than bundled here. ## Risk Assessment - **D=128 codegen does not change, by construction.** Every new operation sits behind `ROWS_PER_INSTR > 1`, and the builder is side-effecting, so a guarded branch emits nothing. At `ROWS_PER_INSTR == 1` the derived values (`K_GROUP=1, K_ROWS_LDS=BN, LDROW=D+PAD, K_GROUP_BYTES=LDROW*2, V_GROUP_BYTES=1*VROW*2`) are each identical to the prior expression. Verified empirically: every D=128 shape the kernel serves in the tracked cohort lowers to identical LLVM IR before and after, under **both** `llvm20` and `llvm22`; only the D=64 cases move. - **Numerically neutral at D=64.** A 19-variant matrix (causal / non-causal, bf16 / fp16, MHA and GQA 4/8/16, persistent and default grids, `block_n` 32/64/128, sliding window, ragged, lazy-rescale off, `waves_per_eu=1`) gives `max_abs` against an fp32 oracle identical to the pre-change baseline to every digit. `varlen` is not in that matrix: it forbids `persistent` and shares the default builder's K path verbatim, so it is covered structurally rather than numerically. - **Static resources at D=64, `block_n=64`** (`llvm-readelf` on the HSACO, no launch — throttle-independent): | builder | pad | VGPR | AGPR | SGPR | spill | LDS | waves/CU | limiter | |---|---|---|---|---|---|---|---|---| | default | 0 | 141 | 0 | 46 | 0 | 32768 B | 12 | VGPR | | default | 8 | 145 | 0 | 46 | 0 | 33792 B | 12 | VGPR | | persistent | 0 | 197 | 0 | 46 | 0 | 32768 B | 8 | VGPR | | persistent | 8 | 197 | 0 | 46 | 0 | 33792 B | 8 | VGPR | Zero spill throughout. The persistent builder's register count is unchanged (the group/sub-row terms are loop-invariant and hoisted); the default builder takes +4 VGPR with no occupancy change. LDS grows 1 KB per workgroup against 160 KB/CU, and the kernel is VGPR-limited rather than LDS-limited at D=64, so the extra kilobyte does not move `waves/CU`. - **Python-only; no C++ mirror required.** This is a `library/kernels/` layout change built from existing `IRBuilder` primitives — no new IR op, no `platform/` helper or emitter touched, so the C++ *lowerer* needs no change. The #9662 / #9198 precedent that these attention kernels are not registered in the platform byte-identity gate no longer holds: develop registered them after this branch was cut. At this PR's base the representative golden held zero `attention_dense` entries across two llvm flavors; it now holds 165 across three, having gained `llvm23`. The D=64 cases are re-blessed here accordingly (see Test Result); the C++ lowerer still needs no change. - **Blast radius is bounded to an opt-in path.** `dispatch/attention.py` matches the `attention_dense` candidate only on an explicit `algorithm="attention_dense"`, never under `auto`. - **No interaction with D=256.** `__post_init__` admits only head sizes 64 and 128; D=256 is a different kernel family. ## ASIC Coverage gfx950 only. `attention_dense` is gfx950-only by construction (`supports_attention_dense` rejects every other arch), and the change is further gated behind `ROWS_PER_INSTR > 1`, i.e. head size 64. Validated on a physical gfx950 node. ## Test Plan - Byte-identity A/B over every shape the dense kernel serves in the tracked cohort, pristine vs patched, for `llvm20` and `llvm22`. - D=64 numeric vs an fp32 oracle across the 19-variant matrix, baseline vs patched, compared digit-for-digit. - K group-pad sweep {0, 8, 16, 24, 32}, forward and reverse (ABBA) to cancel ordering drift, plus a higher-rep 8-vs-16 tiebreak. - Geometry A/B at D=64: `block_n` {32, 64, 128} x `waves_per_eu` {1, 2, 3} x persistent {on, off}. - LDS bank-conflict counters (`SQ_LDS_BANK_CONFLICT`, `SQ_LDS_IDX_ACTIVE`) at the shipped vs padded layout. - Static occupancy probe (`llvm-readelf`) at D=64 for both builders, per pad. - D=128 null control: the same binary timed with the pad knob at 0 and 8, to bound same-session A/B repeatability. - `test_attention_dense_d64_lds.py`, on bf16 and fp16; `black --check`; `ruff check` (no `--fix` — the emitter-code autofix hazard `platform/AGENTS.md` describes applies equally to library kernels, and this file has pre-existing `F841`s whose assignments have op-emission side effects). ## Test Result - **Byte-identity: PASS.** All D=128 cohort shapes identical on both llvm flavors; only the D=64 cases changed; no build errors. - **Numeric: PASS, 19/19**, every `max_abs` identical to the pre-change baseline. - **LDS bank conflicts drop substantially** at the padded layout. Note the whole-wave model predicts a 4x reduction and the measured drop is larger; the counters are whole-dispatch and cover both the K `ds_read_b128` and the still-unpadded transposed-V reads, so the residual conflict is consistent with V — the lever tracked as follow-up. - **D=64 throughput improves materially** at the tracked configurations, monotonically with sequence length. Quantitative results are in the protected internal page (AICK-1828). - **Same-session repeatability**, from the D=128 null control where the pad knob is inert, is a few tenths of a percent. Absolute throughput on this part swings far more across sessions, so only same-session ratios were used. - **Static resources:** zero spill, occupancy unchanged (table above). - **Geometry: no change warranted.** The shipped `bn=64, waves_per_eu=2, persistent` stays optimal. `block_n=128` is now feasible at D=64 (it is VGPR-blocked at D=128) but does not pay for its doubled LDS. - **New test: 63/63 pass.** The 30 dtype-dependent cases now run on both bf16 and fp16; the 3 spec-validation cases stay bf16, since that path never reads the dtype. `black --check` clean. `ruff check` reports the same pre-existing findings as an unpatched build and none from this change. - **Benchmark wiring: verified as far as a CPU-only environment allows.** The script is GPU-only, so neither CI nor this environment can run a sweep. Checked by importing it against a stubbed torch: `--lds-k-group-pad` is present in `--help`, and both the success and failure record paths carry the knob fields. The varlen cache collision was confirmed directly rather than inferred, by building the B=4 and B=6 specs and comparing them: identical `kernel_name()`, different lowered-IR SHA. - **Platform golden re-blessed for D=64, and validated before writing.** The CI-wired `platform/tests/golden/rocke_representative_ir_sha256.json` now covers `attention_dense`, so the intended D=64 IR change trips it. Exactly three cases move -- `fp16_h64_sq512`, `persist_h64_sq512`, `ragged_h64_sq500` -- on each of `llvm20`, `llvm22` and `llvm23`, for 27 changed fields and nothing else. To confirm the blessing environment was trustworthy rather than merely convenient, the harness was first run against unpatched develop, where it reproduced the committed golden with zero drift on all three flavors; each re-blessed hash then matched the value CI reported independently. - **Every D=128 `attention_dense` entry in that gate is unchanged**, on all three flavors. That is an independent confirmation of the D=128 byte-identity claim above, and it extends it to `llvm23`, which was not available when the original A/B was run. - **Not run:** `test_attention_dense_golden.py::test_attention_dense_cpp_python_byte_identity` requires a built `rocke_engine`, which was unavailable in this environment. No new IR op is introduced, so no C++ change is expected; a reviewer with the engine built should confirm. ## Notes: pre-existing, not introduced here `library/tests/golden/attention_dense_ir_sha256.json` does not match a pristine build of `develop` for either llvm flavor (0 of 18 cases), which is consistent with the fixture being deliberately not CI-wired — its own docstring says so, and drift is therefore unpoliced. **This PR does not re-bless it**: doing so would mix an unrelated whole-fixture refresh into a D=64 change, and would land on an already-red baseline. The byte-identity evidence above is the pristine-vs-patched diff rather than the committed fixture. Refreshing the fixture (and deciding whether to CI-wire it) is better handled on its own. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Changes are behind the feature flag HIPKERNELPROVIDER_ENABLE_ROCKE which is defaulted to OFF. JIRA ID : AICK-1828 --------- Co-authored-by: Cursor <cursoragent@cursor.com>
AnaghaRaoAMD
pushed a commit
that referenced
this pull request
Aug 14, 2026
…ll kernel (#10337) ISSUE ID : AICK-1664 Ports the gfx950 dense flash-attention prefill kernel (`attention_dense`, #9480 / AICK-1663) to **gfx942 (MI300, CDNA3)** and productionizes it there: kernel, builder, live benchmark, dispatch arm, torch entry point, and the golden-IR + on-GPU numeric CI lanes. Opt-in only — `auto` routing and the gfx950 arm are untouched. ## Why a separate kernel, not an arch branch in the gfx950 file The gfx950 body bakes in CDNA4-only primitives, so the algorithm genuinely diverges on CDNA3. `dsl_docs/architecture/multi_arch_data_layout.md` puts that case in a per-gfx module, and it keeps the gfx950 golden IR byte-identical by construction. | Concern | gfx950 (CDNA4) | gfx942 (CDNA3) | |---|---|---| | MFMA atom | `mfma_f32_32x32x16` (K=16) | `mfma_f32_32x32x8` (K=8) → 2× per K=16 tile, A/B repack; **C-layout identical** | | Conflict-free V | `ds_read_b64_tr_b16` (transpose read) | **no `ds_read_tr16`** → `perm_b32` store-path transpose | | Cross-half exchange | `permlane32_swap` | **absent** → `perm_b32` / `ds_bpermute` | | LDS / CUs | 160 KB, one CU count | 64 KB, two CU counts → retune occupancy, `num_persistent` | | Tile barrier | bare `s_barrier` is safe (NBUF=2) | single buffer → the tile barrier **must** drain `lgkmcnt` (`sync_lds_only`) | ## Two CDNA3 hazards worth reviewer attention - **`sync_lds_only()`, not a bare `s_barrier`.** gfx90a+ set `FeatureBackOffBarrier`, so SIInsertWaitcnts skips the conservative pre-barrier drain and places the lgkm wait at the `ds_read`'s consumer, which the scheduler may sink past the barrier. Disassembly before the fix showed `ds_read_u16 x4 -> s_barrier -> s_waitcnt lgkmcnt(0)`; with NBUF=1 the next iteration's DMA can then overwrite `V_lds` under an in-flight read. There is an IR-level regression test for this, covering both the conflict-free-V and naive tile-start structures. - **`supports_attention_dense` is the single gate.** `supports(spec)[0] is True` implies `build_attention_dense(spec)` succeeds; a parametrized test asserts that contract in both directions, with a guard against the grid collapsing to one branch and making it vacuous. Divisibility invariants that guard a *silent wrong answer* raise `ValueError` rather than `assert`, because `python -O` strips asserts. ## Commit series One revertable lever per commit, each with its rationale in the body: | | | |---|---| | `refactor` | share the attention LDS-capacity accessor | | `feat` | the gfx942 dense prefill kernel + arch alias export | | `feat` | opt-in dispatch candidate | | `feat` | builder + live benchmark | | `test` | CPU surface + dispatch-wiring tests | | `perf` | lift conflict-free-V (D128 fp16) | | `perf` | `exp2_fast` in the softmax | | `perf` | fuse the softmax rescale; widen `exp2_fast` to bf16 D64 | | `perf` | tune `waves-per-eu` → 2 WG/CU at bf16 D64 | | `perf` | persistent grid-stride variant | | `perf` | adopt the D64 K-bank-conflict pad | | `feat` | `run_attention_dense_torch` entry point | | `perf` | default-off `_P0_IGLP` scheduler-hint knob | | `test` | golden-IR lane | | `test` | wire the library tests into a collectable CI lane | | `docs` | re-baseline the kernel + README to the measured bound | | `test` | correct the bf16-D128 sliced-K ring gate test *(separable — see below)* | ## Notes for the reviewer - **`kernels/gfx950/attention_dense.py` is touched (+12) only for the shared `d64_kpad` spec field** (default `False`). gfx950 codegen is unchanged and its golden is never re-blessed. This is not cross-arch creep. - **Per-config tuning.** Levers are gated per `(head_size, dtype)`, not globally; the builder `README.md` carries the full record of every lever evaluated — adopted *and* rejected — with the mechanism and the reason each won or lost. Module-level knobs (`_P0_D64_KPAD`, `_P0_IGLP`) default off, so non-dispatch builds are byte-identical. - **The last commit is deliberately separable.** `test_bf16_ring_enabled_for_d128` asserts that bf16 D128 takes the sliced-K ring; the code it tests says the opposite (`# bf16 D128 stays off the ring` → `return False`) and has since #9647/#9662. The test is red on `develop` today, independently of this branch. It is corrected here as the final commit touching one test file, so it can be dropped or split into its own PR without disturbing the series. - Per `AGENTS.md`, no measured performance data appears in the code, the commit messages, or this description. Lever effects are described qualitatively; magnitudes live in the protected results page. ## Test plan - [x] **CPU lane** — `python rocke/library/tests/run_all.py`. All gfx942 dense surface, dispatch wiring, and golden-IR tests pass. - [x] **Golden-IR (gfx942)** — new separate file + fixture; the gfx950 fixture is never written. Covers D64/D128 × bf16/fp16 × default/persistent × GQA/MHA × both persistent decodes, at both llvm20 and llvm22. Two cases are built through the dispatch factory so they guard the *shipped* tuning rather than hard-coded levers. - [x] **On-GPU numeric** — `python rocke/library/tests/run_all.py --gpu-only` → **5/5 pass** on gfx942 (ROCm 7.2, torch 2.12.1+rocm7.2), driving the public `run_attention_dense_torch` against an fp32 SDPA oracle for both grid variants. Verified on the rebased tree. - [x] **Resource guard** — 0 VGPR/SGPR spill and within the VGPR/LDS budget across all shipped configs (`builders/gfx942/attention/prefill/_p0_occupancy.py`). - [x] **Cohort parity** — validated against an fp32 SDPA reference across the in-scope cohort (bf16+fp16 × D64+D128, MHA + GQA including non-power-of-2, causal and full, batch > 1) on **both** gfx942 CU counts. - [x] `pre-commit run --from-ref origin/develop --to-ref HEAD` passes. ### Deferred / not covered by this PR - **hipDNN reachability** (epic DoD) is still open and needs coordination with the hipDNN team. - **The on-GPU numeric lane is `gpu`-marked and self-skips on CPU CI.** It gates on `gcnArchName` containing `gfx942` rather than the marketing name, so it correctly covers MI308X. A CPU-only runner will report it skipped, not passed — please read it that way. - **CI registration of `rocke/library/tests/run_all.py` is left to the CI owners** (CPU lane on every runner, `--gpu` on a gfx942 runner), consistent with the gfx950 precedent. - **No C++ engine mirror**: `library/kernels/` has none, so the dual-engine byte-identity gate does not apply here — the Python lowering is the ground truth. - D256 is out of scope (AICK-1495 / AICK-1496). Varlen, ragged and sliding-window are rejected with a structured reason by `supports_attention_dense`. ### Known-red on `develop`, not caused by this PR Running the full library suite **on a GPU runner** shows 8 failures. All 8 reproduce on a clean `develop` checkout, which shows 9 — this branch has one *fewer*, because the final commit fixes the ring-gate test. - `tests/test_attention_builds.py` (3) — fail identically on clean `develop`; untouched here. - `tests/test_lds_budget_resolver.py` (5) — a pre-existing **test-order** bug, not a code bug. The file passes 9/9 in isolation. `tests/test_attn_bf16_d128_ring.py` writes the module-level cache `_RESOLVED_ATTENTION_ARCH = "gfx942"` without restoring it on every path, while the resolver tests patch only the `_resolve_attention_arch` *function* — so they read the leaked value and assert a gfx950 LDS cap against a gfx942 one. It is invisible on CPU because the polluting tests skip without a device. Verified by pairing: this PR's tests followed by the resolver pass (14 and 133 passed); the pre-existing ring test followed by the resolver fails (5). Worth a separate `develop`-side fix; deliberately not bundled here, and deliberately **not** excluded from the new runner, since silencing them would hide `develop` breakage behind this gate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The correctness fix (#9647) put fp16 D128 prefill on the non-ring flash path — correct but ~0.55–0.99× AOTriton flash. This restores the sliced-K ring for fp16 D128 using a depth-2 pipeline that is both correct and faster, recovering (and exceeding) the ring's performance without the bug.
How it fixes the ring
The original bug: depth-3 ring,
kg % 3slot map, k_groups=4 (D128) → slice 3 reuses slot 0, and the reusing DMA was unfenced → corrupted QK accumulation. Two changes:K_SLICED_LDSSEQ.ring_depthspec field. Depth-2 (k%2, 2 slots) keeps the live set{kg, kg+1}with no slot reuse in the k_groups=4 case and uses less LDS → lower occupancy pressure. Measured best for fp16 D128.Risk assessment
Perf (fp16 D128, magnitude randn, MI300X, vs AOTriton flash)
Testing summary
Testing checklist
Technical changes
attention_tiled_2d.py—ring_depthspec field (+ kernel_name tag); unconditional drain-on-reuse fence; depth-2 (k%2) schedule in the ring loop.attention_unified.py—_select_gfx942_flash_ring_depth; ring selector re-enables fp16 D128 (depth-2), keeps bf16 D128 non-ring.attention_spec_builder.py— wirering_depthinto both flash spec builders.test_attn_bf16_d128_ring.py— per-dtype routing assertions + on-GPU numeric gate for both paths.Note: pre-existing gap (not introduced here)
The sliced-K ring (shipping since #9057/#9198) has no C++ engine implementation and no byte-identity gate coverage for attention. This PR follows that existing Python-only pattern. Closing that gap (C++ ring + attention in the gate) is separate follow-up work.
🤖 Generated with Claude Code