Repository navigation
perf(rocke): runtime heads, window and batch for gfx942/gfx950 tiled attention - #15024
Merged
Merged
Conversation
…ime (#14530) JIRA ID : AICK-2676 The gfx942 tiled attention kernels (2D, 4-warp GQA, 3D split-KV and reduce) now read the head counts and the sliding-window size as kernargs and derive the `cu_q` search trip count from `num_seqs`. One compiled binary therefore serves every head pair at a given GQA ratio, every window size and every batch. This is the kernel-side readiness hipDNN AOT packaging needs: a packaged catalog ships one binary per kernel symbol and cannot enumerate head counts, window sizes or batch sizes. Follow-up to #13372 and #13889, which removed raw shape fields from the tiled and scalar cache keys. Builds on #14216 and #14349, merged from develop. The gfx950 counterpart (AICK-2675) reuses the plumbing added here. ## Risk Assessment Risk 4. Every gfx942 tiled kernel changes its argument list and emitted IR, the tiled cache keys change, and the paged K/V descriptor gains runtime strides in both engines. That is a performance-critical path. Mitigations: byte identity is GREEN at llvm20, llvm22 and llvm23 for every family; gfx950 and gfx1250 tiled IR is unchanged (golden); and the gfx942 numerics and an A/B against develop were run on hardware (see Testing). ## ASIC Coverage - **gfx942:** all kernel changes. Verified manually on gfx942: numeric tests, the existing gfx942 GPU suites, and a develop-vs-branch A/B. Superbuild and TheRock CI do not run these GPU numerics on gfx942. - **gfx950 / gfx1250:** host-only changes (signature helpers, cache keys, launch-cache keys). Their tiled spec classes declare no runtime kernargs, so their IR and goldens do not move. PR CI is sufficient. ## Testing Summary - **Byte identity:** full `check_byte_identity.py` GREEN at llvm20, llvm22 and llvm23. The gfx942 2D, 2D stress, 3D and reduce emit configs cover windowed, i64 and sinks specs. - **Goldens:** the representative-IR golden moves only for the gfx942 tiled cases (2D, 3D, reduce and both 4-warp kernels); the strided-KV library golden only for its gfx942 configs. Every gfx950 and gfx1250 entry is unchanged. - **Cache keys (CPU):** - shapes that differ only in head pair (same GQA ratio), window size or batch bucket share one key and lower to byte-identical IR at every flavor, on 2D, 4-warp D256, the 4-warp D128 fold and 3D; - the invariant sweep holds; - gfx950 and gfx1250 still split on the batch bucket. - **Launch signature (CPU):** every production builder route and three gfx942 tuning candidates are checked: the kernel's params equal its launch signature by name and type. - **One binary, many shapes (gfx942 GPU):** `test_attention_gfx942_runtime_shape_numeric.py` runs pairs of shapes that share a key back to back through `run_unified_attention_torch`. The second launch must compile nothing, and both outputs must match an fp32 paged reference (max_abs ≤ 4e-2; bf16/fp16 inputs, randn) with unused block-table slots on a poison block. - Pairs cover 2D default, windowed 2D, D256 4-warp, the D128 window fold, 3D decode and two-token 3D. One pair changes only the head counts at identical token counts. - Checked against mutants: baking the window back in fails the windowed pair (max_abs 0.88); baking the kv head count into the K/V stride faults on the device. - **Existing GPU suites (gfx942):** 69 attention, MLA and strided-KV library test files. Results match develop; the remaining failures fail identically on develop (MLA comgr environment errors, strided-KV and selector-parity cases). **A/B, develop vs this branch (gfx942, 304 CUs)** - **Workload:** 116 problems through `run_unified_attention_torch` with production routing (`backend="auto"`): 98 on the 2D path, 18 on 3D. Cohorts cover d64/d128/d256, bf16/fp16 and fp8 K/V, GQA ratios 1, 4 and 8, sinks, sliding window, the 4-warp D256 and fold paths, and a 64/16 head pair. - **Method:** HIP-event mean over 50 launches after 10 warmup, median over 3 rounds. Arm order alternates per round, and each arm runs in a fresh process each round. - **Toolchain:** torch 2.13.0+rocm10.0.0, LLVM 23 from the ROCm 10 SDK wheels, Python lowering. - **Baseline:** develop at `f37a7056c791`. - **Numeric gate:** every row of both arms is within 4e-2 of an fp32 reference; worst max_abs is 1.24e-2 in both arms. - **Results (branch / develop):** | Measure | Result | |---|---| | Compiled binaries per process | 18 vs 54 | | Total first-launch time | 0.37× | | Steady state, geometric mean over 116 problems | 0.994× (range 0.87–1.06×) | | Per cohort | within 0.98–1.01×, except the two below | | d64 GQA-8 block-32 **sinks** (2 problems, 2D) | 1.05–1.06× | | d64 GQA-8 **fp8 K/V** decode (4 problems, 3D) | 1.01–1.03× | - **Sinks caveat, isolated:** - Timing each commit puts the whole sinks cost in the head-count commit, and only on the narrow 2D kernel. - A timing-only variant with the K/V token stride restored to a constant removes it; restoring the query head count changes nothing. - The per-copy K/V offset now multiplies the token index by the runtime stride inside the KV loop, where the constant stride used to compile to a shift. - Accepted here as a known cost. Hoisting the in-page offsets out of the loop belongs with the planned per-page i32 K/V loader rework. **AOT dry run:** a draft engine manifest pinned to one gfx942 tuning candidate (`transposed_x8`, 4 warps, tile 4× block) expands with `registry_enumerator.py --plan-space`: 24 cells (6 head configurations × bf16/fp16 × block 16/32), production sweep, 20 varying knobs. Result: - 324 instances (within the default per-arch budget of 1000); - 162 distinct kernels on this branch against 324 on develop, because head pairs at one GQA ratio now build the same kernel. ## Testing Checklist - [x] Byte identity - `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py` - Status: Passed - [x] Representative-IR golden - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json` - Status: Passed - [x] Library host tests (attention, MLA, strided KV, dispatch, layering), GPU hidden - `python -m pytest <file>` per file - Status: Passed (3 files fail identically on develop) - [x] Platform dynamic-descriptor tests - `python -m pytest tests/core/test_transforms_dynamic.py` - Status: Passed - [x] gfx942 GPU: new runtime-shape numeric test - `python -m pytest tests/test_attention_gfx942_runtime_shape_numeric.py` - Status: Passed (14/14) - [x] gfx942 GPU: existing attention/MLA/strided-KV library suites (69 files) - Status: Passed, matching develop - [x] gfx942 GPU: develop-vs-branch A/B, 116 problems, 3 rounds - Status: Passed (numeric gate), caveats above ## Technical Changes - **Runtime-kernarg plumbing.** Tiled spec classes declare `RUNTIME_KERNARG_FIELDS`. `_tiled_2d_signature` / `_tiled_3d_signatures` append those fields to the launch signature, and `_runtime_kernarg_values` packs them. Every gfx942 launch site uses the arch-aware signature. The 4-warp GQA kernels now declare `qq_bias_stride_0`, which the shared 2D signature already passed. - **Search trip count.** `binary_search_seq_idx_trip_count` computes `bit_length(num_seqs)` from the `num_seqs` kernarg. It's exact below 2^24 sequences; past that it adds at most one no-op iteration. Specs declare `RUNTIME_SEARCH_ITERS` and the keys drop the batch bucket. The existing helper and its other callers emit the same IR. - **Sliding window.** The size is a kernarg. Whether a window applies still selects structure, so the name token is `sw` and the keys carry `window > 0`. - **Head counts.** `num_query_heads` and `num_kv_heads` are kernargs. The Q/O, segment-output and paged K/V strides are runtime values through `DynamicTensorDescriptor`, which now takes mixed int/Value strides plus runtime strides in `offset_i64_split` (both engines) and `offset_i64` (Python). Names carry `g<Hq/Hkv>`, and the keys carry the GQA ratio. Every selector that reads absolute head counts feeds an output that's already keyed. - **Launch caches.** `_shape_key` (2D grid, 3D grids and workspace) now carries both head counts, since one binary now serves several head pairs. - **i64 name token.** The gfx942 2D name carries `i64` when `use_i64_kv_addr` is set; i32 and i64 builds used to share a symbol. - **C++ twins.** The 2D, 3D and reduce twins mirror all of the above. The gfx942 2D C++ entry also builds the gfx950 fastKV register-P twin, so its runtime changes are gated to gfx942. - **CI enrollment.** `test_attention_tiled_launch_signature.py` joins the installed attention CPU cohort; `test_attention_gfx942_runtime_shape_numeric.py` joins the library GPU cohort. **ABI note:** in this repo, the gfx942 tiled kernels are launched only through `run_unified_attention_torch`, the dispatch tuning candidates and the harnesses updated here. Anything that launches them by name must append `num_query_heads`, `num_kv_heads` and `sliding_window` (2D and 3D) or `num_query_heads` (reduce). **Notes for the packaging work:** - The gfx942 fast paged-KV descriptor path (`use_fast_paged_kv_desc`, reachable only through the `fastkv` tuning knob) still hard-codes 64/8 heads, pinned by its spec validator. A packaged catalog must serve it only at 64/8 or leave that knob out. - gfx942 routes windowed decode to 2D, so production never takes the windowed 3D path; byte identity covers it. - `AttentionTuningSpec.cache_key()` remains field-complete, so it still splits on the runtime fields. That over-keys but is never wrong. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.com>
…ime (#14727) JIRA ID : AICK-2675 ## Motivation hipDNN AOT packaging ships one compiled binary per kernel symbol, and a packaged catalog cannot enumerate head counts, sliding-window sizes or batch sizes. The gfx950 tiled attention kernels (2D, fastKV register-P 2D, 3D split-KV and reduce) baked all three into their IR. So every head pair, every window and every batch bucket compiled its own binary, and the fastKV register-P builder could not be packed at all because it had no spec-first support predicate. This PR applies to gfx950 what #14530 (AICK-2676) did for gfx942. The kernels read the head counts and the window as kernargs and derive the `cu_q` search trip count from `num_seqs`. One gfx950 binary now serves every head pair at a given GQA ratio, every window size and every batch, and kernel names and cache keys carry exactly what stays baked. Stacked on #14530, which adds the shared plumbing this reuses: `RUNTIME_KERNARG_FIELDS` / `RUNTIME_SEARCH_ITERS`, the arch-aware launch signatures, `binary_search_seq_idx_trip_count`, and mixed/i64 runtime strides in `DynamicTensorDescriptor`. ## Technical Details One commit per change; every kernel change is mirrored in the C++ engine in the same commit. - **Batch** (`3e02491`): the gfx950 2D, fastKV register-P and 3D segment kernels take the search trip count from `num_seqs` at runtime, and the tiled cache keys drop the batch bucket. The single-batch schedule flags in the 2D key now collapse to `None` when neither is set, so `num_seqs == 1` and `> 1` no longer split identical kernels. - **Sliding window** (`ddc344e`): the window size is a trailing i32 kernarg. Whether a window applies still selects structure, so the name token becomes `sw` and the keys carry `sliding_window > 0`. The gfx950 benchmark, parity and build-test launch sites switch to `_tiled_2d_signature` / `_tiled_3d_signatures`. - **Head counts** (`74e8aa5`): - The kernels read `num_query_heads` / `num_kv_heads` as kernargs and derive the Q/O, segment-output and paged K/V strides at runtime. Only the GQA ratio stays baked: names carry `g<Hq/Hkv>`, the reduce name drops `h<Hq>`, and the keys carry the ratio. - The fast paged-KV descriptor keeps its pinned 64/8 geometry. `_enable_fast_paged_kv_desc` now resolves it for both the spec builder and the 2D key, so a 32/4 problem at ratio 8 never reuses a 64/8 fast binary. - Every caller of the shared gfx942 2D C++ prologue (which also builds the gfx950 fastKV register-P twin) now reads its shape at runtime, so its `runtime_shape` / `runtime_search` / `runtime_window` gates and the unused 2D `binary_search_iters` helper are gone. - **K/V copy offsets and page-id prefetch** (`25bac46`): with a runtime kv head count, every bf16/fp16 K/V copy in the 2D loop multiplied its token index by the runtime stride. - When a copy's tokens sit in one page, its offset is now built from three parts: - the page base, passed in `soffset`; - the copy's first-token offset, built once before the loop; - the lane's own part, also built once before the loop. - The KV loop carries the page ids of the current and next tile and reads the ids of the tile after next right after the iter-start drain, as the 3D page-id loader does. - Copies that span pages, the fp8 loaders, and the K-ring / grouped-KV2 schedules keep their existing addressing. - New emit config 56 (the shipping bf16 d64/b16 multi-batch spec) keeps byte-identity coverage on the page-spanning path. - **Packaging predicate** (`12d8242`): `supports_fastkv_register_p_2d_spec` reads its parameters off the spec and also refuses what the builder raises on; it shares those checks with the builder. `hkp_pack` maps the builder to it in `_PREDICATE_ALIASES`. - **i64 name token** (`6569f53`): the gfx950 2D and fastKV register-P names carry `i64` when `use_i64_kv_addr` is set, matching gfx942 2D. gfx950 3D already carries `i64kv`. - **Tests** (`5d10262`, `802b732`): - new `test_attention_gfx950_runtime_shape_numeric.py`, enrolled in the library GPU test list in `platform/CMakeLists.txt`; - the cache-key invariant sweep now has two head pairs per GQA ratio and a second window size, so it catches an under-key between them; - new targeted key and launch-geometry tests for gfx950. ## Risk Assessment Risk 4. Every gfx950 tiled kernel changes its argument list and emitted IR, the tiled cache keys change, and the 2D K/V loaders and KV-loop carry change in both engines. That is a performance-critical path. Mitigations: - byte identity is GREEN at llvm20, llvm22 and llvm23 for every family; - only gfx950 golden entries move; - gfx950 numerics and a baseline-vs-branch A/B were run on hardware (below). ## ASIC Coverage - **gfx950:** all kernel changes. Verified manually on gfx950: the new numeric test, the existing gfx950 GPU suites, and an A/B against #14530's head. Superbuild and TheRock CI do not run these GPU numerics on gfx950. - **gfx942:** C++-only cleanup of the shared 2D prologue (constant branches nobody calls any more); the gfx942 Python IR and goldens do not move. - **gfx1250:** test parametrization only. ## Test Plan - **Byte identity:** `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py`, every family. - **Goldens:** - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json`; - the strided-KV library golden; - both re-blessed only where a commit moves gfx950 IR. - **Library host tests (GPU hidden):** `python -m pytest <file>` per file over the attention, MLA, strided-KV, dispatch and layering set (69 files), plus `descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py`. - **Cache keys (CPU):** - the invariant sweep (key equality ⇔ build equality over every selector threshold); - targeted share/split pairs; - an IR-mode A/B over 311,040 gfx950 problems, checking routing and that no problem is served a binary built for a different spec. - **gfx950 GPU numerics:** - `test_attention_gfx950_runtime_shape_numeric.py`: pairs of shapes that share one key run back to back through `run_unified_attention_torch`. The second launch must compile nothing, and both must match an fp32 paged reference (max_abs ≤ 4e-2; bf16/fp16 randn inputs) with unused block-table slots on a poison block. - Pairs cover 2D default, a window, the single-batch combo with and without a window, the 64/8 fast descriptor (batch only) and its 32/4 neighbour, the D256 fast route, 3D decode and two-token 3D. - **Existing gfx950 GPU suites:** unified-mask, additive-bias, LSE, 3D running-max, KV-LDS swizzle, wave page-id load, strided-KV decode and dense. - **A/B:** #14530's head vs this branch, run end to end through `run_unified_attention_torch` with production routing (`backend="auto"`). - **AOT dry run:** a draft gfx950 2D engine manifest, pinned to one tuning candidate, run through `registry_enumerator.py --plan-space`. ## Test Result - **Byte identity:** GREEN at llvm20, llvm22 and llvm23. - **Goldens:** only gfx950 entries move, at every flavor. Representative-IR cases per commit: - batch: 21; - window: 21; - head counts: 24 (adds the 3 reduce cases); - K/V loader: 15 (2D only). - The 7 gfx950 strided-KV configs move with the batch, window and head-count commits. - No case is added or removed, and no gfx942 or gfx1250 entry changes. - **Host tests:** results match develop. The 26 remaining failures in `test_strided_kv_lds.py`, `test_strided_kv_decode.py` and `test_selector_parity.py` fail identically on develop. - **hkp_pack:** `test_hkp_pack_producer_guards.py` passes (24/24), including the fastKV register-P accept/refuse case. The rest of `descriptor-packaging/tests/rocke/` needs the hipDNN plugin SDK data, so it is left to PR CI. - **IR-mode A/B (CPU):** - all 311,040 problems route identically; - no problem is served another spec's binary; - distinct binaries for the sweep drop from 10,662 to 2,094 (2D) and from 5,040 to 3,216 (3D). - **gfx950 GPU numerics:** - runtime-shape test 18/18; - unified-mask 315 passed (16 skipped, other arches); - additive-bias 185/185, LSE 57/57, dense 322/322; - 3D running-max, KV-LDS swizzle, wave page-id load and strided-KV decode pass (other-arch cases skipped). - The 3D-only suites ran before the 2D loader commit, which does not touch 3D. **A/B, #14530 head vs this branch (gfx950, 256 CUs)** - **Workload:** 132 problems across 11 cohorts: - d64/d128/d256, bf16/fp16 and fp8 K/V; - GQA ratios 1, 4 and 8, including 64/8 vs 32/4 and a 64/16 head pair; - sinks, sliding windows (128, 256, 4096 fold), the D256 fast route and the single-batch combo; - 2D and 3D paths. - **Method:** HIP-event mean over 50 launches after 10 warmup, median over 3 rounds. Arm order alternates per round, and each arm runs in a fresh process each round. - **Toolchain:** torch 2.10.0+rocm7.0 wheels with the wheel's comgr (LLVM 20 flavor), Python lowering, the same for both arms. - **Baseline:** #14530's head, `c620a3ee6add`. Candidate: this branch at `25bac467c15a`. - **Numeric gate:** every row of both arms is within 4e-2 of an fp32 reference. Worst max_abs is 1.54e-2 in both arms, and max_abs is identical between arms on every row. - **Results (branch / baseline):** | Measure | Result | |---|---| | Compiled binaries per process | 35 vs 73 | | Total first-launch time | 0.32× | | Steady state, geometric mean over 132 problems | 0.980× (range 0.82–1.06×) | | D256 fast route (6 problems) | 0.947× | | bf16 d128 GQA-4 (40 problems) | 0.961× | | bf16 d128, 64/16 heads (12 problems) | 0.978× | | Sliding-window cohorts (3 cohorts, 7 problems) | cohort geomeans 0.887–0.955× | | d64 GQA-8 block-32, incl. the 64/8 fast descriptor (35 problems) | 1.007× | | Other cohorts | within 0.99–1.02× | - **Outliers:** one problem is above 1.05×: bf16 d128 GQA-4, 4 sequences, q 257, k 4096, at 1.064×. - **Without the loader commit:** per-commit timing on the same harness put the head-count change at 1.49–1.82× on the D256 fast route and 1.05–1.10× on long multi-batch d128 prefill, all from the per-copy runtime stride multiply. The `25bac46` loader change recovers both. **AOT dry run (`--plan-space`)** - **Manifest:** a draft gfx950 2D engine manifest pinned to `gfx950_u2d_transposed32_nw4_mw32_t2xb_llvm`, 24 cells (6 head configurations × bf16/fp16 × block 16/32), with the 21 varying gfx950 knobs classified as performance fields. - **Result:** 445 instances, within the default per-arch budget of 1000. They build 211 distinct kernels on this branch against 421 on the baseline, because head pairs at one GQA ratio now share a kernel. ## Testing Checklist - [x] Byte identity - `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py` - Status: Passed - [x] Representative-IR golden - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json` - Status: Passed - [x] Library host tests, GPU hidden - `python -m pytest <file>` per file - Status: Passed (3 files fail identically on develop) - [x] hkp_pack producer guards - `python -m pytest descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py` - Status: Passed - [x] gfx950 GPU: new runtime-shape numeric test - Status: Passed (18/18) - [x] gfx950 GPU: existing attention suites - Status: Passed - [x] gfx950 GPU: baseline-vs-branch A/B, 132 problems, 3 rounds - Status: Passed (numeric gate) ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.com>
github-actions
Bot
added
documentation
project: hip-kernel-provider
rocke-lib: attention
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
AnaghaRaoAMD
approved these changes
Nov 4, 2026
Conflicts came from develop's attention work landing in the same code as the runtime head-count, window and batch changes (#14530, #14727): - #14616 2D LSE: lse_ptr is declared after the runtime kernargs in both engines, and _tiled_2d_signature takes emit_lse. The 2D and 3D LSE stores index the buffer by the runtime head count rather than a baked NUM_QH, in Python and the C++ twins. - #14367: VALID_ROWS guards combine with the runtime head bound, and the head_size 192 K/V path uses the runtime token stride. - Emit configs: develop keeps its indices; ours move to gfx950_attention_tiled_2d 68-70 and fastkv_regp 7-9. - Goldens: develop's representative shards with the attention, attention_d128_swa and attention_d256 shards re-blessed; the differential golden re-blessed for the gfx942/gfx950 tiled and strided-KV families. - The LSE dispatch test keeps develop's pinned configs under their attention ABI v2 tuning ids; every spec hash is develop's. Co-Authored-By: Claude <noreply@anthropic.com>
…e-aot-readiness Adapt the gfx950 unified attention ingestor engines from develop to the runtime shape kernargs: the 2D and 3D segment kernels take num_query_heads, num_kv_heads and sliding_window after qq_bias_stride_0, and the reduce kernel takes num_query_heads after seq_lens_ptr. The signature tests and fixtures follow. Co-Authored-By: Claude <noreply@anthropic.com>
…s with runtime kernargs The gfx950 unified 2D kernel now takes eleven scalars, so the like-slot exchange test sees C(11, 2) scalar pairs instead of C(8, 2). Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15024 — perf(rocke): runtime heads, window and batch for gfx942/gfx950 tiled attention (dlangbe, AICK-2676, AICK-2675) — 🟠 First-round: C++ crash paths, twin gaps
Utility: 🟢 Positive — one binary per GQA ratio serves every head pair, window and batch Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted First-round review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with develop, build and ctest, byte identity, library parity, both goldens, PR and installed tests, compile, exports and ABI consumers, saved probes, CI logs, code quality. Flags whether a closer look is needed; not an approval or rejection. |
- Head counts: one check_head_counts rule (Python _fmha_common, C++ attention_unified_selectors) runs first in every gfx942/gfx950 tiled validator; spec.num_queries_per_kv raises ValueError instead of dividing by zero. Both 2D C++ _new entries init the builder before naming, so the lower no longer frees an uninitialised builder on 0 KV heads. - gfx950 2D K/V loaders: page-id and copy-offset tables are sized from the spec on the arena instead of fixed 64-entry stack/ctx arrays. - fastKV register-P: the kwargs predicate, C++ predicate and C++ builder refuse non-gfx950 targets with the Python text. - Descriptors: TensorDescriptor gets _term/_term_i64 stride hooks so DynamicTensorDescriptor stops copying the four offset/move loops; C++ _naive_dynamic wraps _mixed, which returns the base view and sets the builder error on failure. New dynamic_helpers parity case (config 10). - Runtime-shape numerics: one torch-free, arch-parametrized module (DeviceArray + _launch_3d_pipeline) replaces the two torch-only files. - gfx950 2D hipDNN pack: drop the 511-sequence cap left over from the baked binary search, and its dead capacity helpers. Byte identity GREEN at llvm20/22/23; no golden moved. Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
|
Addressed the first-round review in ca119e0. The branch is also merged with develop (cb42d0e, clean; develop's new commits don't touch hip-kernel-provider).
Verification
Known, not in this PR: the gfx950 2D builder's other fixed tables ( |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical defects cause incorrect K/V addressing for 192-wide, one-warp heads and omit a required kernel argument from dense segmented launches.
5 open issues
What changed in this PR
Updates rocKE tiled attention on gfx942/gfx950 to share compiled kernels across compatible head counts, sliding-window sizes and batches, supporting AOT packaging.
Changes:
- Updates runtime kernargs, native launch signatures, search bounds and kernel identities.
- Extends dynamic-stride descriptors and gfx950 page-ID prefetch carries.
- Adds packaging validation and expands signature, parity, golden and CI coverage.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified3DSignature.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DSignature.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified3DTestGraphs.hpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DTestGraphs.hpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified3DNative.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DNative.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention_d256.json | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention_d128_swa.json | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/attention.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_loop_epilogue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_compute.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.helpers.attention.h | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/attention_ext.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_strided_kv_decode.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_running_max.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py | Updated as part of this pull request. |
Review effort: Max
…e launches - gfx950 2D page copies: require KV_HALVES_PER_CALL % head_size == 0 in both engines. A one-warp head_size 192 copy (512 halves, 2 2/3 rows) truncated to 2 tokens and addressed the wrong row and dim; those shapes keep the paged-KV descriptor. New cfg71 (the D192 one-warp spec dispatch picks) in both emitters and a D192 one-warp numeric pair. - Dense segment combine: launch the tiled reduce with its runtime kernargs (num_query_heads); it was launched with the pointer-only signature. - Binding parity mirror: add cfg68-71 so prove_parity_binding's completeness check matches the emitters. - Launch-signature test: check the "dg" token (split on "_" drops the underscore, so "_dg" never matched). - Differential golden: re-bless gfx942_attention_tiled_2d_stress (cfg21, 32/5 heads, is now refused by both engines under the head-count rule) and gfx950_attention_tiled_2d (cfg71), all flavors + canonical. Byte identity GREEN at llvm20/22/23; representative-IR and strided-KV goldens unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
|
Correction to my earlier comment: Unrelated to these commits, but seen while running |
amd-khushbu
approved these changes
Nov 4, 2026
AnaghaRaoAMD
enabled auto-merge (squash)
November 5, 2026 00:18
3 of 4 tasks
AviralGoelAMD
added a commit
that referenced
this pull request
Nov 5, 2026
…fied 3D engine - registry_enumerator: an enumerate.also section inherits selector_fields and cell_keyed_by; only a different list is refused (N9); test the documented spec_overrides refusal (N6). - Re-record the differential golden at llvm24 and canonical for gfx942_attention_tiled_3d and attention_strided_kv only (N3). - One split-KV scratch-layout helper in UnifiedAttentionGeometry.hpp for the gfx942 and gfx950 3D packs; gfx942 test graphs reuse the shared helpers (N8). - Carried #15024 code, converged with #15024: num_kv_heads guard in the gfx942 3D kernel_name, operator.index stride validation, the dg-token assertion, and the cache-key ctest split with its YAML categories (Copilot). - Rename leftovers after the #14242 de-dup, one home for the stream rationale, and one shared bundle-coverage helper for the four coverage tests.
Contributor
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15024 — perf(rocke): runtime heads, window and batch for gfx942/gfx950 tiled attention (dlangbe, AICK-2676, AICK-2675) — 🟠 Follow-up: conflict, CI budget, test gaps
Utility: 🟢 Positive — one binary per GQA ratio; every earlier item fixed Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: trial merge with today's develop and golden re-record, tests on that merge, build and ctest, byte identity, library parity, both goldens, installed tests, compile, ABI consumers, saved probes, CI logs, code quality. Flags whether a closer look is needed; not an approval or rejection. |
cgmillette
disabled auto-merge
November 5, 2026 03:39
Contributor
|
Automated merge-queue status (rocKE merge monitor) Hi @dlangbe, I turned auto-merge off on this PR for now. Before a rocKE PR lands, the rocKE merge monitor checks it on top of today's develop branch, and this version would not land cleanly:
The open items from our review, each with its fix:
What to do: address the open items in our latest review (#15024 (comment)). After that, turn auto-merge back on; the monitor re-checks every push. Nothing else changed: the monitor does not approve, reject or merge PRs. It only turned auto-merge off to keep develop green. This comment is updated in place. |
…e-aot-readiness # Conflicts: # dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention_d128_swa.json
…me kernargs The segment and reduce kernels gained runtime kernargs (num_query_heads, num_kv_heads and sliding_window on the segment, num_query_heads on the reduce), so the packed signatures are 22 and 6 arguments and the reduce's last is num_query_heads. TestGfx950AttentionUnified3DSignature already says so; the catalog census still expected 19 and 5. Co-Authored-By: Claude <noreply@anthropic.com>
…C++ head-count rule - The C++ head-count check moves into _fmha_common, like Python's check_head_counts: rocke_fmha_validate_common_spec calls it instead of repeating both texts, and the tiled entries include it from there. - The C++ fastKV register-P builder runs the head-count check before the shared scaffold divides by num_kv_heads, as the Python build raises. - New host CTest rocke_tiled_attention_entry_guards: the gfx942/gfx950 2D and 3D lowers and the fastKV register-P lower refuse 32/0 and 32/6 heads, and the fastKV lower and predicate refuse gfx942, with the Python text. The pytest twin needs the rocke_engine binding, which no CI lane builds. - Dense segment combine: combine_launch_abi() gives the reduce launch signature and runtime values; a torch-free test checks it against the reduce kernel's params, so the CPU lanes run it. Co-Authored-By: Claude <noreply@anthropic.com>
- 2D on gfx942/gfx950 sweeps num_seqs 1 and 2: a body that searches cu_q at runtime reads num_seqs only as the single-batch cohort. 3D (segment sizing reads the batch) and gfx1250 (baked search) keep 1 to 4. - The routing-env states pool on 2D only; on 3D every flipped state reproduced the default sweep's keys and kernels. - _norm_spec reads the spec fields directly instead of dataclasses.asdict. - validate_tiled_attention_arch is cached per arch string; every tiled support gate calls it and it depends only on the arch. Dropping each of the 44 2D key entries in turn, the trimmed sweep catches the same 12 (gfx942) and 18 (gfx950) as the full one. Co-Authored-By: Claude <noreply@anthropic.com>
dlangbe
enabled auto-merge (squash)
November 5, 2026 04:41
Contributor
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15024 — perf(rocke): runtime heads, window and batch for gfx942/gfx950 tiled attention (dlangbe, AICK-2676, AICK-2675) — 🟠 Follow-up: C++ builders skip head-count rule
Utility: 🟢 Positive — one binary per GQA ratio; every earlier item fixed this round Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with today's develop, build and ctest, byte identity, library parity, both goldens, tests at llvm22 and llvm24, installed tests, compile of every touched family, saved probes, engine-binding probes, ABI, CI, code quality. Flags whether a closer look is needed; not an approval or rejection. |
cgmillette
disabled auto-merge
November 5, 2026 06:19
jefyang1
added a commit
that referenced
this pull request
Nov 5, 2026
Conflict: the differential golden (llvm_gfx_all.json). Resolved as develop's file plus this branch's own changes since the last merge (93 entries). The 8 entries both sides changed, gfx950_attention_tiled_2d[27] and [31] at llvm20/22/23/24, are re-recorded from the merged tree: develop's #15024 re-blessed them for its tiled-attention change, and this branch's LDS-alignment rule (8-bit types count as 1 byte; align bounded by gcd(vec, innermost dim)) then lowers the align of four <8 x i8> LDS loads in each from 16 to 8. Nothing else in their IR differs from develop's. The binary ABI stays at 9 (develop is at 8). Golden check passes for every family (ll at four flavors, canonical); byte identity GREEN at llvm20/22/23/24. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
ThruptiRajLakshmanaGowda
added a commit
that referenced
this pull request
Nov 5, 2026
…ke/unified-2d3d-followups Conflict in gfx950/attention_tiled_3d.py supports_tiled_reduce_spec: keep this branch's delegation to supports_split_kv_reduce and port develop's (#15024) check_head_counts head-count rule into the shared predicate. Co-Authored-By: Claude <noreply@anthropic.com>
ThruptiRajLakshmanaGowda
added a commit
that referenced
this pull request
Nov 5, 2026
…cfg71 #15024's later review fix (ca119e0, merged as part of #15024) added its own cfg71 (bf16 D192 one warp) to gfx950_attention_tiled_2d_emit.{c,py} after our LSE phase was authored against an earlier snapshot of that branch. The C switch caught the collision (duplicate `case 71`, a build failure); the Python dict assignment did not -- `_CONFIGS[71] = {...}` silently overwrote cfg71's dict entry with no error. Move our three LSE head-major twins (idx64-66 + head-major) from 71-73 to 72-74, after cfg71. The parity bindings list (platform/cpp/bindings/attention_tiled_parity_configs.py) is positional, not index-keyed, so it already had the correct order (71=cfg71, 72-74=LSE) and needed no change; `check_byte_identity --only attention_tiled` is GREEN across all 6 families, gfx950_attention_tiled_2d now 75/75. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
|
Update: this is fixed in the merged form of this PR — |
ThruptiRajLakshmanaGowda
added a commit
that referenced
this pull request
Nov 5, 2026
…no-key-row O = 0, re-record attention_d128_swa and attention_d256 strided_decode: develop's _tiled_3d_signatures(arch, dtype, strided_kv=True) with this PR's attention_problem. gfx942 4-warp GQA bodies: develop's runtime num_query_heads (hq_p) in the O index, this PR's no-key-row O = 0 select. Golden shards re-blessed with the IR parity harness (llvm20/22/23/24). Co-Authored-By: Claude <noreply@anthropic.com>
4 tasks done
dlangbe
added a commit
that referenced
this pull request
Nov 5, 2026
…e-gfx942-i64-cfvst-loader Resolve conflicts with #15024 (runtime kv head count) and #14842 (buffer_rsrc operand checks). The i64 cfvst V page offset and its C++ twin now use the runtime token stride. Re-bless gfx942_attention_tiled_2d cfg18/19 in the differential golden. Co-Authored-By: Claude <noreply@anthropic.com>
AviralGoelAMD
added a commit
that referenced
this pull request
Nov 5, 2026
#15024 landed as 6ea2887 at the head this branch already merged (ffcfd98, in 982f771), so the four conflicts are that same content seen twice: - rocke_compile.py, test_hkp_pack_producer_guards.py, kernels/gfx950/attention_tiled_3d.py: ours (develop has no change to these files after ffcfd98; ours = #15024 + this PR's pieces). - differential golden llvm_gfx_all.json: develop's (it moved after #15024, #14843), then gfx942_attention_tiled_3d re-recorded at llvm20/22/23/24 and ir_canonical: +50 lines (configs 23-32), none removed. Every other file: merge result = 982f771 + develop's own changes after ffcfd98 (diff-checked file by file, 181 files).
AnaghaRaoAMD
added a commit
that referenced
this pull request
Nov 5, 2026
This branch gates every engine manifest's launch surface. After the develop merge, gfx950_attention_unified_2d's 'capacity' surface names binary_search_iters, which #15024 removed when it made the cu_q search trip count a runtime argument; develop's own gate checks only the gfx950 dense manifest. The manifest is skipped, with the reason, until its owners update the surface. Co-Authored-By: Claude <noreply@anthropic.com>
1 task
AviralGoelAMD
added a commit
that referenced
this pull request
Nov 6, 2026
…engine (#14259) ISSUE ID : AICK-2683 > **Merged with #14242** (gfx950 unified 2D/3D engines, on `develop` since 2026-11-04) in `58573717ad1`, and **de-duplicated against it** in `09549500acb` and `1fa151a1239`. The pack now uses #14242's SDPA graph helpers, problem struct and launch struct (`SdpaGraphMatch.hpp`, `UnifiedAttentionGeometry.hpp`; pack-local `Declined` stays). The Python module follows its reduce pairing (`reduce_spec_for` via `split_kv_reduce_spec`, one shared reduce-rule function used by both arches' `supports_tiled_reduce_spec`) and its names (`build_unified_attention_3d_tiled_module`, `supports_unified_attention_3d_tiled_module`). The gfx942 module test is `test_attention_tiled_3d_module_gfx942.py`. The shipped KDP embedded the old builder name 1,120 times, so it was edited in place (ids unchanged) and re-pushed to DVC in `f5a636cbb60`. > **On top of #15024** (merged to `develop` as `6ea2887e51e`). This branch merged #15024's final head `ffcfd98dd64` and then `develop` (`d638d9e1e40`); the carried copy `e151a6445dd` is gone: every file it touched is now #15024's code, and the gfx942 3D C++ instances, helpers and ABI string (v2) equal `develop`. What this PR still changes in rocKE: `resolved_num_queries_per_kv()` and the 3D module builder, predicates and `reduce_spec_for` (`kernels/gfx942/attention_tiled_3d.py`, `common/attention_unified.py`); the shared reduce rule `tiled_reduce_unsupported_reason` (`common/attention_split_kv_reduce.py`, used by both arches); strided parity configs 23-32 and their golden entries (+50 lines, none changed); the module test `test_attention_tiled_3d_module_gfx942.py`. ## What this adds hipDNN can now run rocKE's gfx942 split-KV **decode** attention (one query token per sequence) as the kernel-ingestor engine `hipkernel:Gfx942AttentionUnified3d`. Each SdpaFwd graph launches two kernels from one packed code object: rocKE's segment kernel, then its reduce kernel (#14158 `sources`). Scope (v1): dense K/V `[B,Hkv,Skv,D]`, `seqlen_q == 1`, D 64/128/256, fp16/bf16, GQA ratio 1/2/4/8/16 at **any head counts**, no mask or bottom-right causal. The catalog is **560 kernels** (30 cells), within the 1,000 AOT budget. The kernels read the head counts at run time and bake only the GQA ratio, so one kernel serves every head pair of its ratio; and every kernel loads K/V with the `stream` cache policy. ## What changed | Area / files | Crux of change | |---|---| | `rocke/library/kernels/{common/attention_unified.py, common/attention_split_kv_reduce.py, gfx942/attention_tiled_3d.py}` + parity configs | Runtime head counts in the 3D kernels are #15024's (on `develop`). This PR adds `build_unified_attention_3d_tiled_module` (segment primary + `reduce` companion), its spec-first predicates, `reduce_spec_for`, `resolved_num_queries_per_kv()`, one reduce rule shared with gfx950; strided parity configs 23-32. | | `IngestorGenerator/tools/registry_enumerator.py` (+ README, tests) | Opt-in `selector_fields` / `neutral_fields`; **`spec_overrides`** (pin a spec field the candidate picks, AICK-2894) and **`cell_keyed_by`** (key merged cells by an attribute of the pick, here the GQA ratio). Absent keys leave every output byte-identical. They apply per `enumerate.also` section (develop's, from #14242, which also supplies `launches`); sections inherit `selector_fields` and `cell_keyed_by` from `enumerate`; an explicit, different list is refused, since a shape's coverage is one key. | | `IngestorGenerator/configs/gfx942_attention_unified_3d.*`, `descriptors/rocKE/gfx942_attention_unified_3d/` | Manifest + selection: 29,808 request combos → **30 cells, 560 kernels**; head pairs are selector fields, head counts neutral, cells keyed by `num_queries_per_kv`, `kv_cache_policy: stream` pinned. Descriptors (KMD `num_queries_per_kv` replaces the head counts), DVC pointer, `launch_surface`, `ground_truth`. | | `kernel_ingestor_engine/packs/Gfx942AttentionUnified3d*` (+ tests, census, registration) | Native pack: admission = rocKE `attention_strided_decode.admits()` within v1 scope; matches a kernel on dtype, D and the graph's Hq/Hkv ratio; declines kernels whose i32 indices would overflow; scorer restates rocKE's split rule; workspace + always-run reduce; segment 30 arguments (graph Hq, Hkv, window 0 at 19–21), reduce 6 (Hq). Admission, arguments and split rule are checked against tables rocKE answers, generated into the build tree at build time (never committed; a rocKE change regenerates them). | | `integration_tests/kernel_ingestor_engine/IntegrationGpuGfx942AttentionUnified3dKnobs.cpp` (+ harness hook) | Device cases in `hip_kernel_provider_integration_tests`, registered under `hkp_gfx942_attention_unified_3d_available()`: forced `block_size`/`num_segments` must select the one expected shipped kernel, then O matches the CPU reference over two executes; BHSD/BSHD/padded-row/mixed layouts, poisoned workspace, a 288 MiB K/V case, and **Hq 24 / Hkv 6, a head pair no cell lists**. | | `integration-test-bundles/{quick,standard}/SdpaFwd/` (+ lane, TOML, YAML) | 18 decode cases (incl. a 16/16 head pair) claimed only by this engine on gfx942; GPU-reference lane. | - `num_seqs` and the head counts are declared IR-neutral, proven by rocKE's `test_num_seqs_is_ir_neutral_on_the_decode_grid` and `test_head_counts_are_ir_neutral_at_a_fixed_gqa_ratio`. - Block size and segment count are the performance arms; feature cells are D × dtype × GQA ratio. - `stream` only: rocKE's own rule would load small K/V (≤ 256 MiB) with `all`. A cold-memory gfx942 A/B of the two policies, 57 decode shapes × 2 rounds with buffers rotated past 1 GiB, found `all` never faster than the noise floor on any small shape (precommitted rule; AICK-2892), so the second copy of the catalog is not shipped. ## Why it works The engine admits what rocKE's own strided-decode candidate admits within v1 scope (checked row-by-row against a generated table; rocKE also admits GQA ratios 3–16 and D192, which v1 declares as exceptions, AICK-2893), launches the same two kernels with the same 30 + 6 arguments rocKE passes (checked against rocKE's `KvTensorLayout` and runtime kernarg values), and always runs the reduce, so the non-zeroed hipDNN workspace is safe: the segment kernel writes every slot, empty segments as `-inf/0/0`. ## How we validated Head is `5883c421570`: `d638d9e1e40` (the tree below plus #15024 and `develop` up to `90ef5ad9489`) merged with `develop` `bc1eeeea292`; its one conflict, the generated `SUPPORT_MATRIX_OVERVIEW.md`, was regenerated (`feature_matrix.py --check` passes). On `5883c421570`: pre-commit clean, gfx942 g1-g6 13 passed, support-matrix tests 184 passed, build exit 0, provider host tests the same 17 local failures (none new), census 10/10. Rows below are `d638d9e1e40`, host only (this host's GPU is gfx90a): | Gate | Tree | Result | |---|---|---| | Merge check | `d638d9e1e40` | 4 conflicts, all #15024's own text seen twice; every other file = prepared tree + `develop`'s later changes, checked file by file (181 files). | | IngestorGenerator | `d638d9e1e40` | gfx942 regeneration g1-g6 13 passed against the pulled KDP (catalog bytes unchanged); full suite 1737 passed, 37 skipped, 0 failed. | | Build + pack | `d638d9e1e40` | clang-tidy build exit 0; `hkp_packaging_product` exit 0; `hipdnn_validate_descriptors` on the packed gfx942 tree: success, engine loaded. | | Provider host tests | `d638d9e1e40` | Exit 1: 1590 passed, 73 skipped, 17 failed, 1 disabled. The 17 are the same names that fail on this gfx90a host at `4414e180d36` (1 `TestContainer`, 5 `TestKernelIngestorEngine`, 11 `TestGfx950AttentionDense*`), none new; CI's gfx94X shards passed them at `4414e180d36`. Census gfx942 10/10. | | IR / ABI | `d638d9e1e40` | Byte identity GREEN (`gfx942_attention_tiled_3d` 33 configs); representative IR parity OK llvm20-24; differential golden OK at llvm20/22/23/24 and canonical. | | rocKE / packaging tests | `d638d9e1e40` | 8 attention test files 583 passed, 24 skipped; producer guards + 4 bundle-coverage tests 30 passed, 2 skipped; pre-commit clean. | Not re-run on this head: the gfx942 device cases (CI's gfx94X lanes run them). Earlier rows, kept for history: Earlier head `435a0f80986` (formatting only on top of `f5a636cbb60`) (`17cedf98ba6` merged with `develop` `327b58255be` in `58573717ad1`, then the #14242 de-duplication). Rows tagged `17cedf98ba6` were taken before the merge and **not re-run** on the merge head. | Gate | Tree | Result | |---|---|---| | Correctness (gfx942 device) | `17cedf98ba6` | `IntegrationGpuGfx942AttentionUnified3dKnobs` **6/6** passed, 0 skipped, each selecting the expected ratio-keyed kernel (the Hq 24 / Hkv 6 case runs `…bf16_d128_g4_b16_t16_seg38`, the same kernel as the 32/8 cases); engine-pinned vs GPU reference: quick **6/6**, standard **12/12** decode cases (incl. 16/16 heads), every execute selected and served by this engine, 0 dropped; ctest GPU-reference lane passed; full provider tests 1525 passed / 9 skipped / 0 failed. | | Artifacts | `17cedf98ba6` | `verify_variant_sets --mode full`: **560 descriptors, 560 distinct binaries**, compiled specialization agrees with metadata (host build tree and the gfx942 device build's installed shard); census 5/5 (build tree), installed census 9/9; `coverage_gate` rungs 1–2 PASS. | | Host gates (build) | `17cedf98ba6` | rocKE-on clang-tidy build clean (`hkp_pack` 560 variants); `*Gfx942AttentionUnified3d*` host tests 69 passed; the knob cases skip with the arch reason on a non-gfx942 device; packaging tests (bundle coverage, both-launch pack) pass; pre-commit clean. | | Host gates (source) | `17cedf98ba6` | Registry regeneration g1–g4 pass; `launch_surface --check` PASS; IngestorGenerator suite 1637 passed / 31 skipped; rocKE attention tests: no failure beyond the PR-head baseline (two pre-existing: a dense budget-name drift and host tests that need torch). | | IR / ABI | `17cedf98ba6` | Byte identity GREEN (C++ ≡ Python `.ll`); representative IR parity: only the six gfx942 3D/reduce cases changed, re-blessed; differential golden: only `attention_strided_kv` and `gfx942_attention_tiled_3d` changed, re-blessed, 0 unblessed changes at llvm20/22/23. | Merge head, host only (CPU host without a gfx942 GPU): | Gate | Tree | Result | |---|---|---| | IngestorGenerator suite | `58573717ad1` | 1733 passed, 37 skipped, 0 failed. gfx942 regeneration g1–g6 ran against the pulled KDP and passed; the gfx950 KDPs are not pulled here, so their gates skip. | | Provider host tests | `58573717ad1` | clang-tidy build clean. `*Gfx942AttentionUnified3d*` + `*Gfx950AttentionUnified*` 213/213 passed. Whole binary: 1593 passed, 73 skipped, 17 failed. The 17 are the set recorded on an earlier `develop` baseline (1 `TestContainer`, 5 `TestKernelIngestorEngine`, 11 `TestGfx950AttentionDense*`); I did not re-run `develop` for this merge. | | IR / ABI | `58573717ad1` | Byte identity GREEN, 91 families (+8 gfx1250 at llvm23), including `gfx942_attention_tiled_3d` (33 configs). Representative IR parity OK at llvm20/22/23/24. Differential golden: OK at llvm20/22/23 (no blessed config changed or missing), **FAILS at llvm24 and in canonical mode** (`gfx942_attention_tiled_3d` and strided-KV entries); the golden has not been re-recorded there yet. | | Docs / hooks | `58573717ad1` | `pre-commit` over `origin/develop..HEAD` passed; support-matrix overview current. | After the de-duplication (`f5a636cbb60`, host only): the C++ build packed all 560 variants with the real compiler under the new builder name (build exit 0); packed KDP has 560 kernels, 560 distinct segment and 560 distinct reduce binaries, ids equal the shipped ids; gfx942/gfx950 unified suites 213/213; full provider binary the same 17 failures as above; IngestorGenerator regeneration g1–g6 for gfx942 and the enumerator tests pass (g2/g3 run, not skipped); rocKE module tests 401 passed, 5 skipped; producer-guards + bundle coverage 28 passed; byte identity exit 0; representative IR parity OK at llvm20/22/23/24. Not compared: binary digests before vs after the rename (the pre-rename pack is gone); IR identity of the module and reduce kernels was checked on 48 specs instead. One side effect: gfx950's reduce predicate now reports the LDS reason before the head-grouping reason when a spec fails both (accept/reject unchanged). Not re-run on the merge head: the gfx942 device cases and engine-pinned bundle runs, `verify_variant_sets`, `coverage_gate`, `launch_surface --check`, the census target, and the real-compiler packing tests (9 fail on this host for an environment reason I have not isolated; the earlier host-gates run attributed the same tests to the nested-Ninja setup). The catalog bytes are unchanged by the merge (regeneration g2 reproduces the shipped KDP byte for byte). CI is running on this head; its gfx94X lanes had not reported on `17cedf98ba6`, and I have not seen them on this head. ## Performance No speed claim. The scorer restates rocKE's split rule; a learned ranker over this catalog is AICK-2545 (UHD). Reading the head counts as kernargs instead of constants is #15024's change; its PR reports the rocKE-level effect. Default engine order picks ASM_SDPA for bf16 D128 decode (also numerically correct). ## Notes - The gfx942 KDP is on the DVC `ingestor` remote (verified by a fresh pull into an empty cache), so the superbuild lanes can fetch it. - v1 refusals (not implemented, tracked follow-ups): GQA ratios other than 1/2/4/8/16 and head size 192 (AICK-2893), paged K/V + seq-len tensors, sliding window, stats/LSE, attn-mask tensor, dropout, sinks, FP8, ALiBi. With Sq = 1, bottom-right causal hides nothing, so no case exercises masking. - Follow-ups: AICK-2545 (UHD ranker; trains on this 560-kernel catalog). --------- Co-authored-by: David Langbehn <47395968+dlangbe@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: dlangbe <David.Langbehn@amd.com> Co-authored-by: AnaghaRaoAMD <anarao@amd.com> Co-authored-by: ThruptiRajLakshmanaGowda <Thruptiraj.LakshmanaGowda@amd.com>
cgmillette
pushed a commit
that referenced
this pull request
Nov 6, 2026
JIRA ID : AICK-2606 Partially addresses #14651 (the `TestKvAddrLimit` classification failure on develop). ## Motivation #12446 added a launch-time guard that rejects a paged KV cache larger than the compiled kernel can address. Some gfx942 K/V loaders form an i32 byte offset from the cache base even when the spec is built with `use_i64_kv_addr=True`. The guard caps those specs through `_I32_ONLY_KV_LOADERS`. The conflict-free V store (cfvst) loader is the capped loader that matters in production. It is the default V loader for gfx942 bf16 D64 and fp16 D64/D128 prefill (no sliding window). A sweep of the gfx942 default-route specs one block past 2^31 bytes found that cfvst is the only capped loader any shipped i64 spec selects. Its V offsets wrap past 2^31 bytes. Before #12446 that read wrong data or faulted the GPU; since #12446 the guard raises `ValueError` before launch. Dispatch has no fallback route, so a deployment whose per-layer K cache is over 2^31 bytes fails every prefill on these configurations. This PR gives the cfvst loader an i64 path that costs no more than the i32 path, and removes the cap. That lifts the cap for every shipped gfx942 configuration. ## Technical Details **cfvst V loader** (`library/kernels/gfx942/attention_tiled_2d.py`, `_cfvst_load_v_regs`) - Each cfvst load call covers a fixed range of tokens: 16 at D64 and 8 at D128 with 4 warps. On every geometry the launcher and the tuning space build, that range lies inside one page, and which page is known at build time. - Under `I64_KV_ADDR` the loader loads the tile's page ids once, builds one 64-bit page pointer per page, and indexes each V load with an i32 offset within the page. The per-lane select and divide stay i32, as on the i32 path. The pointers come from `_tile_page_ptrs`, split out of #14349's `_tile_pages` with the same op order. - The in-page byte offset, the kv head's byte offset and the `PAGE_RSRC_I64` rule come from `_kv_page_off`, `_kv_head_page_b` and the shared predicate, which the per-page K/V loaders use too, so a change to the paged K/V layout reaches cfvst as well. They use the runtime token stride from #15024. - A geometry whose load calls would span pages raises `ValueError` at build time. This follows the existing rule for K/V copies that straddle pages. No such geometry can launch: the launcher accepts block sizes 16, 32 and 64, and the tuning space offers 1, 2 or 4 warps. - It covers every load form: the vector `global_load_vN` path and the `HIPDNN_GFX942_CFV_STORE_SEPOFF` / `_SCALAR_LOAD` diagnostic paths. - The i32 cfvst loader emits the same IR as before. - An earlier commit in this PR used full per-lane i64 offsets (`offset_i64`). That was correct, but at fp16 D128, which already sits at the VGPR limit, it added register spills and made the i64 kernel noticeably slower than the i32 one. The page-pointer version replaces it. **Running-max carry on the transposed softmax path** (`attention_tiled_2d.py`, the loop yield) - After merging develop, the fp16 D128 i64 cfvst kernel was 1.3x slower than its i32 form. A bisect over the merged develop commits found #14542, which carries the raw running max. On the transposed path the raw max then stayed live through the whole body beside the clamped exponent reference. This kernel already sits at the VGPR ceiling, so the one extra value made the compiler route the per-tile compares through VCC one at a time. - The carry is now rebuilt at the yield as `new_l == 0 ? -inf : m_new`. The raw max is -inf exactly while the row has seen no key, which is when l is 0 (a seen key adds P = 1 at the max), so the carried value is the same and #14542's behaviour is kept. Testing l rather than the max compare matters: LLVM folds a rebuild from the compare back into the raw max. - This changes the IR of every gfx942 tiled 2D kernel on the transposed path, i32 included. The non-transposed paths and gfx950 are unchanged. **Scalar `global_load` lowering** (`core/lower_llvm.py`, `cpp/core/lower_llvm/mem.cpp`) - `global_load_typed` always wrote its GEP index as `i32`, so an i64 element index gave invalid IR. It now uses the index operand's own type, as `global_load_vN` already does. This was found while writing the first i64 loader. The final cfvst loader no longer needs it, but the lowering was wrong for any i64 index. Every existing caller passes an i32 index, so no existing kernel changes. **C++ twin** - The page-pointer cfvst loader and the same build error are mirrored in `attention_tiled_2d_conflict_free_v.cpp`. `rocke_gfx942_attn2d_tile_page_ptrs`, `rocke_gfx942_attn2d_kv_page_off`, `rocke_gfx942_attn2d_kv_head_page_b` and `rocke_gfx942_attn2d_page_rsrc_i64` are shared with the per-page K/V loaders. - The rebuilt running-max carry is mirrored in `attention_tiled_2d_kv_body_qk_softmax.cpp` and `attention_tiled_2d_kv_body_pv.cpp`, only when that body builds for gfx942. gfx950 fastkv_regp shares the C++ body but its Python (`gfx950/attention_tiled_2d.py`) still carries the raw max, so it keeps that carry. - No emit config covered cfvst before, and the twin did not match Python. It is now byte-identical: - It computes the per-tile values once and reuses the cached `max_seq_prefix_len`. - It places the split load and store where the transposed softmax issues them. - It fixes C argument evaluation order (unspecified, right-to-left in practice) in `v_t_slot`, the transposed PV token base, and the cfvst block guard and prezero loop. **Lift the cap** - `use_conflict_free_v_store` is removed from `_I32_ONLY_KV_LOADERS` and classified as i64-safe in `test_every_tiled_flag_is_classified_for_kv_addressing`. - `test_gfx942_cfvst_meta_has_no_limit`, `test_limit_memo_is_keyed_per_kernel` and `TestGfx942RouteGuard.test_cfvst_route_launches_past_2gib` now expect no limit. The route test launches at 65535, 65536, 65537 and 69632 blocks and checks the spec is i64 cfvst. - `use_additive_bias` (#13120) is also classified (`_NO_KV_ADDRESS_FLAGS`). The classification test already failed on develop without it. **Byte-identity coverage** - `library/tests/parity/gfx942_attention_tiled_2d_emit.{py,c}`: cfg18 is i32 cfvst and cfg19 is i64 cfvst, both on the bf16 D64 GQA 32/8 prefill geometry. The shipped spec also sets the sliced-K ring, direct-Q-global and the mask-limit stack. The C twin does not port those yet, so the configs leave them off. They still cover the new V loader code. - `llvm_gfx_all.json`: cfg18/19 blessed at llvm20, llvm22, llvm23, llvm24 and in canonical IR. The running-max change re-blesses the 18 transposed-path configs of `gfx942_attention_tiled_2d` and `gfx942_attention_tiled_2d_stress`; no other family moves. - Representative-IR golden: only `attention/gfx942/2d_lse_transposed32_bf16_d128` moves, from the running-max change. **Tests and docs** - `platform/tests/core/test_global_load_index_width.{py,cpp}` check the GEP index width; the C++ one is a host CTest. - `library/tests/test_gfx942_cfvst_i64_kv_addr.py` (CPU): - every V load of the shipped i64 cfvst spec indexes a page pointer with an i32 offset, including both diagnostic forms; - the i32 spec keeps i32 offsets from the cache base; - a page-spanning geometry is refused; - every cfvst geometry in the gfx942 tuning space builds with i64. - `library/tests/test_gfx942_cfvst_i64_kv_addr_numeric.py` (gfx942 GPU lane): a CPU routing anchor per shipped shape, plus an on-GPU check for bf16 D64, fp16 D64 and fp16 D128. Each runs `run_unified_attention_torch` on a 2.125 GiB GQA 32/8 causal cache with every block-table entry above 2^31 bytes, against an fp32 reference. It gates on `get_device_arch() == "gfx942"`. - `library/tests/test_attention_unified_masked_tile_ir.py` accepts the rebuilt carry in its exact form and still rejects a clamped carry. - CI lanes: `test_gfx942_cfvst_i64_kv_addr.py` runs in `rocke_library_attention_cpu_pytest`, and the numeric file in the library GPU lane. `rocke_global_load_index_width` is registered in `tests/cmake/CoreTests.cmake` and listed in the quick and ffm-quick tiers of `ROCKE_ENGINE_test_categories_external.yaml`. - `builders/gfx950/attention/README.md`: describes the cfvst page-pointer addressing; the paged-cache K/V limit list no longer names cfvst. **Out of scope (follow-ups)** - The other `_I32_ONLY_KV_LOADERS` entries: cfv, vhbm and khbm on gfx942, and fp8-PV on gfx942 and gfx950. No shipped spec selects them under i64. - Under i64 the sliced-K ring K loader still builds a buffer resource per copy, the pattern #14349 replaced in the other K/V loaders. - `global_store_typed` still writes its GEP index as `i32` in both engines, so an i64 store index would give invalid IR. No caller passes one. - Running `test_selector_parity.py` in a CI lane (the rest of #14651). - Replacing the hardcoded `0x7FFF0000` K/V `num_records` with one named limit (changes IR and goldens). - Exposing the cache-size limit to `supports_native_unified_attention*`. - Porting the sliced-K ring, direct-Q-global and the mask-limit stack to the gfx942 C twin, so an emit config can match a full shipped spec. ## Test Plan CPU: - Byte-identity gate at every flavor: `ROCKE_LLVM_FLAVOR=llvm20|llvm22|llvm23|llvm24 python tools/check_byte_identity.py` - Differential golden, as CI runs it: `run_diff.py --mode ll --python-only --check-golden --flavor llvm20,llvm22,llvm23,llvm24` `run_diff.py --mode ir --canonical --python-only --check-golden` - Representative-IR golden: `rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256 --flavor <flavor>` at each flavor. - `test_selector_parity.py`, `test_gfx942_cfvst_i64_kv_addr.py`, `test_attention_unified_masked_tile_ir.py`, `test_attention_gfx942_kv_pages.py`, `test_global_load_index_width.py` and the attention spec/selector suites, with the GPU hidden. - `rocke_global_load_index_width` CTest. gfx942 GPU (torch 2.13.0+rocm10.0.0 wheels, Python lowering): - `test_gfx942_cfvst_i64_kv_addr_numeric.py`. - #14542's `test_attention_unified_2d_masked_tile_numeric.py` and `test_attention_unified_scalar_masked_tile_numeric.py` (fully masked first tile, NaN rows), plus `test_attention_tiled_runtime_shape_numeric.py` and `test_gfx942_4wgqa_bounds_numeric.py`. - `builders/gfx942/attention/prefill/parity_unified_attention.py --scenario default`. - Kernel resources (VGPR and SGPR counts, spills) from the code object notes of the shipped cfvst kernels, i32 and i64. - Bisect of the fp16 D128 i64 slowdown: each develop commit merged into the pre-merge head, compiled, and its SGPR count read; the boundary then timed. - Performance on the shipped prefill shapes, i32 vs i64 on this branch. The workload is 2 sequences of 4096 tokens, causal, GQA 32/8, 32-token pages. - Both arms run the same small cache. The i64 arm only claims a larger block count, which selects the i64 spec. - Each arm is first checked against an fp32 reference (declared gate max_abs ≤ 4e-2). - Time is the median of 50 launches after 10 warm-up launches, measured with HIP events, over two repeats. ## Test Result - Byte-identity (`check_byte_identity.py`, every family): GREEN at llvm20, llvm22, llvm23 and llvm24, including cfg18/19. - Differential golden: GOLDEN OK at every flavor and in canonical IR. Representative-IR golden: OK at all four flavors. - CPU tests pass. The new CPU tests fail on develop and pass here. - gfx942 numeric, at this head: all three shipped shapes pass on caches past 2^31 bytes (bf16 gate 4e-2, fp16 gate 2e-2). #14542's masked-tile tests pass, including the NaN rows on the transposed path. - gfx942 parity (default group): 22 PASS, 6 SKIP, 0 FAIL. The skips are D256 shapes over the LDS budget. - Kernel resources: the fp16 D128 i64 kernel uses 252 VGPRs and 101 SGPRs with no spills; before the running-max change it used 254 and 49. - Performance, at these measured configurations, as time ratios. Every arm passed the numeric gate (max_abs 7.9e-3 bf16, 1.1e-3 fp16). | shape | branch i64 / branch i32 | fp16 D128 i64, running-max change / without it | |---|---|---| | bf16 D64 | 1.01–1.02 | | | fp16 D64 | 0.99–1.00 | | | fp16 D128 | 0.95 | 0.72 | The running-max change leaves the i32 kernels' times unchanged within run-to-run noise. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.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 ID : AICK-2676, AICK-2675
Motivation
hipDNN AOT packaging ships one compiled binary per kernel symbol, and a packaged catalog cannot enumerate head counts, sliding-window sizes or batch sizes. The gfx942 and gfx950 tiled attention kernels baked all three into their IR, so every head pair, window and batch bucket compiled its own binary.
This lands #14530 (gfx942) and #14727 (gfx950), both reviewed and merged into
users/dlangbe/rocke-aot-readiness. The tiled kernels (2D, 4-warp GQA, fastKV register-P 2D, 3D split-KV and reduce) now read the head counts and the window as kernargs and derive thecu_qsearch trip count fromnum_seqs. One binary serves every head pair at a given GQA ratio, every window size and every batch.Technical Details
RUNTIME_KERNARG_FIELDS;_tiled_2d_signature/_tiled_3d_signaturesappend them and_runtime_kernarg_valuespacks them. Every gfx942/gfx950 launch site uses these arch-aware signatures.num_seqsat runtime (binary_search_seq_idx_trip_count); the keys drop the batch bucket.window > 0stays baked (name tokensw).DynamicTensorDescriptor, which now takes mixed int/Value and i64 strides in both engines. Names carryg<Hq/Hkv>and keys the GQA ratio. The fast paged-KV descriptor path still bakes 64/8, so its names carryh64kv8.soffset), a per-copy token offset and a lane part, built once before the loop. The KV loop carries the page ids of the current and next tile; tiles of more than 16 pages read them in place. This recovers the runtime-stride cost on the D256 fast route and long multi-batch d128 prefill.i64for i64 K/V addressing._shape_key(2D grid, 3D grids and workspace) carries both head counts.supports_fastkv_register_p_2d_specletshkp_packpackage the gfx950 fastKV register-P builder, for gfx950 only.test_attention_tiled_launch_signature.pyjoins the attention CPU list. The cache-key tests run as a quick entry plus a standard-tier invariant sweep, both listed inROCKE_ENGINE_test_categories_external.yaml. The runtime-shape numeric module,test_attention_tiled_runtime_shape_numeric.py, joins the library GPU list. It is torch-free, so CI's GPU lanes run it.ABI notes
num_query_heads,num_kv_headsandsliding_window(2D/3D) ornum_query_heads(reduce). In this repo those arerun_unified_attention_torch, the dispatch tuning candidates, the harnesses updated here, and two hipDNN packs that already append them:hipkernel:Gfx950AttentionUnified2D(Gfx950AttentionUnified2DNative.cpp, which no longer caps batch at 511 sequences) andhipkernel:Gfx950AttentionUnified3D(Gfx950AttentionUnified3DNative.cpp). No hipDNN pack launches the gfx942 tiled kernels.ATTENTION_ABI_VERSIONand the strided-KVabi_versionmove to v2, so stored gfx942/gfx950 attention tuning records and compile keys from before this change no longer resolve.ROCKE_ABI_VERSIONis unchanged: the removedbinary_search_itersC accessors are kept exported.Risk Assessment
Risk 4: every gfx942/gfx950 tiled kernel changes its argument list and IR, the cache keys change, and the 2D K/V loaders change in both engines. Byte identity is GREEN at every flavor, gfx1250 IR is unchanged, and both arches were run on hardware.
Test Plan
ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py.--check, the strided-KV library golden, and the differential golden (run_diff.py --check-golden, ll and canonical).test_hkp_pack_producer_guards.py.test_attention_tiled_runtime_shape_numeric.py, one arch-parametrized module that launches from numpyDeviceArrays against an fp32 numpy reference. Shapes that share a key run back to back; the second must compile nothing and both must match an fp32 paged reference (max_abs ≤ 4e-2) with unused block-table slots on a poison block. The existing attention GPU suites for each arch were also run.run_unified_attention_torchwith production routing (backend="auto"), per arch.Test Result
test_strided_kv_decode.py,test_strided_kv_lds.py, onetest_selector_parity.pycase) fail identically on develop.test_attention_tiled_runtime_shape_numeric.pyatca119e0be16. gfx950 was 18/18 on the per-arch file this replaces; the re-run on the merged module is pending. The existing suites match develop.f37a7056c791c620a3ee6addregistry_enumerator.py --plan-space, one pinned tuning candidate, 24 cells): gfx942 162 distinct kernels vs 324, gfx950 211 vs 421, because head pairs at one GQA ratio share a kernel.Submission Checklist