Skip to content

feat(rocke): L2-locality-aware block orders for dense attention - #12714

Merged
AnaghaRaoAMD merged 31 commits into
developfrom
users/agorenko/rocke/improve-attention-l2-cache-locality
Oct 26, 2026
Merged

AnaghaRaoAMD merged 31 commits into
developfrom
users/agorenko/rocke/improve-attention-l2-cache-locality

Conversation

@ex-rzr

@ex-rzr ex-rzr commented Sep 29, 2026 •

Copy link
Copy Markdown
Contributor

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

  • Decode classes (library/kernels/common/attention_dense_decode.py): each block
    order 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 this
    module; gqa_pair* stays gfx950-only.
  • Runtime and baked shapes: a decode reads the problem shape through 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.
  • New non-persistent knob nonpersist_decode, the counterpart of persist_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_xcds spec field (default 8).
  • New persistent decodes: bt_hkv_minor, and hq_minor_swz (Swizzled Head-first
    on the persistent grid).
  • Tuning space: nonpersist_decode is a dense tuning knob on both arches (on gfx950,
    of the grid body), with --nonpersist-decode on the builders. The persist_decode /
    nonpersist_decode axes and the builder choices come from each arch's spec
    (supported_persist_decodes() / supported_nonpersist_decodes()).
    chiplet_num_xcds is a problem field (not swept). Adding nonpersist_decode changes the
    dense defaults fingerprint, so dense tuning_ids / config_keys change. gfx950's grid
    body also gets an opt-in variant fixed at qb_minor, gfx950_dense_qbminor (see Fixes);
    there nonpersist_decode is not a knob.
  • Automatic causal balancing:
    • non-persistent orders run query blocks longest-first;
    • persistent orders fold;
    • persistent hq_minor_swz and hkv_major use a pair fold when the query-block count
      is even: one workgroup runs both blocks of {p, NQB-1-p} on consecutive
      grid-stride steps, so every unit costs the same;
    • an explicit interleave still wins;
    • none of this under a sliding window shorter than the keys (it makes the work per
      query block about uniform): there qb_major and hkv_major keep develop's orders. A
      window at least as long counts as plain causal.
  • Auto selection:
    • non-persistent: aligned causal → bt_hkv_minor; large causal MHA → hq_minor_swz;
      otherwise qb_minor;
    • persistent, aligned causal: MHA with more than 8 grid-stride rounds of work →
      hq_minor_swz; batch below the XCD count → bt_hkv_minor; otherwise develop's rule
      (hkv_major for GQA with enough work per KV head, else qb_major, now folded);
    • gqa_pair* are no longer selected automatically;
    • grid choice: gfx942's default dense spec takes the persistent grid at D128 and the
      non-persistent grid at D64.
  • The resolved block orders are part of the kernel name and the cache key, so
    shape-dependent auto choices can't collide on the runtime-shape path.
  • Fixes:
    • the gfx950 dense live benchmark and prefill builder bind arguments via
      attention_dense_signature: their own copies of the ABI broke varlen and
      sliding-window grid launches;
    • the hipDNN gfx950 dense ingestor engine launches every packed kernel with the
      qb_minor grid, so its catalog now comes from gfx950_dense_qbminor (the engine
      manifest'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). A
      quick descriptor-packaging test holds every shipped build_attention_dense spec to
      the C++ launch geometry.
  • Docs:
    • a new case study,
      platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md,
      covering the background, measurements, L2 counters, and the alternatives that were
      tried;
    • updates to the prefill READMEs, gfx950 ALGORITHM.md, the dispatch notes and the
      optimization runbook.

Test Plan

  • Host tests: a new test_attention_dense_decode.py runs every decode's real
    grid/emit_decode on 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, and
    runtime-shape guards.
  • Dense tuning space (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 algorithm
    and spec_id pinned, and gfx950_dense_qbminor stays at qb_minor.
  • IR goldens re-blessed at llvm20/22/23 for the intended changes: gfx942 dense, and gfx950
    persistent representative cases; the representative golden gains gfx942 dense cases
    (both grids: bt_hkv_minor, hq_minor_swz; persistent hkv_major) and gfx950
    hq_minor_swz on both grids. check_byte_identity.py GREEN at llvm20/22/23 (its own
    family list, without attention_dense); the cross-engine check of the harness's
    attention_dense (gfx942, gfx950) and attention_d256 cases,
    test_attention_ir_cpp_parity.py, is identical at llvm20/22/23.
  • hip-kernel-provider: descriptor-packaging and IngestorGenerator suites; the real pack
    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.
  • GPU: the dense numeric suites pass on gfx942 and gfx950; every decode
    × causal/non-causal was checked against a PyTorch SDPA reference, and all orders of a
    shape produce bit-identical outputs.
  • Performance, develop vs this PR, end to end through the dispatcher: 99 dense prefill
    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

  • All tests listed above pass.
  • Dense attention is faster at the measured configurations (geomean of per-shape
    throughput ratios, this PR vs develop, [95% interval]):
    • default specs: +8.5% [+8.4, +8.6] and +9.3% [+8.8, +9.8] on the two gfx942
      parts, +2.7% [+2.5, +3.0] on gfx950 prefill (decode, whose kernels are
      unchanged: −0.6% [−0.9, −0.4]);
    • tuned specs: +9.0% [+8.6, +9.3] on gfx942, +1.4% [+1.1, +1.8] on gfx950;
    • per shape, by whether the 95% interval excludes 0 (faster / within noise / slower):
      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 = 512 launches (the extra per-launch host work; gfx950 S ≤ 1K
      −1.2% [−1.8, −0.5]), a few sliding-window shapes, and on gfx942 40-head S = 4K
      shapes (persistent tail); worst −11.3% (gfx950, S = 512);
    • A/A: a whole round moves the geomean by at most 0.9%;
    • numerically unchanged everywhere (same max abs error as develop on every shape).
  • Per-slice results with intervals, the earlier sweeps' intervals, L2 counters and
    caveats are in the case study (§6).

Submission Checklist

@ex-rzr ex-rzr added the organization: streamhpc contributors from streamhpc label Sep 29, 2026
@github-actions github-actions Bot added documentation project: hip-kernel-provider rocke-lib: attention Touches rocKE library attention/FMHA kernels rocke: platform Touches rocKE platform labels Sep 29, 2026
@therock-pr-bot

therock-pr-bot Bot commented Sep 29, 2026 •

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

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

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

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

🙋 Wish to Override Policy?

@therock-pr-bot

Copy link
Copy Markdown

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

  • python -m pip install pre-commit
  • pre-commit install
  • pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

@therock-pr-bot

therock-pr-bot Bot commented Sep 29, 2026 •

Copy link
Copy Markdown

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

@ex-rzr
ex-rzr force-pushed the users/agorenko/rocke/improve-attention-l2-cache-locality branch from 28e3f53 to 2ddc71e Compare September 30, 2026 06:33
@ex-rzr
ex-rzr force-pushed the users/agorenko/rocke/improve-attention-l2-cache-locality branch from 7b8288d to 9cff6df Compare October 8, 2026 11:28
@ex-rzr
ex-rzr marked this pull request as ready for review October 8, 2026 13:23
@ex-rzr
ex-rzr requested review from a team as code owners October 8, 2026 13:23
@ex-rzr
ex-rzr force-pushed the users/agorenko/rocke/improve-attention-l2-cache-locality branch from cef56f0 to 3ae3041 Compare October 9, 2026 12:01

@cgmillette cgmillette left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@ex-rzr
ex-rzr force-pushed the users/agorenko/rocke/improve-attention-l2-cache-locality branch from 3ae3041 to af09420 Compare October 12, 2026 09:05
@ex-rzr
ex-rzr requested a review from a team as a code owner October 12, 2026 09:05
@ex-rzr

ex-rzr commented Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Conflicts with #13140 are resolved and the bug with hipdnn integration is fixed by pinning nonpersist_decode.

@codecov

codecov Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

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     
Flag Coverage Δ *Carryforward flag
TensileLite-CPP 48.93% <ø> (ø) Carriedforward from bb21093
TensileLite-Unit 83.68% <ø> (-0.04%) ⬇️ Carriedforward from bb21093
hipBLAS 90.62% <ø> (ø) Carriedforward from bb21093
hipBLASLt 64.43% <ø> (-1.52%) ⬇️ Carriedforward from bb21093
hipCUB 82.68% <ø> (ø) Carriedforward from bb21093
hipDNN 86.85% <ø> (ø)
hipFFT 42.00% <ø> (-0.56%) ⬇️ Carriedforward from bb21093
hipRAND 76.12% <ø> (ø) Carriedforward from bb21093
hipSOLVER 69.67% <ø> (ø) Carriedforward from bb21093
hipSPARSE 86.99% <ø> (ø) Carriedforward from bb21093
rocBLAS 68.73% <ø> (+0.01%) ⬆️ Carriedforward from bb21093
rocFFT 48.10% <ø> (+2.42%) ⬆️ Carriedforward from bb21093
rocRAND 57.30% <ø> (-0.01%) ⬇️ Carriedforward from bb21093
rocSOLVER 76.74% <ø> (ø) Carriedforward from bb21093
rocSPARSE 76.92% <ø> (-0.24%) ⬇️ Carriedforward from bb21093
rocThrust 93.15% <ø> (ø) Carriedforward from bb21093

*This pull request uses carry forward flags. Click here to find out more.
see 128 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@ex-rzr
ex-rzr force-pushed the users/agorenko/rocke/improve-attention-l2-cache-locality branch from af09420 to 9fc40bf Compare October 13, 2026 05:31

@CongMa13 CongMa13 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@AnaghaRaoAMD
AnaghaRaoAMD requested a review from yraparti October 23, 2026 16:15
@cgmillette

Copy link
Copy Markdown
Contributor

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

Prior item Status
Ours (10-21) 🟡: the re-uploaded DVC kdp could not be checked from the repo 🟢 Addressed by a test rather than a command. The body says the generator config and the shipped kdp pin nonpersist_decode: qb_minor, and test_desk_check_rocke_corpus.py now holds every shipped dense spec to the C++ launch geometry (read). We did not run that test here: it needs the hipDNN tree, which this checkout omits
yraparti (10-22, 10-23) §Q.2 HIGH: the new block orders lacked harness cases (gfx950, then gfx942) 🟢 FIXED in f8f887a5 (gfx950 hq_minor_swz, both grids) and 9bce73b1 (gfx942 bt_hkv_minor / hq_minor_swz / hkv_major), re-blessed at llvm20 / 22 / 23. The regenerated golden on the merge confirms every new entry (ran)
yraparti (10-23) §D.11 HIGH: gfx950_dense_qbminor admits unpinned requests 🟢 Not reproduced. The variant is built by make_tuned_candidate, which refuses unless both algorithm and spec_id are pinned. 9bce73b1 adds test_every_variant_needs_both_pins, which passes (ran)
yraparti (10-23) §Q.1 HIGH: name the lane that runs test_attention_dense_decode.py 🟢 FIXED: it is in rocke_library_attention_cpu_pytest (f8f887a5) and named in the Test Plan. CI runs it: 94 cases PASSED in a gfx94X shard (read)
yraparti (10-22) §E.15 LOW: unused shape parameter 🟢 FIXED in f8f887a5
New 7da712b0: a sweep or pinned tuning id could move the ingestor's qbminor variant off its qb_minor grid 🟢 Real fix found by the author. The knob now applies only where the variant leaves the order to auto; test_qbminor_sweep_keeps_the_ingestor_grid passes (ran)
CongMa13 APPROVED (10-19) at badbe571 🟡 Given before the last four commits, including the qbminor fix and the new goldens
Bots / CI 🟢 Green at head eb97120d: TheRock Multi-Arch (all four Test hipkernelprovider shards, including rocke_pytest and the attention CPU, host and GPU lanes), ASAN and the hipDNN superbuilds (read). GitHub now reports the PR as conflicting
Our checks 🟢 Range-diff from 188573a1: commits 1–20 unchanged, 3 new; the rest are develop merges (ran). On a merge with today's develop (golden resolved as N1 says): the dense attention suites pass, 576 vs develop's 457, 0 failures. Native build and 30 ctests pass. The C++ attention parity test fails at the PR head and on develop identically: one case's IR fails to parse in the engine on this host, so it is not caused by this PR (ran). rule_scan is clean (ran)
# Sev file:line Finding
N1 🟠 platform/tests/golden/rocke_representative_ir_sha256.json The only conflict with today's develop. Develop added its own cases since the base, so the resolution is: keep all of develop's entries, take this PR's 24 re-blessed attention_dense/gfx950 values and 7 new entries, and regenerate the summary block. On the merge, a fresh --write produces exactly that set at llvm20 / 22 / 23: no other entry drifts (ran) → rebase, take both sides, and re-run --write to refresh the summary.

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>
@AnaghaRaoAMD
AnaghaRaoAMD enabled auto-merge (squash) October 25, 2026 19:40
@AnaghaRaoAMD
AnaghaRaoAMD merged commit 98737a5 into develop Oct 26, 2026
175 of 179 checks passed
@AnaghaRaoAMD
AnaghaRaoAMD deleted the users/agorenko/rocke/improve-attention-l2-cache-locality branch October 26, 2026 06:22
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Oct 26, 2026
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>
@assistant-librarian

Copy link
Copy Markdown
Contributor

TheRock Submodule Bump Activity

Newest first

ThruptiRajLakshmanaGowda added a commit that referenced this pull request Oct 26, 2026
…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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Oct 30, 2026
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>
AnaghaRaoAMD added a commit that referenced this pull request Oct 30, 2026
…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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Nov 1, 2026
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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Nov 2, 2026
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>
ThruptiRajLakshmanaGowda added a commit that referenced this pull request Nov 3, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation organization: streamhpc contributors from streamhpc project: hip-kernel-provider project: hipdnn rocke: platform Touches rocKE platform rocke-lib: attention Touches rocKE library attention/FMHA kernels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants