Repository navigation
feat(rocke): L2-locality-aware block orders for dense attention - #12714
AnaghaRaoAMD merged 31 commits into
Conversation
✅ 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 |
|
🎉 All checks passed! This PR is ready for review. |
28e3f53 to
2ddc71e
Compare
7b8288d to
9cff6df
Compare
cef56f0 to
3ae3041
Compare
cgmillette
left a comment
There was a problem hiding this comment.
PR #12714 — L2-aware dense attention block orders (ex-rzr, AICK-2033) — 🔴 REJECT: gfx950 C++ grid drift
| Prior item | Status |
|---|---|
| AviralGoelAMD (10-09): gfx950 C++ ingestor launches the old grid | 🔴 OPEN — verified: geometry fixed at (nqb, Hq, B); 840 kdp specs unpinned; no reply |
| # | Sev | file:line | Finding |
|---|---|---|---|
| N1 | 🟡 | gfx942/gfx950 attention_dense.py, gfx942 golden |
Conflicts with open #13140 in 4 files; second lander re-blesses |
Utility: 🟢 Positive — well-tested L2-aware orders; host tests and goldens pass (ran); pin kdp, then approve.
3ae3041 to
af09420
Compare
|
Conflicts with #13140 are resolved and the bug with hipdnn integration is fixed by pinning nonpersist_decode. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #12714 +/- ##
===========================================
+ Coverage 75.88% 75.96% +0.08%
===========================================
Files 2883 2880 -3
Lines 480622 479065 -1557
Branches 72281 71977 -304
===========================================
- Hits 364704 363911 -793
+ Misses 90872 90296 -576
+ Partials 25046 24858 -188
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
af09420 to
9fc40bf
Compare
CongMa13
left a comment
There was a problem hiding this comment.
The gfx950 ingestor-grid issue raised in the existing discussion appears addressed in the current head: the catalog pins qb_minor, and the new descriptor check compares the shipped launch geometry with the ingestor geometry. I found one remaining gap in the benchmark evidence for AICK-2033.
…n_dense_signature The live benchmark and the prefill builder re-derived the kernel ABI and gated the batch/seqlen scalars on runtime_shape, but the gfx950 body declares them on _has_shape_params. The two differ for varlen, ragged, paged, sliding-window and moving bottom-right causal, so those launches packed misaligned kernargs. Use the kernel module's attention_dense_signature and _has_shape_params instead, on both the uniform and the varlen paths, and test the gfx950 signature against the built kernel's params.
Each persist_decode value is now one PersistDecode class in kernels/common/attention_dense_decode.py, holding its IR emission, legality check, kernel-name tag and fast-division divisors; PERSIST_DECODES maps names to them. gfx942 and both gfx950 decode sites call the same emit_decode, replacing three hand-written copies that had to mirror each other. gqa_pair and gqa_pair_2phase move there too and stay gfx950-only: each arch's spec lists the decodes its body supports (supported_persist_decodes). A decode reads the problem shape through DecodeShape: the runtime shape params and their fast-division kernargs in the gfx942 persistent body, constants where the shape is baked, so one emission serves both. The ISA of the gfx942 and gfx950 persistent kernels is identical before and after (qb_major with and without interleave, hkv_major, sliding window, D64, LSE output, gqa_pair*; on gfx942 also with sinks). The IR is not: each digit's quotient and remainder are now separate fast divisions, which LLVM merges, and the gfx950 epilogue re-emits the full decode, whose unused outputs LLVM removes. The 11 gfx942 persistent cases in attention_dense_gfx942_ir_sha256.json and the 14 gfx950 persistent cases in rocke_representative_ir_sha256.json are re-blessed (llvm20/22/23).
…rders gfx950 dense had no test that its variants are opt-in (gfx942 has one): assert, for every variant, that unpinned and half-pinned requests are refused and that both pins route to it. The representative IR golden had no gfx942 dense kernel: add the block orders of both grids (bt_hkv_minor, hq_minor_swz; persistent hkv_major) and re-bless at llvm20/22/23. The cross-engine parity test picks them up through the harness.
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #12714 — L2-locality-aware block orders for dense attention — 🟠 Follow-up: not merge-ready — rebase needed — the review items are fixed and CI is green; the only conflict is the representative golden, and keeping both sides' entries is exactly what the merged code produces
Utility: 🟢 Positive — L2-friendly block orders for dense attention on gfx942 / gfx950. The hipDNN ingestor's catalog stays on the grid it launches, and the new orders are pinned in the IR golden on both engines. Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: a local qwen3.6 model made the independent first pass. The orchestrator (Claude Opus 5.5) verified it. The checks were a range-diff against the head we reviewed, and a merge with today's develop whose golden was regenerated and compared entry by entry with both sides. They also covered the dense attention suites on the merge and develop, the native build, the C++ parity test on the merge, the head and develop, and the CI logs at head. The orchestrator owns this merge-worthiness call. It flags whether a closer look is needed; it is not an approval or a rejection. |
…-cache-locality Regenerate rocke_representative_ir_sha256.json on the merged tree with rocke_ir_parity_harness.py --write instead of hand-merging it. Co-Authored-By: Claude <noreply@anthropic.com>
The dense pin moved to gfx950_dense_qbminor with the L2-locality block orders (#12714); the expected rocke_spec_id now follows the manifest. Co-Authored-By: Claude <noreply@anthropic.com>
TheRock Submodule Bump ActivityNewest first
|
…metry corpus test #12714's test_real_attention_dense_bundles_launch_with_the_ingestor_geometry reads k["kernel_source"] on every shipped rocKE kernel, but a multi-launch kernel (the gfx1151 decode pack from #13897) carries `sources` instead, so hkp-pack-quick fails on develop with KeyError: 'kernel_source'. The hunk is identical to #14412's, so whichever lands first, the other merges cleanly. Co-Authored-By: Claude <noreply@anthropic.com>
The dense pin moved to gfx950_dense_qbminor with the L2-locality block orders (#12714); the expected rocke_spec_id now follows the manifest. Co-Authored-By: Claude <noreply@anthropic.com>
…orward attention on gfx942/gfx950 (#14367) ISSUE ID : AICK-2772 ## Summary Forward-attention GPU-time optimizations on gfx942 and gfx950 for both rocKE KV layouts: the unified split-KV kernels (paged KV cache, plus non-paged strided-KV decode, which runs on the same 3D kernel) and the dense kernels (contiguous, non-paged Q/K/V). Measured against the 251-shape SDPA forward suite (decode, causal prefill, chunked prefill, sliding window with sinks, FP8 KV, head_size 192). Several changes build on #14216, #12714, #14276 and #14349; all four have landed on develop and this branch merges develop, so none of them appear in the diff. The results table below was measured at `afbf60ed85c`, before the latest develop merge; a refresh against the new develop is in progress. Python and C++ engines change together; byte-identity is green at every LLVM flavor. ## Optimizations **Decode (unified split-KV)** - Split-KV for GQA ratios that do not divide the 16-row q-block (e.g. 5, 6, 12), with padding-row store masking in the 3D kernel. - Route decode to the split-KV kernel when the 2D body is occupancy-bound (gfx942 and gfx950). - Head-size-256 decode split sizing (at least two pages per segment; about two workgroups per CU on gfx942), and a decode segment cap that counts one q-block per sequence. - LLC-aware K/V cache policy: stream K/V only when the per-call K+V read exceeds the last-level cache. **Dense prefill** - gfx942: 128-row query tile when the 256-row grid under-fills the device; transposed-V path with a NaN-free RNE pack for bf16 head_size 128. - gfx950 persistent body: softmax VALU trim (folded scale FMA, pair-reduced row sums) and K/V DMA overlapped with compute through LDS alias scopes (head_size 128). - Bottom-right causal on gfx942 dense, and a KV-segment split for dense prefill on both arches with a wave-filling segment count, so chunked prefill (Sq < Skv) runs on the dense kernels. - gfx942 sliding window: narrow-window geometry, a band-phase split, and a wave-uniform skip of fully masked tiles; the 128-row under-fill rule also covers head_size 64 and sinks. - gfx950 sliding window and sinks: the single-batch 32x32 combo admits windows and sinks at head_size 64/128, and dense windows get a 128-row default tile when the 256-row grid is not already full. - gfx950 persistent wide-DMA body: 128-row tile when the 256-row work count leaves the device under-filled, and default dispatch routes under-filled plain-causal head_size 128 to it. - gfx950 head_size 256 unified prefill: query blocks run longest-first, so the costly causal tail no longer finishes last. **Coverage and dispatch** - head_size 192 in the unified tiled kernels on both arches: single-batch combo on gfx950; on gfx942, prefill on the lean 4-warp body and decode with Q staged through the V LDS slot (3 workgroups per CU instead of 2). - FP8-KV prefill on the gfx950 single-batch 32x32 combo; FP8 prefill without a 2D kernel now takes the split-KV path instead of the scalar fallback. - Dense auto-selection for contiguous-layout prefill in `dispatch_attention`; contiguous-layout decode (Sq=1) now routes to strided split-KV decode, and the strided-decode support check follows the kernel gate, so GQA ratios 5/6/12 and head_size 192 dispatch too. **Fixes found along the way** - FP8 K/V scales are now forwarded through the explicit tuning binding. - The gfx950 3D page-id K/V loader handles partial-row copies (head_size 192). - The 3D K/V LDS swizzle is rejected when the chunks per row are not a power of two. - The dense KV-segment count is capped so the fp32 workspace index stays in int32 range. - The gfx942 lean 4-warp body (head_size 256 and 192) bounds every K/V, block-table and Q read to the sequence; odd-length sequences previously read past the cache. - Sink requests on the transposed path keep scalar softmax state, so sinks land on the right rows. - `DenseSpace.defaults()` fingerprints only the knobs a shape does not record, so a tuning ID replays across chunk shapes. ## Results rocKE against itself, GPU-only, geomean speedup of the faster of the unified and dense kernel per shape (higher is better). Each cell gives geomean / median speedup and the number of shapes faster / slower by more than 2%. All arms passed the on-silicon numeric check (max |error| <= 0.05 against an fp64 reference) in the measured runs. | Arch | Config | develop -> cherry-picks (#14276, #14349) | develop -> this PR | cherry-picks -> this PR | Newly covered shapes | |---|---|---|---|---|---| | gfx950 | tuned per-shape configs | +0.1% / +0.0% (faster 8, slower 6 of 240) | **+7.0%** / +2.5% (faster 128, slower 1 of 240) | **+6.9%** / +2.7% (faster 130, slower 2 of 240) | 9 | | gfx950 | default dispatch | +0.0% / +0.0% (faster 6, slower 7 of 240) | **+34.0%** / +7.7% (faster 160, slower 1 of 240) | **+33.9%** / +7.2% (faster 165, slower 1 of 240) | 9 | | gfx942 | tuned per-shape configs | +0.5% / -0.2% (faster 5, slower 15 of 238) | **+57.5%** / +19.1% (faster 170, slower 2 of 238) | **+56.6%** / +19.7% (faster 175, slower 1 of 238) | 11 | | gfx942 | default dispatch | +0.6% / +0.0% (faster 16, slower 25 of 238) | **+57.0%** / +18.3% (faster 170, slower 6 of 238) | **+56.1%** / +18.8% (faster 169, slower 3 of 238) | 9 | By workload, cherry-picks -> this PR (gfx950 / gfx942, default dispatch unless noted): - decode: +63.3% / +148.5% (tuned configs: +0.6% / +149.4%, see caveats) - causal prefill d128/d256: +7.5% / +19.5% (tuned configs: +10.4% / +19.5%) - prefill d64: +5.1% / -0.1% (tuned configs: +4.9% / +0.9%) - chunked prefill: +50.2% / +4.4%; with the dense kernel serving Sq < Skv on gfx942, +225.6% (2.2x to 6.1x on 9 shapes, measured separately) - sliding window and sinks: +49.4% / +21.6% (tuned configs: +3.5% / +24.0%) - FP8 KV prefill: +513% on gfx950 (2 shapes); on gfx942, 2 shapes newly covered under tuned configs - head_size 192: newly covered (9 shapes per arch: 4 prefill, 5 decode) - non-paged (strided-KV) decode, measured separately on all decode shapes: +7.7% (gfx950) / +12.4% (gfx942) geomean, plus 20 shapes the cherry-pick base rejects (GQA ratios 5/6/12 and head_size 192) **Method.** Each kernel call is captured 64 times in a HIP graph; 1 warm replay and 4 timed replays, upper median per shape. Three trees: develop (`98737a5c5f4`, the develop this PR merged), cherry-picks (develop plus #14276 and #14349, `8617b561496`), and this PR (`afbf60ed85c`). Per arch and config, all three trees ran one after another in the same job, on the same node; both gfx950 jobs ran on one node and both gfx942 jobs on another, one sample per cell. On gfx942 the four FP8 shapes are excluded from the default-dispatch rows, because develop and the cherry-picks run them on the scalar fallback. Software: torch 2.13 / ROCm 10.0 wheels, LLVM 23 comgr, `ROCKE_LLVM_FLAVOR=llvm23`. "Tuned per-shape configs" pins each shape's kernel configuration from a tuning run that predates this work; "default dispatch" calls `dispatch_attention` with no pinned configuration. **Caveats.** These results hold for the measured configurations only. Because the pinned per-shape configurations predate this work, they still select the earlier decode and chunked-prefill kernels on gfx950 (decode +0.6% tuned against +63.3% default); the default-dispatch rows show those gains, and re-tuning the configurations would expose them. On gfx942, the dense path for chunked prefill is reached only when the caller routes Sq < Skv to dense; the +225.6% figure uses such a caller. ## Verification - Byte-identity gate (Python vs C++ engine LLVM IR) GREEN at llvm20, llvm22 and llvm23, all families, with new emitter configs for the GQA-ratio-5 and head_size 192 3D paths. `test_rocke_ci_static.py` 5/5. - Python vs C++ lowerer parity (`test_attention_ir_cpp_parity.py`, with the `rocke_engine` extension built): every `attention_dense` and `attention_d256` case byte-identical at all three flavors. - CPU pytest after the develop merge: `library/tests -k "attention and not bwd"` 1961 passed, 3 failed; `library/tests/dispatch` 1593 passed, 2 failed. Two of the failures were `test_unified_lse_dispatch.py` config keys that moved with the K/V cache-policy change; they are re-pinned in the last commit, and that file now passes 31/31. The remaining failure (`TestMaskScalar::test_scalar_matches_reference[mixed-no_mask]`) also fails on develop. - Off-suite GPU parity stress on both arches (non-power-of-two GQA ratios with multi-block prefill, bottom-right offsets with KV segments, sliding-window boundaries, fully masked rows, large logits, FP8, head_size 192; at least 3 seeds): 0 failures, 0 NaN/Inf; positive controls with a deliberately shifted mask fail as expected. - Known pre-existing issue, unchanged by this PR: a comgr codegen failure for gfx950 head_size 256 with block_size 64. --------- Co-authored-by: Adam Osewski <Adam.Osewski@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: dlangbe <David.Langbehn@amd.com> Co-authored-by: Anton Gorenko <anton@streamhpc.com>
The dense pin moved to gfx950_dense_qbminor with the L2-locality block orders (#12714); the expected rocke_spec_id now follows the manifest. Co-Authored-By: Claude <noreply@anthropic.com>
The dense pin moved to gfx950_dense_qbminor with the L2-locality block orders (#12714); the expected rocke_spec_id now follows the manifest. Co-Authored-By: Claude <noreply@anthropic.com>
The dense pin moved to gfx950_dense_qbminor with the L2-locality block orders (#12714); the expected rocke_spec_id now follows the manifest. Co-Authored-By: Claude <noreply@anthropic.com>
ISSUE ID : AICK-2033
Motivation
Dense attention's block orders had suboptimal L2 locality and causal balance: the
non-persistent grid had a single fixed order, and the persistent auto order wasn't tuned
for L2 reuse. On gfx942 the non-persistent grid was well behind the persistent one.
This PR adds a small set of L2-locality-aware block orders for gfx942 and gfx950, selects
them automatically, and balances causal work automatically (reverse, fold or pair fold).
The persistent and non-persistent paths now have matching knobs.
Technical Details
library/kernels/common/attention_dense_decode.py): each blockorder is one class (
grid+emit_decode) shared by both arches. Persistent decodes(
qb_major,hkv_major,gqa_pair*) move from three hand-written copies into thismodule;
gqa_pair*stays gfx950-only.DecodeShape:on gfx942's runtime-shape bodies from the shape kernargs (and, on the persistent grid,
each decode's own fast-division kernargs), elsewhere as constants. One emission serves
both; with a baked shape LLVM folds it as before. On the runtime-shape bodies the
shape-dependent parts of an order (causal balancing under a window, the pair fold's
query-block parity) are decided per launch, so one binary still serves every batch,
sequence length and head count.
nonpersist_decode, the counterpart ofpersist_decode:qb_minor: develop's grid, unchanged;bt_hkv_minor: batch and KV head fastest;hq_minor_swz: Swizzled Head-first (arXiv:2511.02132),with the new
chiplet_num_xcdsspec field (default 8).bt_hkv_minor, andhq_minor_swz(Swizzled Head-firston the persistent grid).
nonpersist_decodeis a dense tuning knob on both arches (on gfx950,of the grid body), with
--nonpersist-decodeon the builders. Thepersist_decode/nonpersist_decodeaxes and the builder choices come from each arch's spec(
supported_persist_decodes()/supported_nonpersist_decodes()).chiplet_num_xcdsis a problem field (not swept). Addingnonpersist_decodechanges thedense defaults fingerprint, so dense
tuning_ids /config_keys change. gfx950's gridbody also gets an opt-in variant fixed at
qb_minor,gfx950_dense_qbminor(see Fixes);there
nonpersist_decodeis not a knob.hq_minor_swzandhkv_majoruse a pair fold when the query-block countis even: one workgroup runs both blocks of
{p, NQB-1-p}on consecutivegrid-stride steps, so every unit costs the same;
interleavestill wins;query block about uniform): there
qb_majorandhkv_majorkeep develop's orders. Awindow at least as long counts as plain causal.
bt_hkv_minor; large causal MHA →hq_minor_swz;otherwise
qb_minor;hq_minor_swz; batch below the XCD count →bt_hkv_minor; otherwise develop's rule(
hkv_majorfor GQA with enough work per KV head, elseqb_major, now folded);gqa_pair*are no longer selected automatically;non-persistent grid at D64.
shape-dependent auto choices can't collide on the runtime-shape path.
attention_dense_signature: their own copies of the ABI broke varlen andsliding-window grid launches;
qb_minorgrid, so its catalog now comes fromgfx950_dense_qbminor(the enginemanifest's pin), and the generator config and the shipped kdp (DVC-stored, ids kept)
pin
nonpersist_decode: qb_minor. The packed kernels are unchanged (names, IR). Aquick descriptor-packaging test holds every shipped
build_attention_densespec tothe C++ launch geometry.
platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md,covering the background, measurements, L2 counters, and the alternatives that were
tried;
ALGORITHM.md, the dispatch notes and theoptimization runbook.
Test Plan
test_attention_dense_decode.pyruns every decode's realgrid/emit_decodeon integers and checks the order, coverage (bijection),reverse/fold/pair fold, the Swizzled Head-first layout, and the sliding-window rule.
Runtime shapes are checked as a runtime-shape binary runs: a decode emitted for one
spec and fed another launch's shape and fast-division kernargs must give that launch's
own mapping. It runs in the installed CTest entry
rocke_library_attention_cpu_pytest.test_attention_builds.py: kernel names, cache keys, auto-resolution tables, andruntime-shape guards.
test_tuning_space.py) and the gfx942/gfx950 dense wiring tests:field classification, production sweep, redundant-knob pruning, the new knob and the
auto choices; every gfx950 dense variant admits a request only with both
algorithmand
spec_idpinned, andgfx950_dense_qbminorstays atqb_minor.persistent representative cases; the representative golden gains gfx942 dense cases
(both grids:
bt_hkv_minor,hq_minor_swz; persistenthkv_major) and gfx950hq_minor_swzon both grids.check_byte_identity.pyGREEN at llvm20/22/23 (its ownfamily list, without
attention_dense); the cross-engine check of the harness'sattention_dense(gfx942, gfx950) andattention_d256cases,test_attention_ir_cpp_parity.py, is identical at llvm20/22/23.of the gfx950 dense catalog with the full desk check; the hipDNN gfx950 dense engine
suites on gfx950 (C++ unit, pack census, GPU knob suite, gpu-ref, engine-pinned SdpaFwd
bundles). Without the pin, the GPU suites fault in the causal kernels.
× causal/non-causal was checked against a PyTorch SDPA reference, and all orders of a
shape produce bit-identical outputs.
shapes on gfx942 (two parts) and 114 prefill + 46 decode shapes on gfx950 (prefill Sq = Skv =
512–16K, B mostly 1, MHA and GQA, D = 128 and a few D = 64, fp16/bf16, causal and
sliding window), with default dense specs; and with per-shape tuned dense specs whose
grid and block order are re-chosen for this PR (gfx942 99, gfx950 120 shapes, 6 of
them short-window shapes the default specs do not cover). Five
rounds per part, develop and this PR back to back in each, alternating order; 95%
intervals are t-intervals over the rounds, per shape and per aggregate, with an A/A
check (the same tree in consecutive rounds). Every output, both sides, every round,
checked against an fp64 reference (max abs error ≤ 5e-2).
Test Result
throughput ratios, this PR vs develop, [95% interval]):
parts, +2.7% [+2.5, +3.0] on gfx950 prefill (decode, whose kernels are
unchanged: −0.6% [−0.9, −0.4]);
default 74 / 14 / 11 and 75 / 16 / 8 on gfx942, 65 / 24 / 25 on gfx950 prefill;
tuned 66 / 20 / 13 on gfx942, 62 / 29 / 29 on gfx950. Most slower shapes are
host-bound
S = 512launches (the extra per-launch host work; gfx950S ≤ 1K−1.2% [−1.8, −0.5]), a few sliding-window shapes, and on gfx942 40-head
S = 4Kshapes (persistent tail); worst −11.3% (gfx950,
S = 512);caveats are in the case study (§6).
Submission Checklist