Repository navigation
fix(rocke/attention): exclude gfx942 D128 from sliced-K ring (correctness) - #9647
Merged
AnaghaRaoAMD merged 2 commits intoJul 22, 2026
Merged
Conversation
…ness) #9198 enabled the gfx942 bf16 D128 sliced-K ring for fp16 parity on the strength of a "max_abs 0.00049 verified" numeric check. That check used a uniform_(-0.1, 0.1) oracle whose near-uniform softmax never exercises the online-softmax running-max rescale across the four K slices (k_groups = HD/32 = 4), so the defect was masked. With realistic randn (unit-variance) inputs the D128 ring is numerically wrong -- max_abs ~0.5-1.3 vs the fp32 paged-attn oracle for BOTH bf16 and fp16, GQA and MHA, even at S512 when the ring is forced. The non-ring D128 flash path passes at max_abs ~0.0156 (bf16) / ~0.002 (fp16) on every S; the D64 ring (k_groups=2) is clean. Restore the D128 ring exclusion for BOTH dtypes (broader than #9057's bf16-only exclusion, because fp16 D128 ring is also broken). This is a Python selector-only change (the ring selector lives only in the dispatcher, not the C++ engine), so per-spec IR emission is byte-identical and the byte-identity gate is unaffected. With the ring off, the production dispatcher routes D128 to the correct non-ring geometry (bf16 nw2/T64, fp16 nw4/T64); verified 32/32 configs PASS at magnitude (bf16+fp16, GQA 32/8·64/8·128/8 + MHA 16/16, S512-4096) on MI300X. D64 ring unchanged. Rewrite test_attn_bf16_d128_ring.py: assert D128 ring OFF (both dtypes, incl. force-on env) and add an on-GPU numeric guardrail (randn inputs vs fp32 oracle, skipped without a gfx942 GPU) -- the check #9198 lacked. 30/30 pass. Co-Authored-By: Claude <noreply@anthropic.com>
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
Line-wrapping only (pre-commit black 25.12.0). No behavior change. Co-Authored-By: Claude <noreply@anthropic.com>
|
🎉 All checks passed! This PR is ready for review. |
AnaghaRaoAMD
added a commit
that referenced
this pull request
Jul 28, 2026
…liced-K ring (#9662) ## 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. - **fp16 D128** → depth-2 ring: correct at magnitude + **1.18–1.56× faster** than the non-ring path (0.65–1.55× flash; beats flash at S4096). - **bf16 D128** → stays non-ring (its bs=64 optimum; unchanged from #9647). - **D64** (both dtypes) → depth-3 ring, unchanged. ## How it fixes the ring The original bug: depth-3 ring, `kg % 3` slot map, k_groups=4 (D128) → slice 3 reuses slot 0, and the reusing DMA was **unfenced** → corrupted QK accumulation. Two changes: 1. **Drain-on-reuse fence (now unconditional).** Before a prefetch DMA that reuses the slot slice(kg-1) was just read from, drain LDS reads + barrier. Previously only fired under `K_SLICED_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. Measured best for fp16 D128. ## Risk assessment - **Python-only**, consistent with how the ring already ships: the sliced-K ring schedule lives only in the Python engine (C++ has ring primitives but no schedule loop), and attention is **not in the byte-identity gate**. #9198 enabled the D128 ring with no C++ change for the same reason — no C++ mirror or gate re-bless required. - bf16 D128 and D64 paths are unchanged; the fence refactor is byte-equivalent for k_groups=2 (D64 verified). ## Perf (fp16 D128, magnitude randn, MI300X, vs AOTriton flash) | shape | non-ring (#9647) | depth-2 ring (this PR) | speedup | |---|---|---|---| | 32/8 S1024 | 0.68× | 0.88× | 1.29× | | 32/8 S4096 | 0.92× | **1.42×** | 1.54× | | 128/8 S1024 | 0.70× | 1.06× | 1.51× | | 128/8 S4096 | 0.99× | **1.55×** | 1.56× | | 16/16 S4096 | 0.80× | **1.17×** | 1.46× | ## Testing summary - **Numeric:** fp16 D128 depth-2 ring PASS vs fp32 oracle at magnitude (max_abs ~0.002) — GQA 32/8, 128/8 + MHA 16/16, S1024–4096. - **No regression:** bf16 D128 (non-ring) and D64 depth-3 ring (both dtypes) still PASS at magnitude. - **Suite:** 177 attention tests pass. Guardrail test updated for per-dtype routing; on-GPU numeric check gates both the bf16 non-ring and fp16 depth-2 ring paths. ## Testing checklist - [x] fp16 D128 depth-2 ring numeric vs fp32 oracle at magnitude - [x] fp16 D128 perf recovered (1.18–1.56× vs non-ring; beats flash at S4096) - [x] bf16 D128 non-ring unchanged; D64 depth-3 ring unregressed - [x] launch-meta matches depth-2 spec (nw=4) - [x] 177 attention tests pass - [ ] Reviewer: confirm Python-only ring authoring is acceptable (matches #9198 precedent); C++ ring + gate coverage tracked separately ## Technical changes - `attention_tiled_2d.py` — `ring_depth` spec 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` — wire `ring_depth` into 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](https://claude.com/claude-code) [AICK-1538]: https://amd-hub.atlassian.net/browse/AICK-1538?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ [AICK-1489]: https://amd-hub.atlassian.net/browse/AICK-1489?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ Co-authored-by: Claude <noreply@anthropic.com>
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
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.
Jira: AICK-1662 (bf16 D128 ring regression from #9198) · AICK-1538 (fp16 D128 ring, pre-existing) · epic AICK-1489
Summary
Fixes numerically wrong output from the gfx942 D128 attention prefill sliced-K ring, for both bf16 and fp16, at realistic input magnitude. Excludes D128 from the ring; it falls back to the correct non-ring flash geometry.
One selector change fixes both. D64 ring is unaffected (it is correct).
Root cause
The ring stages K in 32-wide head-dim slices through a 3-slot LDS ring (
kg % 3). D128 hask_groups = HD/32 = 4, so slice 3 wraps onto slot 0 and its prefetch DMA overwrites slot 0 while slice 0's MFMA reads are still in flight — corrupting the QK accumulation. D64 (k_groups=2) never reuses a slot, so it was always clean.#9198's "max_abs 0.00049 verified" check used
uniform(-0.1, 0.1)inputs, whose near-uniform softmax never exercises the online-softmax rescale across the stale slot — so the defect was masked. It reproduces immediately withrandn(unit-variance) inputs.Risk assessment
Perf (D128, magnitude randn inputs, MI300X, vs AOTriton flash)
Testing summary
test_attn_bf16_d128_ring.pyrewritten — asserts D128 ring OFF (both dtypes, incl. force-on env) + adds an on-GPU numeric check (randn vs fp32 oracle, skipped without a gfx942 GPU) — the check feat(rocke/attention): enable gfx942 bf16 D128 prefill ring for fp16 parity #9198 lacked. 30/30 pass.Testing checklist
Technical changes
library/kernels/common/attention_unified.py—_enable_gfx942_flash_k_sliced_ring: early-returnFalsefor gfx942 D128 (both dtypes).library/tests/test_attn_bf16_d128_ring.py— assert ring-OFF for D128 + on-GPU numeric guardrail.Follow-up (stacked PR)
Recoverable perf — the fp16 D128 depth-2 ring (correct + ~0.72–1.05× flash) — will be a separate PR stacked on this one. Reference implementation preserved on
users/anarao/s34-ring-perf-scaffolding.🤖 Generated with Claude Code