[dlangbe]

JIRA ID : AICK-2717

Motivation

#12446 made paged K/V caches past the i32 buffer range (> 0x7FFF0000 bytes) use i64 K/V addressing, which fixed silent zero reads of the cache's top 64 KiB. On gfx942 the tiled 2D i64 path turned out much slower than the i32 path. For bf16 D128 decode on exactly-2-GiB caches (B=256, page size 16; the decode shapes of Llama-3-8B with Skv=4096 and GQA 32/8, Llama-3-70B with Skv=4096 and GQA 64/8, and Qwen3-235B-A22B with Skv=8192 and GQA 64/4), develop's i64 path takes 1.41x to 1.43x as long as its i32 path (measured as described under Performance).

Cause, from the gfx942 ISA: with i64 addressing the K/V loader built a buffer resource for every copy (192 per kernel for these shapes). Each copy reloaded its page id from the block table, waited for it (v_readfirstlane needs the value), and rebuilt the resource. Because vmcnt drains in order, each of those waits also drained the K/V copies issued before it, so the K/V stream ran with little memory-level parallelism. A 64-token tile spans only 4 pages.

Technical Details

Changes

• kernels/gfx942/attention_tiled_2d.py: when i64 addressing is on, a tile holds whole pages (T % BS == 0), and no copy crosses a page ((BS * HD) % KV_HALVES_PER_CALL == 0), the K and V loaders load the tile's page ids once (_tile_pages: one masked block-table load and one readfirstlane per page) and build one resource per page. Each copy takes its page's resource and an offset within the page (_page_voff). Other cases, the fast paged descriptor and the i32 path are unchanged.
• C++ twin (attention_tiled_2d_kv_descriptors_loads.cpp): the same loader, with builder calls in Python's evaluation order. The comments that said i64 never reaches the non-fast loader are corrected.
• _enable_i64_kv_addr docstring: drops a stale measured-cost claim.

Behaviour change

• gfx942 tiled 2D with i64 K/V addressing: same results, fewer page-table loads and resources (for the decode spec: 12 page loads and 12 extra resources per kernel instead of 192 and 193).
• i32 paths, gfx950 and all other kernels: emitted IR unchanged (no existing golden hash moves).
• Also closes a C++-only gap: the C++ twin's non-fast K/V loaders had no i64 path and would have emitted i32 addressing for these specs; no emit config covered it before cfg7.

Coverage

• Parity (gfx942_attention_tiled_2d_emit.{py,c}), byte-identical at llvm20/22/23. Develop's C++ loader drifts on exactly cfg7-10:
  • cfg7: the selector's spec for the i64 decode shape (1 warp, 4 pages per tile).
  • cfg8: i64, 4 warps, 2 pages per tile.
  • cfg9: i64, 1 page per tile.
  • cfg10: cfg7 with K copied by global_load_lds.
• Golden: new representative-IR case attention/gfx942/2d_bf16_d128_b16_decode_i64 (additive diff only).
• tests/test_attention_gfx942_kv_pages.py:
  • CPU structural check, over 1 warp with 4 pages per tile, 4 warps with 2 pages, and 1 page per tile: the i64 kernel issues one page-table load and one resource per page of each tile load, and the same number of copies as i32. All three fail on the per-copy loader. It builds the spec directly and checks counts only; the numeric cases cover addressing.
  • gfx942 numeric, against an fp32 reference with atol=2e-3, rtol=2**-8 (the relative term covers bf16 rounding of large outputs on short sequences):
    • i64 decode, B=8, Skv=1024, GQA 32/8 and 64/4 (num_kv_blocks overstated to force i64). Fails if the in-page offset drops the kv-head term.
    • Ragged lengths 17 to 1000 with a block table 63 pages wide, so the last 64-token tile reads past every row; the table's tail holds an unmapped page id. Removing the block-table bound faults the test.
    • A cache just over 2 GiB per tensor with the sequences' pages straddling 2^31 bytes (skips under about 5 GiB free). Sign-extending the page base, or falling back to the old whole-cache i32 loader, fails it while the small-cache cases still pass.
  • Enrolled in the installed library GPU CTest list and docs/gpu-attention-test-coverage.md. The gfx942 CI lane has no torch, so the numeric cases skip there and only the structural check runs.

