perf(rocke/attention): GPU-time optimizations for unified and dense forward attention on gfx942/gfx950 by AnaghaRaoAMD · Pull Request #14367 · ROCm/rocm-libraries

AnaghaRaoAMD · 2026-10-25T17:21:36Z

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.

…herry-pick of #14276) Squash of the PR's net change (pr/14276 4287837 vs its merge-base with develop). Applied cleanly onto develop; no conflict resolution needed. Co-Authored-By: Claude <noreply@anthropic.com>

The tiled 3D segment kernel rejected num_queries_per_kv values that do not divide BLOCK_M=16 (e.g. 3, 5, 6), so decode with such a ratio fell back to the 2D kernel, whose grid is only num_seqs * num_kv_heads workgroups; long KV decode then ran on a handful of CUs. The 3D kernel already handles a padded q-block: it bounds every q row by the sequence's query length, and a padding row past BLOCK_Q tokens either falls outside that bound or recomputes a row the next q-block owns. Drop the divisibility check in both engines and add an NQK=5 config to the gfx942/gfx950 tiled-3D byte-identity emitters. Co-Authored-By: Claude <noreply@anthropic.com>

…#14216) Squash of the PR's net change (pr/14216 6120fae, 2 commits, vs its merge-base with develop). Applied cleanly; the only file also touched since the merge-base (kernels/gfx942/attention_tiled_3d.py, by #14276) auto-merged in disjoint hunks. Co-Authored-By: Claude <noreply@anthropic.com>