Not in this PR

• The K-sliced-ring slice loader still builds a buffer resource per copy under i64 addressing. It is reachable for caches between 0x7FFF0000 and 2^31 bytes, where its conflict-free V loader is still in range; the C++ twin rejects the ring.
• Page ids are bounded by index (block_table_max_idx) but not checked against the cache's block count before they form a 64-bit base, the same as the i32 path and the fast descriptor. Bounding them needs the block count as a kernel argument, which is an ABI change.
• Each K or V tile load still waits once on its page-table loads before readfirstlane; the per-copy waits are gone.
• gfx950 tiled 2D has the same per-copy pattern (2 pages per tile, 8 copies). Follow-up, needs gfx950 hardware to measure.
• The i32 path also loads a page id per copy; moving it to per-page resources changes IR for every gfx942 2D paged spec and is a separate PR.

Performance

Relative results only. gfx942, ROCm 10.0.0 wheels (torch 2.13.0+rocm10.0.0, LLVM 23 from the ROCm SDK, the attention suite's default toolchain), one idle GPU. Each arm: 5 warmup launches through run_unified_attention_torch, then the median of 5 batches of 50 launches timed with HIP events. Correctness gate in the same run: max_abs against an fp32 PyTorch reference on sampled sequences, tolerance 1e-3; every arm passed with matching max_abs (2.8e-4 to 3.4e-4). Page tables were tested contiguous and random; the ratios below hold for both.

Decode on an exactly-2-GiB cache (B=256, D128, bf16, page size 16), branch on the i64 path:

Shape	develop i64 vs develop i32	branch i64 vs develop i64	branch i64 vs develop i32
Skv 4096, GQA 32/8	1.43x as long	2.15x faster	1.50x faster
Skv 4096, GQA 64/8	1.42x as long	2.11x faster	1.49x faster
Skv 8192, GQA 64/4	1.41x as long	2.11x faster	1.49x faster

The third column is a near-equivalent comparison: the i32 arm uses a cache two pages smaller, the largest i32 can address, with the same kernel otherwise.

Other shapes on the i64 path (forced by overstating num_kv_blocks, median of 5 batches of 20 launches, same correctness gate), branch vs develop: prefill B=4 S=1024 GQA 32/8 1.21x faster; prefill B=1 S=4096 GQA 64/8 1.30x faster; decode B=64 Skv=8192 GQA 64/8 2.07x faster. These are diagnostic configurations, not shipped cache sizes.

Test Plan

• ROCKE_LLVM_FLAVOR=llvm23 python rocke/platform/tools/check_byte_identity.py: GREEN.
• python rocke/tools/run_checks.py --steps gate,parity (llvm20 + llvm22): GREEN.
• Mutation: without the C++ change, gfx942_attention_tiled_2d cfg7 drifts.
• rocke_ir_parity_harness.py --check: OK before the new case; after re-bless the golden diff only adds the new case and count bumps.
• tests/test_attention_gfx942_kv_pages.py on gfx942 with ROCm 10.0.0 wheels (LLVM 23) and torch 2.10.0+rocm7.0 (LLVM 20): 3 passed. Structural test fails on the old loader; numeric test fails with a broken in-page offset.
• library/tests on CPU, one pytest per file with the GPU hidden (backward attention excluded): 169 files pass; the 6 failing files (test_strided_kv_decode, test_strided_kv_lds, test_kda_hostpack, test_conv_abi, test_conv_dgrad_patch, and the gfx950 use_additive_bias case in test_selector_parity) fail the same way on develop.
• platform/tests/test_rocke_ci_static.py: 5 passed.
• gfx942 GPU performance and correctness as in Performance above.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

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

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

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

🙋 Wish to Override Policy?

[therock-pr-bot]

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

[cgmillette]

Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max

PR #14349 — one K/V resource per page on the gfx942 2D i64 path — 🟢 First-round: no closer look needed beyond AnaghaRaoAMD's two open threads — no other kernel's IR moves, both engines agree on the new loader, and taking out either half of the change makes its own check fail

Prior item	Status
AnaghaRaoAMD (10-24) attention_tiled_2d.py:2688: both conditions of PAGE_RSRC_I64 always hold, so assert them instead	🟡 OPEN, no reply yet. Our check agrees for every spec the predicate admits. supports() allows head size 64 / 128 / 256, block size 16 / 32 / 64 and at most 8 warps, which gives at most 1,024 halves per call against a page of at least 1,024 halves, all powers of two (read). The assert would also turn a hand-built spec outside that envelope into a clear error rather than the unhandled i64 case in C++
AnaghaRaoAMD (10-24) test_attention_gfx942_kv_pages.py:159: atol=1e-2 is too loose to catch a partial addressing bug	🟡 OPEN, no reply yet. We have no gfx942 here to measure it, and CI skips this numeric test (below)
Bots / CI	🟢 Green at head 9c32c354: TheRock Multi-Arch, ASAN, the hipDNN superbuild, pre-commit and clang-tidy. In the gfx94X and gfx110X shards, the new structural test passes; both numeric cases skip because the lanes have no torch, as the description says (read)
Our checks	🟢 No other kernel's IR moves: the differential sweep over every anchored family (1,256 configs, llvm20) is identical to develop's (ran). Byte identity is GREEN for all 87 families at llvm22, and for gfx942_attention_tiled_2d at llvm20 and llvm23 too (ran). The native build and 30 ctests pass (ran). The structural test passes with both engines in strict mode (ran). Negative controls (ran): with develop's Python loader, the structural test fails (192 == 4 * 3). With develop's C++ twin, byte identity for gfx942_attention_tiled_2d goes RED on one of its 8 configs. The library CPU tests that import the 2D kernel show no new failure against develop (558 passed) (ran). The test is in the installed GPU CTest list (read). rule_scan is clean (ran)
Merge	🟢 GitHub reports a conflict because develop moved, but only the representative golden conflicts, and the change is additive. Regenerating it on the merge gives develop's entries plus the one new 2d_bf16_d128_b16_decode_i64 case, with no drift, and the golden gate passes at llvm20 / 22 / 23 (ran)

No net-new findings.

Utility: 🟢 Positive — the i64 K/V path on gfx942 2D no longer rebuilds a buffer resource and reloads a page id for every copy. It changes nothing on the i32 path or for any other kernel.

---

Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted

First-round review: a local qwen3.6 model made the independent first pass. The orchestrator (Claude Opus 5.5) verified it. The checks were a regenerated golden for the merge and the native build and ctests. They also included byte identity for every family, a differential sweep of every anchored family against develop, and the new test with both engines in strict mode. Negative controls with develop's Python loader and develop's C++ twin, and the library CPU tests against develop, were run as well. Finally, they covered the CI logs at head. The orchestrator owns this merge-worthiness call. It flags whether a closer look is needed; it is not an approval or a rejection.

[amd-khushbu]

Review checklist pass (rocKE PR-review checklist, two independent passes)

Ran this PR through the team's rocKE PR-review checklist (byte-identity/two-engine parity, coverage/fallback behavior, DRY/blast-radius, dispatcher/spec integration, structure/hygiene, merge-readiness, verification honesty, paged/indirection safety, reusable design principles), then did a second independent pass focused on the addressing math and test coverage. No blockers in either pass. The core change is sound: addressing math is correct (page selection, in-page byte offsets, i64 base arithmetic all check out against the pre-existing stride/guard logic), the C++ twin's emit order matches Python exactly (verified op-by-op), and this is correctly scoped as an emitter-only fix (no selection/routing logic touched).

Worth calling out explicitly: this also fixes a latent C++-only correctness gap

Before this PR, the C++ twin's non-fast paged-KV branch had no i64-addressing code path at all — it always built the plain i32 voffset regardless of I64_KV_ADDR (the removed comment said as much: "I64_KV_ADDR is false in gfx942's buildable space, so the plain i32 voffset is used"). Python's equivalent branch, by contrast, already correctly dispatched to offset_i64_split for i64 specs before this PR. So if the C++-built kernel were ever exercised for a cache large enough to need i64 addressing on this non-fast path (exactly the shape the new cfg7 parity case adds, which didn't exist before this PR), it would have silently used i32 addressing — reintroducing the silent-zero-reads bug #12446 fixed, but only in the C++ engine. "Same results" undersells this a bit; might be worth a line in the PR description saying this also closes a latent C++-only gap (reachability of that exact shape via the shipped C++ engine prior to this PR wasn't established either way, for what it's worth).

MEDIUM

• The K-sliced-ring loader (use_k_sliced_ring, attention_tiled_2d.py:2839-2848) still rebuilds one buffer resource per copy via offset_i64_split under i64 addressing — the same per-copy pattern this PR fixes elsewhere. Nothing in __post_init__ makes use_k_sliced_ring and use_i64_kv_addr mutually exclusive. The PR's "Not in this PR" list names the FP8 loaders, but those are actually unreachable on gfx942 (rejected explicitly) — the K-sliced-ring loader is the sibling that's actually reachable and isn't named.
• The block-table OOB guard and a ragged (non-tile-aligned) last tile are never exercised — the structural and numeric fixtures both use a seqlen that's an exact multiple of tile_size. Deleting the cmp_lt mask, or an off-by-one in first + page, would still pass both tests. (confirmed independently in both review passes)
• No committed test actually reads memory past the 2 GiB boundary — the GPU numeric test's cache is ~17 MB, with num_kv_blocks overstated only to flip the i64 flag. Every real offset in the test stays under 2^31, so a wrong multiply width or an incorrectly-sized resource bound would still pass.
• Coverage is num_warps=1 / 4-pages-per-tile / decode-only. Production also reaches other i64 geometries via _select_2d_num_warps (prefill gets 2 or 4 warps — the PR's own prefill perf numbers ran on this), plus N_BLOCKS_PER_TILE==1 and USE_GLOBAL_LOAD_LDS_K with i64 — none of these have a parity config, golden case, or numeric test, so C++/Python parity for these geometries is unverified.
• The structural test (test_i64_copies_share_one_resource_per_page) is a pure intrinsic-count check; it would not catch a bug that preserves counts but uses the wrong page/offset (e.g. a swapped page index or bad stride in _page_voff). It's also the only new test that runs in CI (the numeric test is GPU-gated and skips on the gfx942 CI lane, which has no torch).
• The new GPU test was added to the shared rocke_library_gpu_pytest CTest entry (~30 files, one TIMEOUT 1800) without re-timing the whole entry on hardware after the addition.
• No VGPR/scratch check (llvm-readelf --notes or rocprofv3) was run despite the change now holding up to 4 live buffer resources per tile instead of 1 — a plausible register-pressure increase that a speedup alone doesn't rule out.
• The root-cause explanation (vmcnt in-order draining serializing the K/V stream) is derived from reading ISA/instruction-ordering semantics rather than a dynamic hardware counter.

LOW

• attention_tiled_2d.py:337-339 (the use_i64_kv_addr field's own docstring) still says "it costs a per-call make_buffer_rsrc" — the PR updated the sibling _enable_i64_kv_addr docstring in attention_unified.py but missed this one.
• The block-table physical id is index-guarded (block_table_max_idx) but not value-range-checked against the real num_kv_blocks before folding into the 64-bit pointer base. A padded/garbage table entry inside the valid range could point the resource base at unmapped memory. This is identical to the pattern in the code it replaces (the removed offset_i64_split call, and _fast_paged_kv_voff's i64 branch) — not new or worsened by this PR, but cheap to harden now since page ids are loaded once per page instead of once per copy.
• The two engines fail differently on an out-of-range page count: C++ has a new ROCKE__MAX_TILE_PAGES = 16 cap, Python has none. Both reviews independently computed the real ceiling of n_blocks_per_tile under supports_tiled_2d's LDS gate and got exactly 16 — provably unreachable today, but worth mirroring the cap on the Python side for clarity.
• A vmcnt drain still happens once per tile load (the page-id load + readfirstlane at the head of each K/V tile load) — the serialization this PR targets is reduced, not eliminated. Possible follow-up, not a defect.
• The n_pages > ROCKE__MAX_TILE_PAGES C++ rejection doesn't follow the codebase's NOT_YET_IMPLEMENTED: convention used elsewhere for implementation-shortcut rejections.
• test_i64_copies_share_one_resource_per_page enters via dataclasses.replace(_DECODE, use_i64_kv_addr=True), bypassing the real selector. Defensible as a structural check (the companion numeric test independently confirms reachability through the selector), but the docstring doesn't say so.
• The "12 loads/resources instead of 192/193" comparison isn't reproducible by any committed test, since the old per-copy loader was deleted. Inherent to any refactor that removes the old code, not unique to this PR.
• The C++-built kernel for this path is never executed on a GPU in this PR's own tests, only byte-identity-diffed (as IR text) against Python's build, which is the one validated on silicon. Matches the repo's existing two-engine test strategy, not unique to this PR.

Open questions (not findings)

• Which develop commit/worktree was the "develop i64 vs develop i32" baseline measured against?
• Is there a live production C++ dispatcher (outside this diff) that actually routes real inference traffic through this exact path, or is C++ validated only via byte-identity today?
• Was the (i64 addressing + non-fast descriptor) combination ever actually reachable via the shipped C++ engine prior to this PR, or was it purely latent/untested until now?

[cgmillette]

Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max

PR #14349 — one K/V resource per page on the gfx942 2D i64 path — 🟢 Follow-up: no closer look needed — both of AnaghaRaoAMD's threads are fixed, the refactor leaves every parity config's IR unchanged, and the per-page resources cost 5 SGPRs while saving 22 VGPRs; one 🟡 doc line

Prior item	Status
AnaghaRaoAMD (10-24) thread 1: assert the two invariants instead of branching on them	🟢 FIXED in 9670fa88. The gate is now I64_KV_ADDR and not FAST_PAGED_KV_DESC, and the unreachable i64 fallback is gone from both engines' K and V loaders. A hand-built spec outside the envelope now gets a ValueError; the C++ twin raises a build error with the same message. At the old head, head size 96 with 8 warps and i64 addressing still built through the removed path; now it is refused with the new message (ran). All 8 gfx942 tiled-2D parity configs emit identical IR at the old head and on the merge, at llvm20 / 22 / 23 (ran)
AnaghaRaoAMD (10-24) thread 2: atol=1e-2 is too loose	🟢 FIXED: the tolerance is now atol=2e-3 (read). The author measured 6.4e-4 to 8.9e-4 on gfx942; we have no gfx942 to repeat it
amd-khushbu (10-26) checklist pass: no blockers; MEDIUM items on coverage and on register pressure	🟡 OPEN; no reply yet. The coverage items: the ragged last tile and the OOB mask are not exercised, no test reads past 2 GiB, several i64 geometries are untested, and the K-sliced-ring loader still builds a resource per copy. Register pressure, measured (ran): the i64 decode kernel (cfg7), compiled for gfx942 with ROCm clang 24, uses 160 VGPRs and 63 SGPRs with no scratch or spills. Develop uses 182 and 58, so the per-page resources save 22 VGPRs and cost 5 SGPRs. The i32 kernel is unchanged (182 / 56)
Bots / CI	🟡 TheRock Multi-Arch has been queued since 16:01 UTC. ASAN, the hipDNN superbuild (which ran before today's develop break), pre-commit, clang-tidy and Component CI are green (read)
Our checks	🟢 The merge with develop is clean (22 behind). The native build and 32 ctests pass (ran). Byte identity is GREEN for all 85 families at llvm22 and for gfx942_attention_tiled_2d at llvm20 / 23 (ran). The golden gate passes ×3, with its new i64 decode case unchanged (ran). The structural test passes with both engines in strict mode (ran). rule_scan is clean (ran)

#	Sev	file:line	Finding
N1	🟡	library/kernels/gfx942/attention_tiled_2d.py:337	The use_i64_kv_addr field comment still says i64 addressing "costs a per-call make_buffer_rsrc". The sibling _enable_i64_kv_addr docstring now describes one resource per page → update the field comment to match. (Worker's catch, verified by reading.)

Utility: 🟢 Positive — the i64 K/V path on gfx942 2D builds one resource per page instead of one per copy. With fewer VGPRs it is cheaper, and the unreachable fallback is now a clear error in both engines.

---

Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted

Follow-up review: a local qwen3.6 model made the independent first pass. The orchestrator (Claude Opus 5.5) verified it. The checks were the IR of every parity config at the old head and on the merge, at three flavors, and out-of-envelope specs at both heads. They also included register and scratch use from gfx942 code objects against develop, the native build and ctests, byte identity, the golden gate, the structural test with both engines and the CI state. The orchestrator owns this merge-worthiness call. It flags whether a closer look is needed; it is not an approval or a rejection.

[dlangbe]

@amd-khushbu
Thanks for the thorough pass. Changes in the latest push, then the rest.

Fixed

• Field comment at attention_tiled_2d.py:337 now says one resource per page (also the bot's N1).
• C++ page cap: removed rather than mirrored. The page pointer/resource arrays are arena-allocated per tile, so both engines accept the same specs and there is no NOT_YET_IMPLEMENTED limit left.
• Ragged tiles and the block-table guard: new test_i64_ragged_lengths_and_block_table_tail (lengths 17..1000, block table 63 pages wide so the last 64-token tile reads past every row). The table is a view whose tail holds 0x7FFFFFFF, so reading it unmasked faults: removing the cmp_lt bound aborts the test with a GPU memory fault.
• Real >2 GiB addresses: new test_i64_decode_reads_past_2gib uses a cache of 2 GiB + 4 MiB per tensor with the sequences' pages straddling 2^31 bytes (skips under ~5 GiB free). Sign-extending the page base, or falling back to the old whole-cache i32 loader, fails it while the small-cache tests still pass. (An i32 multiply that is zero-extended is still correct below 4 GiB, so that particular mutation is not a bug at this size.)
• Other geometries: emit configs cfg8 (4 warps, 2 pages/tile), cfg9 (1 page/tile) and cfg10 (i64 + use_global_load_lds_k), byte-identical at llvm20/22/23; develop's C++ drifts on exactly cfg7-10. The structural test now runs over the 1-warp/4-page, 4-warp/2-page and 1-page geometries, and fails on the per-copy loader for all three.
• Structural test docstring says it builds the spec directly and only counts; the numeric tests cover addressing through the selector.
• Tolerance is now atol=2e-3, rtol=2**-8: the 17-token sequence has |o| up to 1.8, where the bf16 rounding of the output alone is 3.7e-3. Misaddressing errors I measured are 0.09-0.2.

Measured, no change

• Registers (gfx942, ROCm 10, LLVM 23): i64 decode spec 159 VGPR / 63 SGPR vs develop's 180 / 58; 4-warp i64 239 / 96 vs 237 / 95; no spills or scratch; i32 unchanged.
• CTest: the file's 7 tests take about 11 s and pass inside the full library GPU list. On my gfx942 node the full list does not complete on either develop or this branch: the process dies in test_mla_prefill_gfx942_numeric.py::test_mla_prefill_numeric_dispatched[square_h64], which also happens with that file alone on develop. That is unrelated and I'll raise it separately.

Follow-ups, not in this PR (added to the description)

• The K-sliced-ring slice loader still builds a resource per copy under i64 (reachable between 0x7FFF0000 and 2^31 bytes; the C++ twin rejects the ring).
• Page-id value range: not checked against num_kv_blocks, same as the i32 path and the fast descriptor. Bounding it needs the block count as a kernarg, which is an ABI change.
• One page-table load and readfirstlane per tile load remain; the per-copy waits are gone.
• Root cause is from the ISA (per-copy s_waitcnt vmcnt before each readfirstlane) plus the measured speedup; I did not capture hardware counters.

Open questions

• Baseline: develop at 7a1c04c9ccdc for both develop columns.
• C++ at run time: the gfx942 tiled 2D kernel is built in Python and lowered from its serialized IR; the provider has no native gfx942 tiled 2D pack, so the C++ builder is only checked through byte identity.
• Reachability before this PR: the non-fast i64 combination was reachable in Python (the selector sets i64 for caches past the buffer range), so the C++ builder's missing i64 path was a parity gap with no config covering it, not a live runtime bug. Agreed it's worth saying, and the description now does.

[Copilot]

🔵 Needs a closer look

Low-level GPU addressing changes need human confirmation of the reported numerical and Python/C++ parity results, especially because numerical tests skip in the stated CI lane.

No open issues 🎉

What changed in this PR

Optimizes rocKE’s gfx942 tiled 2D attention for large paged K/V caches by reusing one buffer resource per page instead of rebuilding it per copy, following the i64 addressing fix in #12446.

Changes:

• Implements per-page i64 K/V loading in Python and C++, including previously missing C++ i64 handling.
• Adds parity configurations, additive IR goldens, and structural and GPU numerical tests.
• Registers the tests with CTest and updates addressing and coverage documentation.

Paths below are relative to dnn-providers/hip-kernel-provider/rocke/.

File	Description
platform/​tests/​instances/​rocke_ir_parity_harness.py	Adds a representative i64 decode case.
platform/​tests/​golden/​rocke_representative_ir_sha256.json	Adds hashes for the new case across three LLVM versions.
platform/​cpp/​instances/​gfx942/​attention_tiled_2d_kv_descriptors_loads.cpp	Implements matching C++ per-page resources and offsets.
platform/​CMakeLists.txt	Registers the new test file with installed GPU tests.
library/​tests/​test_attention_gfx942_kv_pages.py	Tests resource counts, ragged sequences, and large-cache addressing.
library/​tests/​parity/​gfx942_attention_tiled_2d_emit.py	Adds four Python i64 parity configurations.
library/​tests/​parity/​gfx942_attention_tiled_2d_emit.c	Adds corresponding C parity configurations.
library/​kernels/​gfx942/​attention_tiled_2d.py	Reuses page resources in the i64 K/V loaders.
library/​kernels/​common/​attention_unified.py	Updates the addressing docstring and removes a stale performance claim.
docs/​gpu-attention-test-coverage.md	Documents the new numerical cases and CI coverage limitations.

_{Review effort: Balanced}

[CongMa13]

No new issues found. I agree with the following existing comments:

• Make the whole-page i64 assumptions explicit and reject unsupported specs — AnaghaRaoAMD — #14349 (comment)
• Tighten the numeric test tolerance — AnaghaRaoAMD — #14349 (comment)
• Cover ragged table tails and real large addresses, and add coverage for alternate i64 geometries — amd-khushbu — #14349 (comment)
• Call out that GPU numeric tests skip in the torchless CI lane — copilot-pull-request-reviewer[bot] — #14349 (review)
• Update the use_i64_kv_addr field comment to describe per-page resources — cgmillette — #14349 (comment)
• Track the K-sliced-ring i64 per-copy resource path as follow-up work — amd-khushbu — #14349 (comment)

[cgmillette]

Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max

PR #14349 — one K/V resource per page on the gfx942 2D i64 path — 🟢 Follow-up 2: no closer look needed — every open item is fixed or listed under "Not in this PR". The new geometries are byte-identical ×3, the structural test fails on develop's per-copy loader in all three geometries, and no geometry spills. One 🟡 description line

Prior item	Status
Ours N1 🟡 (10-26): the use_i64_kv_addr field comment still described a per-call resource	🟢 FIXED in c72e80b2 (read)
amd-khushbu (10-26) MEDIUM: the ragged last tile and the block-table guard are never exercised	🟢 FIXED: new test_i64_ragged_lengths_and_block_table_tail (read). We checked its arithmetic. The table is 63 pages wide, so the last 64-token tile of a 1,000-token row reads page index 63. The guard spans the whole table (num_seqs × bt_stride = 504 entries), so it masks the last row's read at index 504, where the tail holds 0x7FFFFFFF. Rows 0–6 over-read the next row's valid ids, which the token mask discards (read). GPU-only; we have no gfx942 to run it
amd-khushbu MEDIUM: no test reads past 2 GiB	🟢 FIXED: new test_i64_decode_reads_past_2gib (read). Pages are 32 KiB and each tensor has 65,664 of them (2 GiB + 4 MiB). The used pages, 65,408–65,663, straddle page 65,536, which starts at 2^31 bytes. num_kv_blocks is the true count, so the selector picks i64 by itself (read). GPU-only
amd-khushbu MEDIUM: other i64 geometries have no parity config	🟢 FIXED: cfg8 (4 warps, 2 pages per tile), cfg9 (1 page) and cfg10 (use_global_load_lds_k). gfx942_attention_tiled_2d byte identity is GREEN for all 11 configs at llvm20 / 22 / 23 (ran). Given the same specs, develop's engine emits identical IR for cfg0–6 and different IR for exactly cfg7–10, at all three flavors (ran, Python engine)
amd-khushbu MEDIUM: the structural test counts only, for one geometry	🟢 FIXED: it now covers three geometries and passes on the merge. Run against develop's per-copy loader, it fails in all three (192, 48 and 192 page loads where 12, 6 and 3 are expected) (ran, negative control). It still checks counts only; the GPU tests cover addressing
amd-khushbu MEDIUM: register pressure unmeasured	🟢 Measured (ran; ROCm clang 24, gfx942, -O3). i64 VGPR / SGPR, merge vs develop: • cfg7: 160 / 63 vs 182 / 58 • cfg8: 241 / 100 vs 238 / 99 • cfg9: 205 / 57 vs 204 / 57 • cfg10: 176 / 63 vs 254 / 79 No config uses scratch or spills, and the i32 kernels are unchanged. The author's LLVM 23 figures point the same way
amd-khushbu MEDIUM + LOW, the rest	🟢 Each is addressed (read): • Recorded under "Not in this PR": the K-sliced ring still builds a resource per copy; the page-id value range is unchecked, and checking it needs an ABI change; one page-table wait per tile load remains; the root cause comes from ISA reading. • C++ page cap removed: pages are now arena-allocated per tile. The native build is clean (ran). • Docstring: now says the structural test builds the spec directly. • CTest timing: the author measured the file at about 11 s within the GPU list. They report that the full list dies in test_mla_prefill_gfx942_numeric.py on develop as well; that is a develop issue, which they will raise
AnaghaRaoAMD (04:55) and amd-khushbu (06:19) APPROVED at 866daebb; Copilot (05:15) "needs a closer look" because the numerics skip in CI; CongMa13 (15:06) COMMENTED at 866daebb: "no new issues", agreeing with the existing items	🟢 Approvals are at head, and every item CongMa13 lists is addressed above (read). Copilot's point is confirmed in the next row
Bots / CI	🟡 All 46 checks pass at head; none fail or are pending (read). The gfx94X lane does collect the file through the GPU CTest list. Its 3 structural cases pass, but the 4 numeric cases skip with No module named 'torch' (ran, CI logs). So far only the author's node has run the numerics, the same torch gap seen on other PRs
Our checks	🟢 • The merge with develop is clean (20 behind) (ran). • Commits 1–3 changed only in rebase context (range-diff). • cfg0–7 emit identical IR at the old head and on the merge at llvm20 / 22 / 23, so commit 4 changes no IR (ran). • Out-of-envelope specs still raise the same ValueError (ran). • The native build and 33 ctests pass; rocke_arch_domain_reproducibility skips under clang 24 as usual (ran). • Byte identity is GREEN for every gfx942 family at llvm22 (ran). • The golden gate passes ×3 (ran). • rule_scan is clean (ran)

#	Sev	file:line	Finding
N1	🟡	PR description, Test Plan	The kv-pages line still says "3 passed", a count from before the four new cases. The file now has 7 cases (3 structural, 4 numeric), and the author's 10-27 reply says all 7 pass → refresh the line

Utility: 🟢 Positive — the i64 K/V path on gfx942 2D builds one resource per page. It needs fewer VGPRs on the decode and global_load_lds variants, both engines agree on all 11 configs, and the new GPU tests cover the ragged block-table tail and real addresses past 2 GiB.

---

Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted

Follow-up review: a local qwen3.6 model made the independent first pass. The orchestrator (Claude Opus 5.5) verified it. The checks were:

• a range-diff against our last review;
• IR of every parity config at the old head, on the merge and on develop, at three flavors;
• the structural test on the merge and against develop's loader;
• register and scratch use from gfx942 code objects against develop;
• the native build and ctests, byte identity, the golden gate, and the CI logs at head.

The orchestrator owns this merge-worthiness call. It flags whether a closer look is needed; it is not an approval or a rejection.

[CongMa13]

No new P0 or P1 issues were identified during this review. The changes are approved.

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-10-28 — Included in TheRock via ROCm/TheRock#9541.