…h (cherry-pick of #14349) Squash of the PR's net change (pr/14349 9c32c35, 2 commits, vs its merge-base with develop). Conflict: platform/tests/golden/rocke_representative_ir_sha256.json, llvm23 summary "ok_cases" only (develop had moved 268 -> 270, the PR 268 -> 269 for its one new case). Resolved to 271 (develop's two cases plus the PR's one). Verified by re-running the IR parity harness with --write into a scratch file: the regenerated golden is identical to the resolved one, so no existing hash moved and only the PR's new attention/gfx942/2d_bf16_d128_b16_decode_i64 entries were added. Co-Authored-By: Claude <noreply@anthropic.com>

…ry-pick of #12714) Squash of the PR's net change under dnn-providers/hip-kernel-provider/rocke (pr/12714 eb97120 vs its merge-base with develop). Left out: the PR's six hip-kernel-provider / hipDNN IngestorGenerator files outside the rocke tree (gfx950 dense catalog pin to the qb_minor grid, Gfx950AttentionDenseGeometry, its test, the gfx950_attention_dense generator configs and artifacts.dvc). This base branch carries rocke changes only. Conflict: platform/tests/golden/rocke_representative_ir_sha256.json, summary counters only (arches gfx942/gfx950, families attention_dense/conv, ok_cases) in all three flavors, where develop and the PR each added cases. Resolved as develop's value plus the PR's delta (e.g. llvm23 ok_cases 271 + 7 = 278). Verified by re-running the IR parity harness with --write into a scratch file: the regenerated golden is identical to the resolved one. Co-Authored-By: Claude <noreply@anthropic.com>

With num_queries_per_kv not dividing BLOCK_M=16 a q-block holds BLOCK_Q * NQK valid rows plus padding. A padding row maps to the first token of the next q-block but is computed under this block's KV bounds, so its segment store could overwrite the next block's correct partials in multi-block prefill. Mask every segment-workspace store to rows below VALID_ROWS = BLOCK_Q * NQK, as the 2D kernel does. The guard is emitted only when VALID_ROWS < BLOCK_M, so IR for dividing ratios is unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

… into users/anarao/perf/gpu-only-p1-decode

…se body Without a sliding window, the persistent dense body now matches the grid body's softmax form: the row max is taken on unscaled scores and each probability is exp2(fma(s, qk_scale, -m)), one VALU op per score instead of a multiply and a subtract. The row sum is a pairwise <2 x f32> tree, the half-wave max/sum exchange uses v_permlane32_swap instead of an LDS ds_bpermute, and the wide-DMA LDS line index is kept in an SGPR so each tile's LDS destination is scalar math. Main-loop VALU per KV tile drops from about 160 to about 100 on the D128 wide-DMA configurations. Window phases are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

…se on gfx950 dense When the causal sliding window spans at least one query block (W // block_n >= block_m // block_n) and the diagonal is unshifted, every window-edge (L) tile ends below the block's first query and every causal-edge (R) tile starts at or past it. The causal compare is then always true in L and the window compare is always true in R, so each boundary phase applies only the bound that can bind. Both bounds stay on when the window is narrower than the query block, and on the prologue tile (the first block's start tile is its diagonal tile). Applies to the default and persistent builders. Re-blessed attention_dense/gfx950/paged_swa_fp16_sq512 (W256 at block_m 256 takes the new path) and added two golden cases (default grid W256, persistent block_m 128 W128 with sinks). The bf16 sliding-window numeric test now also covers W == block_m and W > block_m (non-tile-multiple). Co-Authored-By: Claude <noreply@anthropic.com>

…e 32x32 combo Single-batch (num_seqs == 1) FP8 K/V prefill on gfx950 was excluded from the 4-warp transposed 32x32 combo and fell back to a 16-row-per-CTA geometry, re-loading and re-dequantising every K/V tile for each small query block at one wave per CU. The sync-dequant FP8 loader already stages K/V in LDS at the working dtype, which is what the combo reads, so admit FP8 (without sliding window) into the single-batch combo cohort. The schedule levers the FP8 loader does not wire stay off for it: early-V and V double-buffer are now gated on no FP8 (K single-buffer, softmax interleave and the d128 small tile already were). The C++ selectors mirror the change. Adds two gfx950 tiled-2D emitter configs for the FP8 single-batch combo (d128 T128, d64 T64) and FP8 single-batch rows to the selector parity matrix. Co-Authored-By: Claude <noreply@anthropic.com>

The unified 2D and 3D tiled kernels on gfx942 and gfx950 handle any head size that is a multiple of 32, but the admission gates listed {64, 128, 256} explicitly, so head_size 192 problems had no unified candidate at all. Add 192 to UNIFIED_HEAD_SIZES and to the four per-arch supports_tiled_{2d,3d} gates, mirrored in the C++ entry and public glue checks. Adds head_size 192 emitter configs (the selector's prefill 2D spec and decode 3D spec) to the gfx942 and gfx950 tiled 2D/3D parity emitters. Co-Authored-By: Claude <noreply@anthropic.com>

…ense prefill - gfx942 dense: shift the causal diagonal by seqlen_kv - seqlen_q, read from the shape params, in the mask, the KV-loop bound and the diagonal split, so one binary serves every chunk length. Dispatch admits causal_bottom_right; a bottom-right window and cross-length NO_MASK requests still decline. - gfx942 and gfx950 dense default grids: optional kv_segments (new tuning axis, default 1) splits each work item's KV tiles across several CTAs. Each segment writes its FP32 partial output plus running max and denominator, and the existing split-KV combine kernel merges them. This raises the CTA count for short query chunks over long KV. - Tests: GPU numeric cohorts for both paths, contract and feature-set updates, and new golden cases (additions only; existing hashes unchanged). Co-Authored-By: Claude <noreply@anthropic.com>

…verage Co-Authored-By: Claude <noreply@anthropic.com>

… prefill An unpinned attention request that declares contiguous tensors (layout='dense') with more than one query row now dispatches to the arch's dense kernel at its default spec instead of failing with no candidate. On gfx950 the body is chosen by a work-item rule: the persistent body once ceil(Sq/block_m)*Hq*B covers the persistent CTA count and the causal diagonal is fixed, wide LDS DMA additionally for plain causal D128 problems on tile multiples, the grid body otherwise (including the moving bottom-right diagonal). gfx942 has a single dense candidate whose default spec already picks its body. Paged and unspecified layouts and single-row requests keep the unified route, so existing paged callers see no change. The dense candidates now declare the layout_dense feature, matching the RDNA convention. Dispatch is Python-only; no emission changes. Co-Authored-By: Claude <noreply@anthropic.com>

…or gfx942 dense Two window-only changes to the gfx942 dense prefill path: - Dispatch: a causal band of width W visits about block_m + W keys per query row, so the 256-row default block wastes most of its tiles on masked work for narrow windows and leaves short sequences below one CTA per CU. Window requests now get their own geometry rule: W <= 256 uses 128-row query blocks (64 rows for bf16 D64) on the one-CTA-per-block grid; a wider window whose 256-row grid has fewer query blocks than CUs uses 128-row blocks with the grid choice unchanged. Requests without a window keep the default geometry. - Kernel: when the window spans the query block (W // block_n >= block_m // block_n) the KV loop splits into three phases: window-edge tiles apply only the window bound, interior tiles run mask-free, and diagonal tiles apply only the causal bound. Narrower windows keep the single loop with both bounds. Re-blessed gfx942 goldens dispatch_swa_d128_bf16_w128 (now 128-row blocks and the phase split) and swa_d128_fp16_w256 (W == block_m takes the split); added persist_swa_d128_bf16_w256 and dispatch_swa_sink_d64_bf16_w128. The SWA numeric cohort gains W512 rows that exercise a non-empty mask-free phase. Co-Authored-By: Claude <noreply@anthropic.com>

…50 dense When a bottom-right causal chunk on the gfx950 default grid launches fewer work items than CUs, dispatch now takes the 128-row query tile and splits each work item's KV range into the power of two nearest twice the CU-to-work-item ratio (capped at 8). A grid that already covers the device keeps one segment, and so do requests with sinks or LSE output. kv_segments is recorded in the knob space and scoped to the plain default grid. Co-Authored-By: Claude <noreply@anthropic.com>

bf16 D128 took the naive V path (four element-wise ds_read_u16 per PV operand, 128 LDS reads per KV tile) because the transposed-V store was recorded as spilling past the waves-per-eu=2 cap. With P converted by a NaN-free round-to-nearest-even pair pack (bits + 0x7FFF + lsb, one perm_b32 per element pair) the bf16 body fits at about 220 VGPR with no spill on LLVM 23, so _use_cfvst now enables the transposed-V store for D128 in both dtypes. The pack is bit-identical to the generic f32->bf16 conversion for finite inputs, and P = exp2(s - m) is always finite. Per KV tile the bf16 loop drops from about 930 to about 590 instructions (LDS reads 144 -> 32, lgkmcnt waits 82 -> 28). fp16 IR is unchanged. Golden re-blessed for the 15 bf16 D128 cases only; fp16 and D64 hashes are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

The gfx942 dense default spec always used the shared 256-row query tile, so short causal prefills launched fewer work items than CUs (Sq=512 with 32 heads is 64 items on a 304-CU part) and each item walked a long KV loop. When the 256-row work count (query blocks x heads x batch) is below the CU count, the default now uses a 128-row tile, which doubles the work items and halves the longest causal KV loop. At or above the CU count the 256-row tile stays. Only D128 without sliding window or sinks takes the new rule; block_m is recorded in the tuning knobs because it now depends on the problem. Co-Authored-By: Claude <noreply@anthropic.com>

The gfx942 tiled 2D decode body fits one single-wave workgroup per CU (its LDS footprint), so a large-batch decode grid leaves most SIMDs idle and a small-batch short-KV grid of a few workgroups walks its tiles serially. The shared 2D/3D selector picks 2D for both (2D grid above the subscription target, or KV <= 512 tokens). UnifiedAttentionProblem.select_path() takes an optional arch. On gfx942, all-decode problems in the decode split rule's cohort (fp16/bf16, head size 64/128, 16-token pages, no sliding window, no explicit target_ctas) route to the 3D segment kernel whenever it supports the problem and its K/V loader can address the cache. The runtime and the dispatcher candidates pass the arch; without one, routing is unchanged. Selection only: no kernel IR change. Co-Authored-By: Claude <noreply@anthropic.com>

Extend the per-arch decode routing to gfx950. Its tiled 2D decode body is also one single-wave workgroup per q-block, so large-batch decode grids are bound by workgroups per CU and short-KV small-batch grids walk their tiles serially, while the 3D segment kernel splits KV across the device. gfx950 cohort: fp16/bf16 decode, head size 128 or 256, 16-token pages, no sliding window, no explicit target_ctas, 3D-supported problems (its 3D kernel has an i64 K/V path, so no cache bound). The gfx942 rule is unchanged. The route test moves to test_attn_decode_route_3d.py and covers both arches. Selection only: no kernel IR change. Co-Authored-By: Claude <noreply@anthropic.com>

…decode Head-size-256 gfx950 decode with few CTAs per segment sits outside the decode split rule and takes the pre-bump split, which reaches 128 segments at 2K KV tokens: one 16-token page per segment. Cap the split so each segment keeps at least two pages (max_seqlen_k // 32). Longer KV keeps its split; fp8 K/V and windowed problems are unchanged. Selection only: no kernel IR change. Co-Authored-By: Claude <noreply@anthropic.com>

gfx942 head-size-256 decode sat outside the decode split rule: few-CTA shapes took up to 128 segments (one 16-token page each at 2K KV tokens), many-KV-head shapes 32 segments (about 7 CTAs per CU), and short-KV shapes routed to the 2D body, which runs a few single-wave workgroups. Route gfx942 head-size-256 decode to the 3D kernel like head sizes 64/128, and split it to about 2 CTAs per CU (its LDS footprint halves the resident CTAs of the smaller heads) with at least two pages per segment. fp8 K/V, windowed and explicit target_ctas problems are unchanged. Selection only: no kernel IR change. Co-Authored-By: Claude <noreply@anthropic.com>

…k per sequence Outside the gfx942 decode split rule (few CTAs at long KV), D128 decode segments are capped by the pre-bump split, which counts q-blocks with the prefill upper bound (total_q // block_q + num_seqs). At 16 query heads per KV head (one query row per q-block) that bound doubles the count for decode and halves the cap (32 instead of 64 segments at one sequence and 8 KV heads). Count one q-block per sequence, as the decode launch grid does, for all-decode problems on gfx942. Selection only: no kernel IR change. Co-Authored-By: Claude <noreply@anthropic.com>

…ndependent Rework the gfx942 dense sliding-window geometry rule: - Narrow bands (W <= 256) now also halve the KV tile to 32 keys at D128, which halves the K/V LDS footprint so two CTAs fit per CU. - Drop the wide-window rule keyed on the runtime grid size: it made the compiled geometry depend on batch/seqlen/heads, so one persistent binary no longer served every runtime shape. - Fall back to the default geometry when the narrow one is not admitted (for example a window that leaves the last smaller query block with no key). - Record block_m and block_n with their effective values, since the base spec now resolves them per request. Re-blessed gfx942 golden dispatch_swa_d128_bf16_w128 (now block_n 32). Co-Authored-By: Claude <noreply@anthropic.com>

… a wave-filling KV split For bottom-right causal with seqlen_q != seqlen_k, the gfx942 dense base spec now turns on the causal diagonal split. It also picks a KV-segment count by wave quantization: the smallest power of two (up to 8) whose CTA count fills its last wave on the device to at least 80%. With one segment the persistent grid stays. With more, the non-persistent grid runs and splits each work item across CTAs. Requests with sinks, LSE output or a window keep one segment. Equal-length requests are unchanged. The diagonal split is recorded in the knob space for moving diagonals on gfx942. Co-Authored-By: Claude <noreply@anthropic.com>

…50 dense Reverts 6b61edf. Dropping the redundant bound shrinks the masked loop bodies but changes register allocation unpredictably: on two of the tuned sliding-window configurations the kernel crossed the 256-VGPR occupancy step or spilled more, and those shapes ran slower. Restores the previous goldens and tests. Co-Authored-By: Claude <noreply@anthropic.com>

…sistent dense loop Both K and V of tile j+1 were DMA'd into LDS at the end of iteration j, and because LLVM must assume any LDS read may alias an in-flight LDS DMA, it drained vmcnt(0) before the first K read of the next iteration. The loads were therefore never in flight while the wave had other work, even with partial_vmcnt_prefetch. The persistent body now sets lds_alias_scopes, so a read of one LDS allocation only waits for DMA into that allocation, and without a sliding window the next tile's K load is issued right after this tile's QK reads (the other K buffer was last read before the loop-top barrier). K(j+1) then loads during the softmax and PV, and V(j+1) loads during the next QK. The loop-top wait count is unchanged: V(j+1) is still the newest load. Co-Authored-By: Claude <noreply@anthropic.com>

… page-id loader The gfx950 3D segment kernel's page-id loader assumed each async K/V copy (64 lanes x 8 halves) covers whole rows: it precomputes one per-lane (row, dim) offset and advances it by KV_HALVES_PER_CALL / head_size rows per copy. For head_size 192 that quotient is fractional, so every copy after the first read rows that did not match the LDS slots it filled, and decode output was wrong. When KV_HALVES_PER_CALL is not a multiple of head_size, derive the lane's row and dim per copy instead. The spec now rejects that case with multi-page tiles or the K/V LDS swizzle, where a copy could straddle pages or the row permutation would not be closed. Every head size that divides the copy is emitted unchanged. The C++ engine mirrors both. Co-Authored-By: Claude <noreply@anthropic.com>

…nary-reuse test The default gfx942 dense spec now picks a 128-row query tile for under-filled shapes, so the cohort that checks one persistent binary serves every shape pins the shared 256-row tile explicitly. Co-Authored-By: Claude <noreply@anthropic.com>

…ing binding bind_tuning_attention_torch launched run_unified_attention_torch without k_scale / v_scale, so every pinned or tuned unified launch over an FP8 K/V cache read the cache unscaled (the runner defaults both to 1.0) and returned values in the raw FP8 range. Forward both scales, overridable per launch like softmax_scale and stream. Co-Authored-By: Claude <noreply@anthropic.com>

#14367 develop's gfx942 dense GPU-time work (#14367) made attention_gfx942_dense serve a moving bottom-right diagonal, so this request is now admitted instead of declined. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…C-aware K/V cache policy #14367 made the 3D segment kernels' K/V cache policy "stream" only when a decode's K/V exceeds the last-level cache, and "all" otherwise. Every Unified3D selection kernel is enumerated at batch 1, seqlen_k 4096, so all of them now default to "all"; the selection still pinned "stream", which the registry-regeneration gates (g1-g3) reject. Set kv_cache_policy to "all" in the selection's spec defaults and regenerate the KDP. Only the 123 Unified3D code objects change (and their kernel ids, whose stable key includes the spec); the other descriptors are byte-identical, and the Unified2D and Dense code objects are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

_use_cfvst returns True for every D128 spec regardless of dtype (since ROCm#14367), so use_cfvst=None resolves ON for BF16 D128. Correct the case study, dispatch AGENTS.md, prefill README, module docstring and a test comment that said it resolves off / is opt-in. Only K double buffering and V register prefetch are new opt-in levers in this PR. Tests are unchanged; they already assert the real default.

…PF C++ parity Since ROCm#14367 production resolves BF16 D128 CFVST on, so the bench aborted on all 12 shapes before timing. B1 now pins use_cfvst=False (the CFVST-off reference), B2 is production dispatch, downstream stages inherit the policy, and A pins use_cfvst=False to rebuild the pre-ROCm#14367 Step-0 winners (symbol gains the cfvst0 token). All 12 shapes pass the pre-timing guards. Date the case-study results to cc84e26 (pre-ROCm#14367) instead of implying they were measured at the PR head. test_attention_ir_cpp_parity now also lowers the 12 gfx942 KDB/VPF golden cases through both engines.

…ke/dense-paged-kv Conflicts with #14367 (dense GPU-time optimizations), #14679 (per-family representative-IR golden shards) and the CMake test registration changes. Beyond the textual resolution: - Paged K/V keeps the default 256x64 tile on both dense dispatchers. The smaller tiles from #14367 (gfx942 half-height and narrow-window tiles, gfx950 window tile) were measured and tested on contiguous K/V only, and at 128 query rows a wave owns 16 keys of a 64-key tile, so 8-key pages would be refused. test_paged_keeps_the_default_tile pins it. The gfx942 paged numeric test now takes the shipped paged spec from dispatch (it set paged on the contiguous spec, whose tile can now differ) and compares it with a contiguous twin on the same tile. - Paged + causal_bottom_right is refused in the shared spec (moved from the gfx950 spec): #14367 made bottom-right legal on gfx942, and no paged body is tested with the shifted diagonal. - gfx942 cfvst now also covers bf16 D128 (#14367); paged V keeps the DMA path. - Golden: the attention_dense shard is re-blessed. develop's entries are unchanged; the five paged cases are added. The gfx942 paged sliding-window hashes differ from this branch's previous golden because #14367 changed the window body.

Dispatch now gives paged K/V the smaller tiles from #14367, as it does for contiguous K/V (gfx942 half-height and narrow-window tiles with the wave skip, gfx950 window tile), when the pages fit the tile: whole pages per KV tile and each wave's rows in one page (paged_tile_fits, dense_paged_tile_fits). Otherwise it keeps the default 256x64 tile, as before. The kernels need no change: the paged loaders do one lookup per wave at any wave count, and the wave skip keeps K/V staging in every wave. On its dispatch tiles, paged dense now takes the same page sizes as unified (16, 32, 64). Paged is faster on every smaller tile measured on both arches (fp32-gated A/B), in line with contiguous on the same V path. Tests: dispatch tile pins and the fallback, page 8 refused, the helper against the spec, a compile budget for the new paged kernels, one golden case per new paged tile.

…C-aware K/V cache policy #14367 made the 3D segment kernels' K/V cache policy "stream" only when a decode's K/V exceeds the last-level cache, and "all" otherwise. Every Unified3D selection kernel is enumerated at batch 1, seqlen_k 4096, so all of them now default to "all"; the selection still pinned "stream", which the registry-regeneration gates (g1-g3) reject. Set kv_cache_policy to "all" in the selection's spec defaults and regenerate the KDP. Only the 123 Unified3D code objects change (and their kernel ids, whose stable key includes the spec); the other descriptors are byte-identical, and the Unified2D and Dense code objects are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

…14459) ISSUE ID : AICK-2493 ## Motivation #13492 made the gfx942 and gfx950 3D split-KV decode kernels faster, but the explanation lived in commit messages and a knob table. Nothing described what the segment kernel now does, how the decode split rule picks a segment count, which levers ship by default, or how to replay the checks. This PR writes that down, and corrects statements in the docs it touches that did not match the kernels. ## Technical Details 18 Markdown files, plus comment-only edits to 2 kernel source files: no code line, test, golden or generated-code change. - **Algorithm write-ups** (gfx950 `ALGORITHM.md` §5, §5.1 and §6.1, gfx942 §8): the segment kernel's transposed QK with register-resident P, the decode split rule (notation, scope, and worked examples that the split-rule test pins), and the gfx942 K/V feed. - **Case study** (new, `library/builders/common/attention_3d_decode/CASE_STUDY.md`): what a decode problem gets by default on each architecture, per-lever outcomes (qualitative, transcribed from #13492), honest losses and caveats, what the correctness gates covered, and how to replay the checks and the A/B. - **Runbook and skills**: the two reusable tactics (source-side LDS swizzle for gfx9 async copies, non-temporal read-once streams), the always-applied levers in §12.1.K2, and an arch note. - **Reference docs**: the knob reference, the strided-KV decode doc, the troubleshooting list, the env-var table, the dispatch tuning note and the routing table in `AGENTS.md`. - **Updates for #14367**, which landed on `develop` while this PR was open and changed what several statements describe: the 3D GQA gate now admits every ratio from 1 to 16, `kv_cache_policy` streams K/V only above 256 MiB per call, decode at head size 256 has its own segment count on both architectures, and the K/V LDS swizzle needs a power-of-two chunk count. The case study, both `ALGORITHM.md` files, the knob reference, the runbook and skills, the troubleshooting list, the dispatch note and the MLA design note (`library/builders/mla/DESIGN.md`) are updated to match. - **Spec comments**: the `kv_cache_policy` comments in `attention_tiled_3d.py` for gfx942 and gfx950 now say the dispatcher streams K/V only above the last-level cache size. Comment lines only; the syntax trees before and after are identical. The branch has four commits that add the content and four review follow-ups (each correcting one of the first four). Later commits warn replayers not to revert the overflow fixes that landed after #13492, fix the replay commands and the GQA and parity claims flagged in review, and update the docs for #14367 (the reviewer-requested fixes, the two spec comments with the runbook compliance table, and the MLA design note). Three merges of `develop` keep it current. Open items are listed in the case study (§7); the multi-page tiles and the page-id loader are marked "under revision" because #13128 reworks the same 3D KV loading. ## Test Plan The changes are text and comments, so the checks are on the text, on the commands it quotes, and on the two kernel files being unchanged in code: - `git diff --check`, and `pre-commit run --from-ref origin/develop --to-ref HEAD` on the branch. - Every relative link and `#anchor` in the changed files resolves (checked with a local script that is not part of this PR). - The branch merges into current `develop` without conflicts (`git merge-tree`). - The commands the case study tells readers to run: the representative IR golden check, the CPU split-rule suite, and the knob, strided-KV, cache-policy and decode-routing suites. - `tools/run_checks.py --steps guard` (the relative-path guard). - The two kernel files: no code line differs from the previous commit, and their syntax trees are identical. - No code changes, so no test file is added. ## Test Result - Whitespace and pre-commit hooks pass, 0 broken links or anchors, and the merge into `develop` is clean. - `rocke_ir_parity_harness.py --check` prints `IR parity OK` at llvm20, llvm22, llvm23 and llvm24 (run again after the kernel-comment edits), the split-rule suite passes (57 tests), and the relative-path guard passes. - The knob, strided-KV, numerics, cache-policy and decode-routing suites (the four GPU-numerics files the case study lists, plus the cache-policy and routing tests) give 81 passed and 182 skipped on the gfx90a host used here: their GPU tests need gfx942 or gfx950, and the C++-engine cases need the built engine. The GPU numerics and the forced-3D parity command were not run on gfx942 or gfx950 for this PR. - The two kernel files changed in comment lines only, so the emitted IR cannot change, and the IR golden check above agrees. The Python/C++ differential byte-identity gate was not re-run, since no emitting code changed. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>

Brings in #13955 as merged (the stacked copy drops out), #14367, #14276, #14215, #14109, #14111, #13897, #14679 and the rest of develop. Resolutions: - Parity emitters: develop's #14367 added gfx942 tiled-3D configs 18 and 19; this branch's strided decode configs move from 18-27 to 20-29 in the .py and .c, and the C strided switch moves to idx >= 20. Golden re-recorded for the family only: 0-19 unchanged in every flavor, 20-29 added. - test_registry_regeneration: keep the per-engine parametrization and adopt develop's stable-id, byte-for-byte g2 and --expect-ids g3, with the shipped folder taken per engine instead of the gfx950 globals. - Representative golden: take develop's per-family shards (#14679). - CMake, IngestorPacks.hpp, READMEs: keep both sides; SUPPORT_MATRIX_OVERVIEW re-rendered from the merged claims.

…r the develop merge rocKE #14367 changed _num_segments on gfx942: a head-size-256, block-16 decode rule (about 2 CTAs per CU, at most max_seqlen_k // 32 segments, at least 2), and the D128 pre-bump clamp now counts the decode grid (num_seqs q-blocks). gfx942SplitRuleSegments restates both; the split formula is factored out so the pre-bump can take the decode-grid num_2d. The split-rule table is regenerated and its grid widened (review): GQA ratios 2 and 8, and pinned rows for the Skv 8192 decode-rule threshold, the D128 pre-bump Skv 32768 bound, the D256 block-16 floor and cap, and a qpk-8 row the grid cannot otherwise distinguish. 8,650 -> 14,419 rows; new branch anchors. Develop API and wiring updates the merge left: - TestGfx942AttentionUnified3dDispatch: unsortedDefinitions() returns pointers (#13207); copy the kernel out before the state manager goes. - The gfx942 integration lane reserves its TOML with the census shard, as the gfx950 lane does (#14111 check). - TestKernelIngestorEngine lists the gfx942 pack's dispatch symbol (#13315).

…KE widened rocKE #14367 lets the strided decode candidate admit every GQA ratio from 1 to 16 (it masks the padding rows of the 16-row query tile) and head size 192. The engine keeps v1 scope: head sizes 64/128/256 and GQA ratios 1/2/4/8/16, the only shapes the catalog ships. The admission test now declares both as v1 exceptions next to the window one, each asserted non-vacuous, and the log lines, comments and manifest say these are scope choices, not rocKE limits. The admission table is regenerated with GQA ratios 2 (8/4) and 16 (16/1), which the engine claims (review), and head size 192, so the new exception has rows: 4,800 -> 8,400 rows. The rocKE module test's refused-ratio case moves from 3 (now built) to 32.

…ied 3D catalog rocKE #14367 makes the 3D segment kernels pick kv_cache_policy from the decode K/V size: "stream" above 256 MiB, "all" below. The policy changes only the cache bits of the segment kernel's K/V loads (the reduce kernel is identical), so it is a performance field. The manifest now declares it one, with a string KMD field (default "stream"), a metadata binding and a name suffix, so the catalog holds every policy rocKE picks: 4,566 kernels (2,150 stream, 2,416 all) in the same 138 cells, against a 1,000-kernel AOT budget (warning). The catalog is regenerated from current rocKE with the shipped bundle's ids: 1,886 kernels keep their id, 2,680 are new, and 306 whose segment counts rocKE no longer picks are retired (no retunes). The KDP is rendered whole, so it is no longer splice-ordered and the byte-for-byte g2 gate passes. The new KDP (11.5 MB) still has to be pushed to the DVC ingestor remote. Regeneration gates g1-g4 pass for both engines; bundle coverage and launch_surface --check pass. The native scorer does not read the policy yet.

assistant-librarian · 2026-11-03T11:07:06Z

TheRock Submodule Bump Activity

Newest first

2026-11-03 — Included in TheRock via ROCm/TheRock#9698.

…C-aware K/V cache policy #14367 made the 3D segment kernels' K/V cache policy "stream" only when a decode's K/V exceeds the last-level cache, and "all" otherwise. Every Unified3D selection kernel is enumerated at batch 1, seqlen_k 4096, so all of them now default to "all"; the selection still pinned "stream", which the registry-regeneration gates (g1-g3) reject. Set kv_cache_policy to "all" in the selection's spec defaults and regenerate the KDP. Only the 123 Unified3D code objects change (and their kernel ids, whose stable key includes the spec); the other descriptors are byte-identical, and the Unified2D and Dense code objects are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

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>

…r engines (#14242) ## Summary Packages rocKE's gfx950 unified attention as two hipDNN kernel-ingestor engines: `hipkernel:Gfx950AttentionUnified2D` (prefill and the decode shapes 3D does not cover) and `hipkernel:Gfx950AttentionUnified3D` (split-KV decode, a segment and a reduce launch per kernel through #14158's multi-launch `sources`). Both ship AOT, so an SDPA-forward graph on gfx950 can be served by unified attention with no JIT. Both catalogs are reduced to what a correctness-checked MI355X sweep measured as needed: Unified2D ships 222 of its 792-kernel full space and Unified3D 123 of 348, with no graph losing coverage or more than the measured noise (the final reduced-vs-full sweep is below). JIRA ID : AICK-2542 JIRA ID : AICK-2682 ## Risk Assessment Risk 2 of 5. The PR adds two engines and their descriptors. No existing engine's code, catalog or ranking changes: an unpinned SdpaFwd graph still ranks `hipkernel:Gfx950AttentionDense` first (device test `DenseStillRanksAheadWhenUnpinned`). The shared pieces it touches: - generator tooling: `registry_enumerator` full-space and multi-launch sections, the regeneration gates, provenance stamping, and two fixes in `variant_reduce`. Only `provenance.generated_from` is opt-in. The expanded registry output now always writes per-kernel `provenance.rocke_spec_id` / `rocke_tuning_id`, including Dense's 840 catalog kernels (its manifest pins a `spec_id`). The compact `variants:` selections drop it, so no shipped descriptor changes, but any other consumer of the expanded output sees the new field; - one rocKE builder, `build_unified_attention_3d_tiled_module`; - the ingestor save census's list of dispatch symbols. These are covered by their own tests and by the existing dense regeneration gates, which still pass. 3D is the first production user of `launchNames()` (#14158). ## ASIC Coverage gfx950 only: both engines ship only for gfx950, and other architectures' shards carry neither engine. PR CI has no gfx950 device lane (`hkp-probe-gfx950` packs and asserts but launches nothing), so the device evidence comes from MI355X (and, before the rebase, MI350X) runs. ## Testing Summary - Host: unit and census tests for both packs, the arch-neutral geometry, and the generator suite under `hipdnn_strict_skips`, including regeneration gates g1–g7 for all three gfx950 manifests; static descriptor gates on the installed shards (variant sets, coverage rungs 1–2, validator, launch surface). - Device, MI355X (gfx950), on the shipped (reduced) catalogs at this head's code: - engine-pinned SdpaFwd quick and standard suites against the CPU and GPU references, both engines plus a Dense regression run; - the knobs suite: every shipped kernel class forced and checked against the CPU reference, plus declines, precedence and a NaN-prefilled workspace; - a timing comparison of every engine that serves each corpus graph, and PyTorch; - RUNBOOK stage 7's final sweep: the full catalog against the reduced one in one session. - How the catalogs were reduced: a 3-round dnn-benchmark sweep with correctness on, over 188 prefill graphs (2D) and 88 decode graphs (3D), then `variant_reduce`. ## Packaging and performance statistics All numbers below come from tool output (packaging logs, device suite runs, dnn-benchmark sweeps). **Packaging** (gfx950 shard, shipped catalogs; pack time is the build's own `hkp_pack` command per engine, median of 3, relative to Dense's on the same job) | Engine | Kernels (full space) | AOT budget | Code objects, zstd | KDP | Pack time vs Dense, build default workers | Serial time per kernel vs Dense | |---|---:|---:|---:|---:|---:|---:| | Unified2D | 222 (792) | 1000 | 1.55 MiB | 0.29 MB | 0.11× | 1.58× | | Unified3D | 123 (348), 246 launches | 400 | 0.64 MiB | 0.26 MB | 0.41× | 1.06× | | Dense (catalog unchanged, for scale) | 840 (52 distinct binaries since #14059) | 1000 | 6.16 MiB | 1.26 MB | 1 | 1 | - The gfx950 kpack is 8.49 MiB at this head; the two engines are 2.19 MiB, 26% of it. Since #14059 (runtime head counts for Dense), Dense's 840 catalog kernels compile to 52 distinct binaries, but the kpack still stores one frame per kernel, so Dense grew from 5.79 to 6.16 MiB. Collapsing those entries is #14059's stated follow-up. The pack-time ratios and the full-space figure (14.45 MiB) were measured against Dense before #14059. - Provider load (`hipdnn_list_engines`, median of 5) is 1.24× its time without these descriptors. **Correctness and coverage** (MI355X, shipped catalogs) | Suite | Unified2D | Unified3D | Dense | |---|---|---|---| | SdpaFwd quick + standard, CPU and GPU reference (served / declined / failed) | 43 / 6 / 0 | 4 / 45 / 0 | 45 / 4 / 0 | | Knobs suite (forced kernel classes, declines, precedence) | 30 / 30 (both engines) | | 62 / 62 | | Corpus graphs served and validated against PyTorch | 188 / 188 | 82 / 88 | | Every decline is intentional: Unified2D leaves decode to Unified3D and declines Hkv > Hq; Unified3D serves decode only, and its 6 corpus declines are batch-256 decode graphs whose K/V is over the 32-bit addressing cap (ASM SDPA serves those). **Kernel time** (MI355X, autotuned, median of 3 rounds; ratio = Unified / other, below 1 = Unified faster; every engine's round-to-round drift ≤ 0.4%). The Unified2D rows were re-measured at this head, on the kernels develop's tiled-2D changes recompiled (see the checklist). | | Graphs both serve | Geomean ratio | Time-weighted (other / Unified) | Unified faster on | |---|---:|---:|---:|---:| | Unified3D vs ASM SDPA (decode) | 60 | 0.200 | 8.17 | 60 / 60 | | Unified3D vs PyTorch | 82 | 0.062 | 34.2 | 82 / 82 | | Unified2D vs Dense (prefill) | 179 | 2.37 | 0.38 | 0 / 179 | | Unified2D vs PyTorch | 188 | 1.38 | 0.55 | 39 / 188 | - Unified3D serves 22 decode graphs no existing hipDNN engine serves. - Unified2D is slower than Dense wherever Dense serves, and Dense keeps ranking first there; it adds coverage Dense lacks (the 9 bottom-right causal graphs in these corpora). **Variant reduction** | Engine | Kernels | Chosen eps | Measured drift | Final sweep, reduced / full kernel time (geomean, rounds 1–3) | |---|---|---:|---:|---| | Unified2D | 792 → 222 | 0.02 | 0.62% | 1.000, 0.999, 1.000 (u2d_attn); 1.000, 0.999, 1.000 (u2d_rocke) | | Unified3D | 348 → 123 | 0.03 | 1.08% | 1.004, 1.002, 0.999 (u3d_attn); 1.002, 1.007, 1.004 (u3d_rocke) | Both arms of the final sweep served the same graphs and passed correctness on every corpus. ## Testing Checklist - [x] Unit (host, no GPU) - `HIP_VISIBLE_DEVICES=-1 hip_kernel_provider_tests` - 1454 passed, including `TestKernelIngestorEngine.EveryHandlerReportsSupportsSavingExactlyWhenItSavesAndRestores`. Our build host also fails 11 Dense signature/dispatch tests: it has a `/nonexistent` directory, which those tests assume is absent. They pass in CI - Status: Passed - [x] Census - `ctest -R 'hkp-census|rocke_unified'` - 36/36 (the census inventories are the reduced catalogs') - Status: Passed - [x] Generator suite (strict skips) - `pytest -p hipdnn_strict_skips IngestorGenerator/tests` - 1646 passed, 0 skipped - Status: Passed - [x] Regeneration gates on the reduced selections - `tests/test_registry_regeneration.py` (g1–g7) 39 passed; AOT budget tests 12 passed - Status: Passed - [x] rocKE attention CPU tests and packaging guards - 2D/3D mask_type, 3D module, decode split, KV cache policy, spec-predicate parity: 145 passed, plus 6 develop KV-cache tests that need the compiled `rocke_engine`; `test_hkp_pack_producer_guards.py` and `test_gfx950_attention_unified_3d_bundle_coverage.py`: 26 passed - Status: Passed - [x] Golden IR, before the reduction - `ctest -R rocke_golden_static` - the reduction changes no rocKE code - Status: Passed - [x] Launch surface - `launch_surface.py configs/gfx950_attention_unified_{2d,3d}.engine.yaml --check` - Status: Passed - [x] Variant sets - `verify_variant_sets.py --mode structural` and `--mode full --arch gfx950 --manifest …` - 2D 222/222 distinct binaries, 3D 123/123 - Status: Passed - [x] Coverage gate rungs 1–2 - `coverage_gate.py --mode full --expect-engine …` - Status: Passed - [x] Validator - `hipdnn_validate_descriptors <shard> --expect-engine … --json` - success, no warnings - Status: Passed - [x] Packaging probe - gfx942 configure with `HIPKERNELPROVIDER_ENABLE_PACKAGING_PROBES=ON`, `ctest -R '^hkp-probe-gfx950'` - 12/12; `hkp-probe-gfx950` packs 51 UKDs of the current descriptors in a gfx942-only build - Status: Passed - [x] Device suites, MI355X - engine-pinned SdpaFwd quick + standard with GPU reference (`hip_kernel_provider_gfx950_attention_*_gpu_ref_integration_tests`) and CPU reference (`*-external-integration_{quick,standard}_suite`) - 0 failures: Unified2D 43 served / 6 declined, Unified3D 4 / 45, Dense 45 / 4 - Status: Passed - [x] Device knobs suite, MI355X - `hip_kernel_provider_integration_tests --gtest_filter='*Gfx950AttentionUnified*'` - 30/30; Dense knobs 62/62 - Status: Passed - [x] Fault re-check, MI355X - the whole `u2d_rocke` corpus in one autotuned process, twice and then with serialized launches - 93/93 each, 0 errors - Status: Passed - [x] Reduction sweep + `variant_reduce.py` - 3 rounds, correctness on; 2D `SWEEP_DONE`, 3D `SWEEP_DONE`; 2D 792 → 222, 3D 348 → 123 - Status: Passed - [x] Final sweep, reduced against full (RUNBOOK §7) - same coverage, correctness passed on both arms, ratios above - Status: Passed - [x] Unified3D catalog regenerated for #14367's LLC-aware K/V cache policy - all 123 kernels now `kv_cache_policy: "all"` (was `"stream"`), same kernel classes; MI355X re-run: suites and knobs unchanged, 246/246 decode rows correct, kernel time against the `"stream"` catalog geomean 0.998, time-weighted 0.999 - Status: Passed - [x] Device evidence at this head (rebased onto develop 2026-11-03, after the review-round refactors) - Dense's and Unified3D's code objects are byte-identical to the previous device runs; develop's #13652 recompiled 201 of the 222 Unified2D code objects (same kernels and ids), and the Unified2D/3D matchers moved onto a shared header, so both engines' device runs were repeated at this head: suites Unified2D 43 / 6 / 0, Unified3D 4 / 45 / 0, Dense 45 / 4 / 0 (passed / skipped / failed), knobs 30/30 and 62/62; compare 188/188 Unified2D graphs (564/564 rows) and 82 Unified3D graphs (246/246 rows) correct against PyTorch in all 3 rounds. Kernel time vs the previous runs: Unified2D geomean 0.985 (none of 188 graphs >2% slower), Unified3D 1.003 (ASM_SDPA control 1.006). A final rebase picked up #14059, which recompiled Dense's 840 code objects only (Unified2D/3D byte-identical to the run above); the suites and knobs were re-run at this head with the same results, `DenseStillRanksAheadWhenUnpinned` included - Status: Passed - [x] DVC - both reduced KDPs on the `ingestor` remote; `dvc status -c` in sync - Status: Passed - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - **Unified2D pack** (`Gfx950AttentionUnified2DNative.cpp`): matches BSHD SDPA forward on the identity paged-KV layout and builds block_tables, seq_lens and query_start_len in the workspace. - It declines: batch > 511 (the baked 9-iteration binary search), Skv not a multiple of 16, i32 KV over 0x7FFF0000, any window, and the decode cells Unified3D covers (Sq = 1 with the GQA group dividing 16). - Engine knobs: tile_size, num_warps, block_m_per_warp, use_mfma_32x32 and use_transposed_qk_32x32. - **Unified3D pack** (`Gfx950AttentionUnified3DNative.cpp`): decode only. - `launchNames()` is `{segment, reduce}`, and both definitions are resolved through `launchKernelDefinition` with a signature check before any load. - The segment launch writes `[token, Hq, seg, D]` scratch and the reduce always runs, so the workspace needs no zeroing. - Knob: num_segments. No batch limit. - **Arch-neutral geometry** (`UnifiedAttentionGeometry.hpp`): the launch geometry and workspace layouts take a `UnifiedAttentionArchTraits`, and the gfx950 headers are thin wrappers, so a gfx942 pack (AICK-2544) only supplies traits. - **rocKE**: `build_unified_attention_3d_tiled_module` packs the segment and its derived reduce (`reduce_spec_for`) as one `KernelModule`. The module predicate gates both halves. The 3D pack compiles develop's running-max fix for split-KV (#14276). - **Generator**: - `enumerate.also` sections may use `space: full` and carry manifest `launches`. The enumerator refuses a registry that pairs a reduce when the manifest declares no launches, and keeps one default per cell. - The gates are launch-aware: g1 keys on launches, g3 uses `--subset` for full spaces, g4 compares launch sets, and g6 walks every launch and checks module roles. - g2 and g3 run per engine with the shipped ids reused (#14109), comparing every descriptor byte for byte. - `generate.py --stamp-provenance` stamps `provenance.generated_from` (config, manifest, generator and rocKE commits). It is opt-in, so a default regeneration stays byte-identical. - `variant_reduce` fixes: - it reads the format-2 winner lines the runtime has written since #13028; - it keys a graph on its content with `Graph.id` (minted fresh at each finalization) zeroed, so sweep rounds match. - **Catalogs** (the manifests enumerate the full spaces; the selections are `variant_reduce`'s): - Unified2D: 222 kernels over 198 cells (33 head configs × bf16/fp16 × 3 masks): the narrow nw4/mw16/t64 geometry in every cell, transposed 32x32 nw4/mw32/t64 in 21 and narrow nw2/mw16/t32 in 3. The wide 32x32 geometry won no graph and is not shipped. - Unified3D: 123 kernels: num_segments 16 in all 116 decode cells, 32 or 64 in 7 bf16 bottom-right D128 variants. The K/V cache policy follows #14367's LLC-aware rule, which picks `"all"` for every shipped decode cell. - **Descriptors**: both are generated with develop's generator (post-#13028 UED/UHD, UKD `sources` for 3D), every kernel id reused from the full catalog. A regeneration with the shipped ids reused reproduces them byte for byte. The KDPs live in DVC (`ingestor` remote). - **Shared code** (one home per fact): - `packs/UnifiedAttentionGraphMatch.hpp`: the graph-match helpers both unified packs use (operand/BSHD checks, SDPA node, mask classification, dtype, metadata readers); the tests' graph builder is `UnifiedAttentionTestGraphs.hpp`. - rocKE `kernels/common/attention_mask_type.py` (mask ordinals and band -> kind, used by both tiled specs and re-exported by dispatch) and `attention_split_kv_reduce.py` (the segment -> reduce derivation the dispatcher and the module export share). - CMake registers Dense, Unified2D and Unified3D by one loop per lane (external integration, gpu_ref, census). - **Tests**: - The device knobs suite (`IntegrationGpuGfx950AttentionUnifiedKnobs.cpp`) adds a harness `prepareWorkspace` hook. It forces the default geometry/segment count across the head configurations and each other shipped class on a cell that carries it, and `CoversEveryCatalogKernelClass` holds those tables to the catalog. - The engine-pinned `gpu_ref` targets and external-integration lanes cover both engines. - The ingestor save census names all seven pack dispatch symbols. - The Unified2D/3D integration TOMLs are reserved with the gfx950 census shard (`hkp_reserve_census_shard(... FILES ...)`, as #14573 requires of every shard entry). **Constraints and follow-ups** - #13372 and #13889 are both in develop; this branch no longer carries a copy of #13889. - Heads are catalog cells (baked), window = 0, 2D batch ≤ 511, i32 KV ≤ 0x7FFF0000. D256 is deferred for both engines. The kernarg work is tracked in AICK-2675. - There is no 2D/3D `select_path` crossover replay: the split is the static decode rule above. - Runtimes older than UKD 1.1 skip the 3D descriptors. Strided-KV 3D is a follow-up. - Follow-up: execution-plan save/restore (#13315) for both handlers. Today they report `supportsSaving() == false`, as the RDNA handler does. - One earlier MI355X sweep hit a single GPU fault in Unified2D on `u2d_rocke/rocke_sdpa_035` (bf16, D 128, Hq 64 / Hkv 8, causal, S 512). It has not reproduced since: not in the forced geometries on that cell, 15 single-graph runs, 4 full-corpus runs on MI350X, or the MI355X full-corpus re-check above. - Unified2D kernel time moved with develop's tiled-2D changes (#14542, #14616), not with this PR: at this head it is 1.035× (geomean) its pre-rebase time on the same 188 graphs, up to 1.36× on long causal bf16 D128 prefill (S 8192), while Dense in the same runs is flat (0.997). Dense still ranks first wherever it serves; the table above is at this head. - Measurement note: the pinned dnn-benchmarking (v1 result schema, which `sweep.py` reads) has no bottom-right causal PyTorch reference. The sweeps used upstream dnn-benchmarking #84's reference, with the 9 bottom-right corpus graphs written as `causal_mask: false, right_bound: 0, diagonal_alignment: BOTTOM_RIGHT` (the same band under hipDNN's `resolveDiagonalBand`): #84 reads `causal_mask: true` as top-left regardless of the alignment. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>

…14929) ISSUE ID : AICK-2955 ## Problem During decode tuning on gfx942, `gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm` repeatedly hit illegal memory accesses at bf16, B=64, Sq=1, Skv=4096, Hq=64, Hkv=4, D=128, block size 16. One fault hung the GPU queue, the candidate hit the 300 s timeout, and the tuning job failed. The D128 body of `build_gfx942_4warp_gqa` loads a full TOKBLK-row Q block (128 rows, or 32 tokens when head-folded) and whole KV tiles, and bounds neither. A one-token decode sequence reads up to 127 Q rows past its token, which runs off the end of the Q tensor for the last sequences. A KV length that is not a tile multiple also reads block-table entries past the sequence's pages. #14367 added bounds to the lean D192/D256 body only. ## Fix - Clamp every Q row to the last token of its sequence (plain and head-fold paths). Each lane's clamped row is loop-invariant, so it is computed once before the KV loop; recomputing it per tile measured up to 4% slower on fp16 block-size-64 kernels. - Clamp every K/V block-table index to the sequence's last page. Keys past `klen` then read rows of the last valid page, which are masked to -inf (P = 0); clamped Q rows are never stored. Only load addresses change. - This clamps the page, not the key as the lean body does, so the per-lane page/offset split stays out of the KV loop. A per-tile key clamp measured 3–7% slower on the production fold kernels. - Reject a `block_size` that is not a power of two or does not divide the KV tile at build time. Every admitted size is 16, 32 or 64. - The 4-warp builder has no C++ twin, so this is Python-only. - Re-bless `attention_d128_swa`, the only golden case that goes through the D128 body. Re-blessing `attention_d256` (the lean body) gives byte-identical output. Tests: - **CPU (the CI guard; the gfx94X lane has no torch):** `TestD128BodyBoundsReads` in `test_d256_gfx942_fast_routing.py` checks that every block-table, Q and K/V read of the fold, plain, GQA-8 bs32 and bs64 variants is bounded. These checks fail on develop's D128 body. Host tests also check that the builder rejects a `block_size` that is not a power of two or does not divide the KV tile, and that the page clamp shifts by exactly log2(`block_size`). The file is now in `_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS`. - **GPU:** the lean-body oracle is renamed `test_gfx942_4wgqa_bounds_numeric.py` and gains D128 cases: sliding-window prefill at odd lengths on the production route, and the pinned decode shape. Block-table padding points about 2 GiB past the cache, so a read past the last page faults. Output is checked against an fp32 reference, and each case asserts the fold or plain kernel it builds. Enrolled in `_ROCKE_LIBRARY_GPU_TEST_RELPATHS`. - Both GPU files gate on `gcnArchName` instead of the device marketing name. `test_gfx942_gqa_head_fold_numeric.py` was skipping on gfx942 parts whose name lacks `MI300`. ## Verification | Check | Result | |---|---| | `test_rocke_ci_static.py -k golden` (every family, llvm20/22/23/24) | 8 passed | | `check_byte_identity.py --only attention`, all flavors (merged head, engine rebuilt) | GREEN. This gate compares the Python and C++ engines, so it does not cover the 4-warp kernel. | | Host pytest: 4-warp/fold emit, CPU bounds probe, dispatch, LSE, tuning-space tests | 153 passed, 28 skipped (GPU-only cases) | | Full `platform/tests` (merged with develop at 835ea36) | 6714 passed, 4418 skipped, 4 failed. The 4 failures are `test_engine_linkage.py::test_installed_ctest_entries_get_scoped_environment_contracts` (the `rocke_library_sdpa_pytest` CTest entry has no command); they fail identically with this PR's CMake lines removed. Skips: cross-engine checks that need the `rocke_engine` extension (not built here; this kernel has no C++ twin), GPU tests for other architectures, and native C++ test binaries not configured. | | gfx942 GPU (merged head): `test_gfx942_4wgqa_bounds_numeric.py` | 24 passed | | gfx942 GPU: `test_gfx942_gqa_head_fold_numeric.py` | 6 passed | | `pre-commit run` on the changed files | Passed | | Outputs, develop vs this branch | Bit-identical on 23 kernels (12 production, 11 pinned tuning spec) on one gfx942 part and 5 production kernels on a second | Kernel-only time, this branch vs develop: median device time of 30 launches from `torch.profiler`, both sides back-to-back in one job on the same GPU. LLVM flavor llvm22, torch 2.12.1+rocm7.2. Production route, `run_unified_attention_torch(backend="auto")`, 2 rounds, gfx942: | Production kernel | S4096 w512 | S8192 w4096 | |---|---|---| | bf16 fold bs16 | 1.008 | 1.005 | | bf16 fold bs32 | 1.004 | 1.016 | | bf16 plain bs64 | 0.999 | 1.005 | | fp16 plain bs16 | 0.994 | 0.992 | | fp16 plain bs32 | 0.989 | 1.001 | | fp16 plain bs64 | 1.011 | 1.004 | A second gfx942 part agrees within 1.3% on the five production kernels measured there (fp16 bs64 1.013 / 1.007, bf16 bs64 0.991, fp16 bs16 1.006, bf16 fold bs16 1.010). Pinned 4-warp tuning spec through the dispatcher, 1 round, gfx942: | Kernel | new / develop | |---|---| | Decode at the reported shape (B=64, Hq=64, Hkv=4, bs16) | 0.785 | | Decode B=32 bs32 / B=64 bs64 | 0.864 / 0.943 | | Causal prefill, bs16 | 0.853 | | Other pinned variants (fold, plain, GQA-8, block size 64) | 0.947 – 1.014 | Decode likely gains because the clamped Q loads all hit the sequence's one token instead of streaming 127 extra rows (no hardware counters were collected). --------- Co-authored-by: Claude <noreply@anthropic.com>

…r engines (#14242) ## Summary Packages rocKE's gfx950 unified attention as two hipDNN kernel-ingestor engines: `hipkernel:Gfx950AttentionUnified2D` (prefill and the decode shapes 3D does not cover) and `hipkernel:Gfx950AttentionUnified3D` (split-KV decode, a segment and a reduce launch per kernel through #14158's multi-launch `sources`). Both ship AOT, so an SDPA-forward graph on gfx950 can be served by unified attention with no JIT. Both catalogs are reduced to what a correctness-checked MI355X sweep measured as needed: Unified2D ships 222 of its 792-kernel full space and Unified3D 123 of 348, with no graph losing coverage or more than the measured noise (the final reduced-vs-full sweep is below). JIRA ID : AICK-2542 JIRA ID : AICK-2682 ## Risk Assessment Risk 2 of 5. The PR adds two engines and their descriptors. No existing engine's code, catalog or ranking changes: an unpinned SdpaFwd graph still ranks `hipkernel:Gfx950AttentionDense` first (device test `DenseStillRanksAheadWhenUnpinned`). The shared pieces it touches: - generator tooling: `registry_enumerator` full-space and multi-launch sections, the regeneration gates, provenance stamping, and two fixes in `variant_reduce`. Only `provenance.generated_from` is opt-in. The expanded registry output now always writes per-kernel `provenance.rocke_spec_id` / `rocke_tuning_id`, including Dense's 840 catalog kernels (its manifest pins a `spec_id`). The compact `variants:` selections drop it, so no shipped descriptor changes, but any other consumer of the expanded output sees the new field; - one rocKE builder, `build_unified_attention_3d_tiled_module`; - the ingestor save census's list of dispatch symbols. These are covered by their own tests and by the existing dense regeneration gates, which still pass. 3D is the first production user of `launchNames()` (#14158). ## ASIC Coverage gfx950 only: both engines ship only for gfx950, and other architectures' shards carry neither engine. PR CI has no gfx950 device lane (`hkp-probe-gfx950` packs and asserts but launches nothing), so the device evidence comes from MI355X (and, before the rebase, MI350X) runs. ## Testing Summary - Host: unit and census tests for both packs, the arch-neutral geometry, and the generator suite under `hipdnn_strict_skips`, including regeneration gates g1–g7 for all three gfx950 manifests; static descriptor gates on the installed shards (variant sets, coverage rungs 1–2, validator, launch surface). - Device, MI355X (gfx950), on the shipped (reduced) catalogs at this head's code: - engine-pinned SdpaFwd quick and standard suites against the CPU and GPU references, both engines plus a Dense regression run; - the knobs suite: every shipped kernel class forced and checked against the CPU reference, plus declines, precedence and a NaN-prefilled workspace; - a timing comparison of every engine that serves each corpus graph, and PyTorch; - RUNBOOK stage 7's final sweep: the full catalog against the reduced one in one session. - How the catalogs were reduced: a 3-round dnn-benchmark sweep with correctness on, over 188 prefill graphs (2D) and 88 decode graphs (3D), then `variant_reduce`. ## Packaging and performance statistics All numbers below come from tool output (packaging logs, device suite runs, dnn-benchmark sweeps). **Packaging** (gfx950 shard, shipped catalogs; pack time is the build's own `hkp_pack` command per engine, median of 3, relative to Dense's on the same job) | Engine | Kernels (full space) | AOT budget | Code objects, zstd | KDP | Pack time vs Dense, build default workers | Serial time per kernel vs Dense | |---|---:|---:|---:|---:|---:|---:| | Unified2D | 222 (792) | 1000 | 1.55 MiB | 0.29 MB | 0.11× | 1.58× | | Unified3D | 123 (348), 246 launches | 400 | 0.64 MiB | 0.26 MB | 0.41× | 1.06× | | Dense (catalog unchanged, for scale) | 840 (52 distinct binaries since #14059) | 1000 | 6.16 MiB | 1.26 MB | 1 | 1 | - The gfx950 kpack is 8.49 MiB at this head; the two engines are 2.19 MiB, 26% of it. Since #14059 (runtime head counts for Dense), Dense's 840 catalog kernels compile to 52 distinct binaries, but the kpack still stores one frame per kernel, so Dense grew from 5.79 to 6.16 MiB. Collapsing those entries is #14059's stated follow-up. The pack-time ratios and the full-space figure (14.45 MiB) were measured against Dense before #14059. - Provider load (`hipdnn_list_engines`, median of 5) is 1.24× its time without these descriptors. **Correctness and coverage** (MI355X, shipped catalogs) | Suite | Unified2D | Unified3D | Dense | |---|---|---|---| | SdpaFwd quick + standard, CPU and GPU reference (served / declined / failed) | 43 / 6 / 0 | 4 / 45 / 0 | 45 / 4 / 0 | | Knobs suite (forced kernel classes, declines, precedence) | 30 / 30 (both engines) | | 62 / 62 | | Corpus graphs served and validated against PyTorch | 188 / 188 | 82 / 88 | | Every decline is intentional: Unified2D leaves decode to Unified3D and declines Hkv > Hq; Unified3D serves decode only, and its 6 corpus declines are batch-256 decode graphs whose K/V is over the 32-bit addressing cap (ASM SDPA serves those). **Kernel time** (MI355X, autotuned, median of 3 rounds; ratio = Unified / other, below 1 = Unified faster; every engine's round-to-round drift ≤ 0.4%). The Unified2D rows were re-measured at this head, on the kernels develop's tiled-2D changes recompiled (see the checklist). | | Graphs both serve | Geomean ratio | Time-weighted (other / Unified) | Unified faster on | |---|---:|---:|---:|---:| | Unified3D vs ASM SDPA (decode) | 60 | 0.200 | 8.17 | 60 / 60 | | Unified3D vs PyTorch | 82 | 0.062 | 34.2 | 82 / 82 | | Unified2D vs Dense (prefill) | 179 | 2.37 | 0.38 | 0 / 179 | | Unified2D vs PyTorch | 188 | 1.38 | 0.55 | 39 / 188 | - Unified3D serves 22 decode graphs no existing hipDNN engine serves. - Unified2D is slower than Dense wherever Dense serves, and Dense keeps ranking first there; it adds coverage Dense lacks (the 9 bottom-right causal graphs in these corpora). **Variant reduction** | Engine | Kernels | Chosen eps | Measured drift | Final sweep, reduced / full kernel time (geomean, rounds 1–3) | |---|---|---:|---:|---| | Unified2D | 792 → 222 | 0.02 | 0.62% | 1.000, 0.999, 1.000 (u2d_attn); 1.000, 0.999, 1.000 (u2d_rocke) | | Unified3D | 348 → 123 | 0.03 | 1.08% | 1.004, 1.002, 0.999 (u3d_attn); 1.002, 1.007, 1.004 (u3d_rocke) | Both arms of the final sweep served the same graphs and passed correctness on every corpus. ## Testing Checklist - [x] Unit (host, no GPU) - `HIP_VISIBLE_DEVICES=-1 hip_kernel_provider_tests` - 1454 passed, including `TestKernelIngestorEngine.EveryHandlerReportsSupportsSavingExactlyWhenItSavesAndRestores`. Our build host also fails 11 Dense signature/dispatch tests: it has a `/nonexistent` directory, which those tests assume is absent. They pass in CI - Status: Passed - [x] Census - `ctest -R 'hkp-census|rocke_unified'` - 36/36 (the census inventories are the reduced catalogs') - Status: Passed - [x] Generator suite (strict skips) - `pytest -p hipdnn_strict_skips IngestorGenerator/tests` - 1646 passed, 0 skipped - Status: Passed - [x] Regeneration gates on the reduced selections - `tests/test_registry_regeneration.py` (g1–g7) 39 passed; AOT budget tests 12 passed - Status: Passed - [x] rocKE attention CPU tests and packaging guards - 2D/3D mask_type, 3D module, decode split, KV cache policy, spec-predicate parity: 145 passed, plus 6 develop KV-cache tests that need the compiled `rocke_engine`; `test_hkp_pack_producer_guards.py` and `test_gfx950_attention_unified_3d_bundle_coverage.py`: 26 passed - Status: Passed - [x] Golden IR, before the reduction - `ctest -R rocke_golden_static` - the reduction changes no rocKE code - Status: Passed - [x] Launch surface - `launch_surface.py configs/gfx950_attention_unified_{2d,3d}.engine.yaml --check` - Status: Passed - [x] Variant sets - `verify_variant_sets.py --mode structural` and `--mode full --arch gfx950 --manifest …` - 2D 222/222 distinct binaries, 3D 123/123 - Status: Passed - [x] Coverage gate rungs 1–2 - `coverage_gate.py --mode full --expect-engine …` - Status: Passed - [x] Validator - `hipdnn_validate_descriptors <shard> --expect-engine … --json` - success, no warnings - Status: Passed - [x] Packaging probe - gfx942 configure with `HIPKERNELPROVIDER_ENABLE_PACKAGING_PROBES=ON`, `ctest -R '^hkp-probe-gfx950'` - 12/12; `hkp-probe-gfx950` packs 51 UKDs of the current descriptors in a gfx942-only build - Status: Passed - [x] Device suites, MI355X - engine-pinned SdpaFwd quick + standard with GPU reference (`hip_kernel_provider_gfx950_attention_*_gpu_ref_integration_tests`) and CPU reference (`*-external-integration_{quick,standard}_suite`) - 0 failures: Unified2D 43 served / 6 declined, Unified3D 4 / 45, Dense 45 / 4 - Status: Passed - [x] Device knobs suite, MI355X - `hip_kernel_provider_integration_tests --gtest_filter='*Gfx950AttentionUnified*'` - 30/30; Dense knobs 62/62 - Status: Passed - [x] Fault re-check, MI355X - the whole `u2d_rocke` corpus in one autotuned process, twice and then with serialized launches - 93/93 each, 0 errors - Status: Passed - [x] Reduction sweep + `variant_reduce.py` - 3 rounds, correctness on; 2D `SWEEP_DONE`, 3D `SWEEP_DONE`; 2D 792 → 222, 3D 348 → 123 - Status: Passed - [x] Final sweep, reduced against full (RUNBOOK §7) - same coverage, correctness passed on both arms, ratios above - Status: Passed - [x] Unified3D catalog regenerated for #14367's LLC-aware K/V cache policy - all 123 kernels now `kv_cache_policy: "all"` (was `"stream"`), same kernel classes; MI355X re-run: suites and knobs unchanged, 246/246 decode rows correct, kernel time against the `"stream"` catalog geomean 0.998, time-weighted 0.999 - Status: Passed - [x] Device evidence at this head (rebased onto develop 2026-11-03, after the review-round refactors) - Dense's and Unified3D's code objects are byte-identical to the previous device runs; develop's #13652 recompiled 201 of the 222 Unified2D code objects (same kernels and ids), and the Unified2D/3D matchers moved onto a shared header, so both engines' device runs were repeated at this head: suites Unified2D 43 / 6 / 0, Unified3D 4 / 45 / 0, Dense 45 / 4 / 0 (passed / skipped / failed), knobs 30/30 and 62/62; compare 188/188 Unified2D graphs (564/564 rows) and 82 Unified3D graphs (246/246 rows) correct against PyTorch in all 3 rounds. Kernel time vs the previous runs: Unified2D geomean 0.985 (none of 188 graphs >2% slower), Unified3D 1.003 (ASM_SDPA control 1.006). A final rebase picked up #14059, which recompiled Dense's 840 code objects only (Unified2D/3D byte-identical to the run above); the suites and knobs were re-run at this head with the same results, `DenseStillRanksAheadWhenUnpinned` included - Status: Passed - [x] DVC - both reduced KDPs on the `ingestor` remote; `dvc status -c` in sync - Status: Passed - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - **Unified2D pack** (`Gfx950AttentionUnified2DNative.cpp`): matches BSHD SDPA forward on the identity paged-KV layout and builds block_tables, seq_lens and query_start_len in the workspace. - It declines: batch > 511 (the baked 9-iteration binary search), Skv not a multiple of 16, i32 KV over 0x7FFF0000, any window, and the decode cells Unified3D covers (Sq = 1 with the GQA group dividing 16). - Engine knobs: tile_size, num_warps, block_m_per_warp, use_mfma_32x32 and use_transposed_qk_32x32. - **Unified3D pack** (`Gfx950AttentionUnified3DNative.cpp`): decode only. - `launchNames()` is `{segment, reduce}`, and both definitions are resolved through `launchKernelDefinition` with a signature check before any load. - The segment launch writes `[token, Hq, seg, D]` scratch and the reduce always runs, so the workspace needs no zeroing. - Knob: num_segments. No batch limit. - **Arch-neutral geometry** (`UnifiedAttentionGeometry.hpp`): the launch geometry and workspace layouts take a `UnifiedAttentionArchTraits`, and the gfx950 headers are thin wrappers, so a gfx942 pack (AICK-2544) only supplies traits. - **rocKE**: `build_unified_attention_3d_tiled_module` packs the segment and its derived reduce (`reduce_spec_for`) as one `KernelModule`. The module predicate gates both halves. The 3D pack compiles develop's running-max fix for split-KV (#14276). - **Generator**: - `enumerate.also` sections may use `space: full` and carry manifest `launches`. The enumerator refuses a registry that pairs a reduce when the manifest declares no launches, and keeps one default per cell. - The gates are launch-aware: g1 keys on launches, g3 uses `--subset` for full spaces, g4 compares launch sets, and g6 walks every launch and checks module roles. - g2 and g3 run per engine with the shipped ids reused (#14109), comparing every descriptor byte for byte. - `generate.py --stamp-provenance` stamps `provenance.generated_from` (config, manifest, generator and rocKE commits). It is opt-in, so a default regeneration stays byte-identical. - `variant_reduce` fixes: - it reads the format-2 winner lines the runtime has written since #13028; - it keys a graph on its content with `Graph.id` (minted fresh at each finalization) zeroed, so sweep rounds match. - **Catalogs** (the manifests enumerate the full spaces; the selections are `variant_reduce`'s): - Unified2D: 222 kernels over 198 cells (33 head configs × bf16/fp16 × 3 masks): the narrow nw4/mw16/t64 geometry in every cell, transposed 32x32 nw4/mw32/t64 in 21 and narrow nw2/mw16/t32 in 3. The wide 32x32 geometry won no graph and is not shipped. - Unified3D: 123 kernels: num_segments 16 in all 116 decode cells, 32 or 64 in 7 bf16 bottom-right D128 variants. The K/V cache policy follows #14367's LLC-aware rule, which picks `"all"` for every shipped decode cell. - **Descriptors**: both are generated with develop's generator (post-#13028 UED/UHD, UKD `sources` for 3D), every kernel id reused from the full catalog. A regeneration with the shipped ids reused reproduces them byte for byte. The KDPs live in DVC (`ingestor` remote). - **Shared code** (one home per fact): - `packs/UnifiedAttentionGraphMatch.hpp`: the graph-match helpers both unified packs use (operand/BSHD checks, SDPA node, mask classification, dtype, metadata readers); the tests' graph builder is `UnifiedAttentionTestGraphs.hpp`. - rocKE `kernels/common/attention_mask_type.py` (mask ordinals and band -> kind, used by both tiled specs and re-exported by dispatch) and `attention_split_kv_reduce.py` (the segment -> reduce derivation the dispatcher and the module export share). - CMake registers Dense, Unified2D and Unified3D by one loop per lane (external integration, gpu_ref, census). - **Tests**: - The device knobs suite (`IntegrationGpuGfx950AttentionUnifiedKnobs.cpp`) adds a harness `prepareWorkspace` hook. It forces the default geometry/segment count across the head configurations and each other shipped class on a cell that carries it, and `CoversEveryCatalogKernelClass` holds those tables to the catalog. - The engine-pinned `gpu_ref` targets and external-integration lanes cover both engines. - The ingestor save census names all seven pack dispatch symbols. - The Unified2D/3D integration TOMLs are reserved with the gfx950 census shard (`hkp_reserve_census_shard(... FILES ...)`, as #14573 requires of every shard entry). **Constraints and follow-ups** - #13372 and #13889 are both in develop; this branch no longer carries a copy of #13889. - Heads are catalog cells (baked), window = 0, 2D batch ≤ 511, i32 KV ≤ 0x7FFF0000. D256 is deferred for both engines. The kernarg work is tracked in AICK-2675. - There is no 2D/3D `select_path` crossover replay: the split is the static decode rule above. - Runtimes older than UKD 1.1 skip the 3D descriptors. Strided-KV 3D is a follow-up. - Follow-up: execution-plan save/restore (#13315) for both handlers. Today they report `supportsSaving() == false`, as the RDNA handler does. - One earlier MI355X sweep hit a single GPU fault in Unified2D on `u2d_rocke/rocke_sdpa_035` (bf16, D 128, Hq 64 / Hkv 8, causal, S 512). It has not reproduced since: not in the forced geometries on that cell, 15 single-graph runs, 4 full-corpus runs on MI350X, or the MI355X full-corpus re-check above. - Unified2D kernel time moved with develop's tiled-2D changes (#14542, #14616), not with this PR: at this head it is 1.035× (geomean) its pre-rebase time on the same 188 graphs, up to 1.36× on long causal bf16 D128 prefill (S 8192), while Dense in the same runs is flat (0.997). Dense still ranks first wherever it serves; the table above is at this head. - Measurement note: the pinned dnn-benchmarking (v1 result schema, which `sweep.py` reads) has no bottom-right causal PyTorch reference. The sweeps used upstream dnn-benchmarking #84's reference, with the 9 bottom-right corpus graphs written as `causal_mask: false, right_bound: 0, diagonal_alignment: BOTTOM_RIGHT` (the same band under hipDNN's `resolveDiagonalBand`): #84 reads `causal_mask: true` as top-left regardless of the alignment. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>

…14929) ISSUE ID : AICK-2955 ## Problem During decode tuning on gfx942, `gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm` repeatedly hit illegal memory accesses at bf16, B=64, Sq=1, Skv=4096, Hq=64, Hkv=4, D=128, block size 16. One fault hung the GPU queue, the candidate hit the 300 s timeout, and the tuning job failed. The D128 body of `build_gfx942_4warp_gqa` loads a full TOKBLK-row Q block (128 rows, or 32 tokens when head-folded) and whole KV tiles, and bounds neither. A one-token decode sequence reads up to 127 Q rows past its token, which runs off the end of the Q tensor for the last sequences. A KV length that is not a tile multiple also reads block-table entries past the sequence's pages. #14367 added bounds to the lean D192/D256 body only. ## Fix - Clamp every Q row to the last token of its sequence (plain and head-fold paths). Each lane's clamped row is loop-invariant, so it is computed once before the KV loop; recomputing it per tile measured up to 4% slower on fp16 block-size-64 kernels. - Clamp every K/V block-table index to the sequence's last page. Keys past `klen` then read rows of the last valid page, which are masked to -inf (P = 0); clamped Q rows are never stored. Only load addresses change. - This clamps the page, not the key as the lean body does, so the per-lane page/offset split stays out of the KV loop. A per-tile key clamp measured 3–7% slower on the production fold kernels. - Reject a `block_size` that is not a power of two or does not divide the KV tile at build time. Every admitted size is 16, 32 or 64. - The 4-warp builder has no C++ twin, so this is Python-only. - Re-bless `attention_d128_swa`, the only golden case that goes through the D128 body. Re-blessing `attention_d256` (the lean body) gives byte-identical output. Tests: - **CPU (the CI guard; the gfx94X lane has no torch):** `TestD128BodyBoundsReads` in `test_d256_gfx942_fast_routing.py` checks that every block-table, Q and K/V read of the fold, plain, GQA-8 bs32 and bs64 variants is bounded. These checks fail on develop's D128 body. Host tests also check that the builder rejects a `block_size` that is not a power of two or does not divide the KV tile, and that the page clamp shifts by exactly log2(`block_size`). The file is now in `_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS`. - **GPU:** the lean-body oracle is renamed `test_gfx942_4wgqa_bounds_numeric.py` and gains D128 cases: sliding-window prefill at odd lengths on the production route, and the pinned decode shape. Block-table padding points about 2 GiB past the cache, so a read past the last page faults. Output is checked against an fp32 reference, and each case asserts the fold or plain kernel it builds. Enrolled in `_ROCKE_LIBRARY_GPU_TEST_RELPATHS`. - Both GPU files gate on `gcnArchName` instead of the device marketing name. `test_gfx942_gqa_head_fold_numeric.py` was skipping on gfx942 parts whose name lacks `MI300`. ## Verification | Check | Result | |---|---| | `test_rocke_ci_static.py -k golden` (every family, llvm20/22/23/24) | 8 passed | | `check_byte_identity.py --only attention`, all flavors (merged head, engine rebuilt) | GREEN. This gate compares the Python and C++ engines, so it does not cover the 4-warp kernel. | | Host pytest: 4-warp/fold emit, CPU bounds probe, dispatch, LSE, tuning-space tests | 153 passed, 28 skipped (GPU-only cases) | | Full `platform/tests` (merged with develop at 835ea36) | 6714 passed, 4418 skipped, 4 failed. The 4 failures are `test_engine_linkage.py::test_installed_ctest_entries_get_scoped_environment_contracts` (the `rocke_library_sdpa_pytest` CTest entry has no command); they fail identically with this PR's CMake lines removed. Skips: cross-engine checks that need the `rocke_engine` extension (not built here; this kernel has no C++ twin), GPU tests for other architectures, and native C++ test binaries not configured. | | gfx942 GPU (merged head): `test_gfx942_4wgqa_bounds_numeric.py` | 24 passed | | gfx942 GPU: `test_gfx942_gqa_head_fold_numeric.py` | 6 passed | | `pre-commit run` on the changed files | Passed | | Outputs, develop vs this branch | Bit-identical on 23 kernels (12 production, 11 pinned tuning spec) on one gfx942 part and 5 production kernels on a second | Kernel-only time, this branch vs develop: median device time of 30 launches from `torch.profiler`, both sides back-to-back in one job on the same GPU. LLVM flavor llvm22, torch 2.12.1+rocm7.2. Production route, `run_unified_attention_torch(backend="auto")`, 2 rounds, gfx942: | Production kernel | S4096 w512 | S8192 w4096 | |---|---|---| | bf16 fold bs16 | 1.008 | 1.005 | | bf16 fold bs32 | 1.004 | 1.016 | | bf16 plain bs64 | 0.999 | 1.005 | | fp16 plain bs16 | 0.994 | 0.992 | | fp16 plain bs32 | 0.989 | 1.001 | | fp16 plain bs64 | 1.011 | 1.004 | A second gfx942 part agrees within 1.3% on the five production kernels measured there (fp16 bs64 1.013 / 1.007, bf16 bs64 0.991, fp16 bs16 1.006, bf16 fold bs16 1.010). Pinned 4-warp tuning spec through the dispatcher, 1 round, gfx942: | Kernel | new / develop | |---|---| | Decode at the reported shape (B=64, Hq=64, Hkv=4, bs16) | 0.785 | | Decode B=32 bs32 / B=64 bs64 | 0.864 / 0.943 | | Causal prefill, bs16 | 0.853 | | Other pinned variants (fold, plain, GQA-8, block size 64) | 0.947 – 1.014 | Decode likely gains because the clamped Q loads all hit the sequence's one token instead of streaming 127 extra rows (no hardware counters were collected). --------- Co-authored-by: Claude <noreply@anthropic.com>

…lpers (#14244) ## Summary C++ mirror of the shared dense attention spec (`library/kernels/common/attention_dense_spec.py`) and of the dense-attention LSE helpers in `_lse_store.py`. First step of [AICK-2711](https://amd-hub.atlassian.net/browse/AICK-2711) (gfx950 dense attention C++ mirror); the gfx950 spec layer and both builders follow in separate PRs. Nothing in the shipped provider calls the new C code yet; the gfx950 and gfx942 builder ports will. ## What's ported | Python | C | |---|---| | `AttentionDenseSpec` fields and defaults | `rocke_attention_dense_spec_t`, `rocke_attention_dense_spec_default()` (`include/rocke/instance_attention_dense_spec.h`) | | `DENSE_TILE_GEOMETRIES` | `ROCKE_DENSE_TILE_GEOMETRIES` | | `__post_init__` | `rocke_attention_dense_spec_validate()` | | `resolved_persist_decode` / `resolved_nonpersist_decode` | `rocke_attention_dense_spec_resolved_*_decode()` | | `kernel_name()` and its `_*_name_parts` hooks | `rocke_attention_dense_kernel_name()`, `rocke_attention_dense_name_parts_t`, `rocke_attention_dense_name_parts_base()` | | `check_dense_spec_preflight` | `rocke_attention_dense_spec_preflight()` | | `PERSIST_DECODES` / `NONPERSIST_DECODES` (name, tag, check) | one `{name, tag, check}` table per grid in `instances/common/attention_dense_spec.cpp` | | `dense_row_has_key`, `zero_keyless_row` | `rocke_dense_row_has_key`, `rocke_zero_keyless_row` in `include/rocke/instance_lse_store_internal.h`, next to develop's `rocke_lse_store` (#14616), which the dense path reuses | Shared helpers this PR adds to or reuses: `rocke_spec_set_reason` (now marks a cut reason with `...`) and `rocke_strbuf_append_py_repr` (Python `repr(str)`, now also used by the IR printer). **Not ported here:** - `attention_dense_cache_key` (Python launcher-cache key; no C caller). - `validate_dense_lse` / `validate_unified_lse` (host-side torch tensor checks; no C caller). - `kernels/common/attention_dense_kv_split.py` (#14367: `check_kv_segments`, `segment_tile_range`, `emit_segment_store`, `declare_segment_params`, `MAX_KV_SEGMENTS`) and the `kv_segments` field / `kvs{N}` token. These belong to the gfx950 (AICK-2711) and gfx942 (AICK-2710) builder PRs, whose parity pairs must cover `kv_segments > 1`. Recorded on both tickets. - The runtime-shape builder surface from #14059 (`runtime_shape`, `fastdiv_divisors`, `runtime_kernarg_fields` on the spec; `make_fastdiv` / `fastdiv_kernargs` in the decode module). It feeds kernel params and the persistent fast division, so it goes with the builder ports. ## Design Subclass hooks are passed in explicitly: - per-class answers (`supports_*`, both decode sets, type name) through `rocke_attention_dense_traits_t`; - the per-spec `_aligned_causal` override (gfx950: `right_bound != 0`) through a `lookahead` argument; - name parts through `rocke_attention_dense_name_parts_t`. Preflight runs the base checks on the spec it is given, then an optional arch validator (`rocke_attention_dense_validate_fn`) for the arch's own fields, as both Python subclasses do with `super().__post_init__()`. A decode a traits list names but the C table doesn't implement is rejected, and `name_parts_base` returns an error for it rather than the default order's empty tag. `kernel_name` returns `ROCKE_ERR_VALUE` instead of a truncated name. The accessors, `name_parts_base` and `kernel_name` require a spec that passed preflight; on one that didn't, they return 0 / NULL / `ROCKE_ERR_VALUE` instead of dividing by zero. Every `ValueError` reason matches the Python text byte for byte. The two paged geometries where Python raises `ZeroDivisionError` (`block_m < 32`, or more waves than `block_n` rows) return a fixed `"division by zero"`; CPython's text varies by version, so those are compared by verdict only. ## How it's pinned **Spec layer: the `rocke_attention_dense_spec` ctest.** The reasons, block orders and names don't all reach emitted IR, so byte identity can't see them. Both engines are pinned to `platform/tests/golden/attention_dense_spec_vectors.json`. `library/tests/test_attention_dense_spec_vectors.py` recomputes it with the current Python (also in `rocke_library_attention_cpu_pytest`) and owns `--regenerate`; the ctest embeds the same file and requires the same answers from C. The file holds: - every field default and the tile table; - 129 hand-picked base cases and 400 cases drawn with a fixed seed (2711) around the auto block-order thresholds. `test_cases_reach_every_reject_site` traces the Python and requires every reject the C port mirrors to be reached by some case, and each of the two paged divisions to raise `ZeroDivisionError`; - a `Gfx950AttentionDenseSpec` hook table (16 cases) and a `Gfx942AttentionDenseSpec` hook table (9 cases). `test_hook_tables_exercise_every_hook` reverts each hook a class overrides and requires a recorded answer to change. Changing one token of `attention_dense_spec.cpp` at a time (thresholds, decode preconditions, ignoring a traits list or `lookahead`, name tags, defaults, the tile table, preflight's base and arch halves, accessor guards) fails the ctest in every case tried. **LSE helpers: byte identity.** The `attention_dense_lse` parity pair (`library/tests/parity/attention_dense_lse_emit.{py,c}`, 8 configs) is GREEN in `check_byte_identity.py`, as are all 97 families. ## Testing - All 36 host ctests pass (build tree); `rocke_attention_dense_spec` and `rocke_library_attention_cpu_pytest` pass from the installed tree. - `test_attention_dense_spec_vectors.py`, `test_attention_dense_lse.py`, `test_lse_store_ir.py` pass. - `check_byte_identity.py`: GREEN, all families. - pre-commit clean. ## Land order #14059 landed first; the vectors are regenerated on top of it (only gfx950 hook-case name parts moved). Any later change to the shared spec or a subclass's hooks needs `python test_attention_dense_spec_vectors.py --regenerate`. ## Known gaps (pre-existing in Python, tracked separately) - Preflight checks Q/O in elements, but gfx950's ragged path addresses Q in bytes: [AICK-2876](https://amd-hub.atlassian.net/browse/AICK-2876). This PR mirrors the current Python check; the fix goes to Python first. - gfx950's `right_bound` has no upper bound, so `q + right_bound` can overflow i32: [AICK-2915](https://amd-hub.atlassian.net/browse/AICK-2915). gfx950's validator isn't in this PR; the fix lands with the gfx950 spec layer PR, Python first. [AICK-2711]: https://amd-hub.atlassian.net/browse/AICK-2711?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ [AICK-2876]: https://amd-hub.atlassian.net/browse/AICK-2876?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: Chris Millette <chris.millette@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

aosewski and others added 30 commits October 24, 2026 00:18

Merge remote-tracking branch 'origin/users/anarao/perf/gpu-only-base'…

46eb70d

… into users/anarao/perf/gpu-only-p1-decode

test(rocke/attention): cover head_size 192 in the declared unified co…

cf74cab

…verage Co-Authored-By: Claude <noreply@anthropic.com>

AnaghaRaoAMD mentioned this pull request Oct 31, 2026

feat(hip-kernel-provider): add gfx950 unified attention 2D/3D ingestor engines #14242

Merged

19 of 20 tasks

This was referenced Oct 31, 2026

perf(rocke): add opt-in gfx942 BF16 D128 K LDS double buffering and V register prefetch #13205

Merged

feat(rocke): runtime head counts for gfx950 dense attention #14059

Merged

ex-rzr mentioned this pull request Nov 2, 2026

feat(rocke): paged K/V for dense prefill attention on gfx942 and gfx950 #14552

Open

1 task done

aosewski mentioned this pull request Nov 2, 2026

feat(rocke): additive bias (fp32 attn_mask) for gfx950 dense attention (AICK-2554) #14904

Open

13 of 21 tasks

AviralGoelAMD mentioned this pull request Nov 2, 2026

feat(hipdnn): gfx942 unified 3D split-KV decode as a hipDNN ingestor engine #14259

Merged

AnaghaRaoAMD mentioned this pull request Nov 3, 2026

fix(rocke): bound Q and K/V reads in the gfx942 4-warp GQA D128 body #14929

Merged

ex-rzr mentioned this pull request Nov 3, 2026

fix(rocke): stop OOB reads of K/V and the block table in unified attention #14967

Open

1 task done

amd-khushbu mentioned this pull request Nov 3, 2026

feat(hip-kernel-provider): add gfx942 dense attention AOT ingestor engine #14908

Merged

AnaghaRaoAMD mentioned this pull request Nov 4, 2026

fix(rocke): name seqlen and tile when dense attention rejects a partial tile #14335

Merged

1 task done

Navigation Menu

Repository navigation

perf(rocke/attention): GPU-time optimizations for unified and dense forward attention on gfx942/gfx950 - #14367

perf(rocke/attention): GPU-time optimizations for unified and dense forward attention on gfx942/gfx950#14367
AnaghaRaoAMD merged 89 commits into
developfrom
users/anarao/perf/gpu-only-integ

AnaghaRaoAMD commented Oct 25, 2026 •

edited

Loading

Uh oh!

assistant-librarian Bot commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

AnaghaRaoAMD commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Optimizations

Results

Verification

Uh oh!

assistant-librarian Bot commented Nov 3, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

AnaghaRaoAMD commented Oct 25, 2026 •

edited

Loading