[amd-khushbu]

JIRA ID : AICK-2712

Summary

This is PR 5 of a stack of 9 that brings the rocKE C++ unified-attention engine to parity with Python (see the stack below). It sits on PR 4.

The gfx942 fp16/bf16 flash prefill specs Python selects use the transposed-32x32 VALU stack (mask-limit, invariant hoist) and the sliced-K ring. The C++ builder rejected both, after PR 1 made it stop silently ignoring them. This PR ports them, plus the causal mask phase split and HBM-direct K/V, so C++ builds the production gfx942 flash specs byte-identically.

Changes

All paths are under dnn-providers/hip-kernel-provider/rocke/platform/cpp unless noted.

• Mask-limit and invariant hoist.
  • The per-lane query-row invariants (and the ALiBi slope) are hoisted before the KV loop.
  • The mask-limit threshold compare honours skip_mask.
  • The ALiBi slope is picked as Python picks it; plain transposed + ALiBi previously failed to build.
  • The first parity configs to combine the production flash stack exposed other gaps on this path, fixed here: sinks seeding under transposed scalar state, the K single-buffer and split-cfvst branches, and C argument-order drift at four call sites.
• Sliced-K ring (depth 2 or 3).
  • The transposed QK streams K in k_slice_hd-wide slices through ring_depth LDS slots, with ring_depth - 1 slices in flight behind a partial vmcnt fence.
  • When a prefetch targets the slot the previous slice was read from, the slot is drained before reuse. The depth-3 ring uses the CK LdsSeq slot maps.
  • Slice geometries Python rejects are rejected the same way.
• Causal mask phase split, V/K from HBM. use_causal_mask_phase_split, use_v_hbm_direct and use_k_hbm_direct are added to the shared spec, defaulted off; gfx950 ignores them. Their name tokens (cmps, vhbm, khbm) match Python.
• Gate: the unified tiled-2D gate no longer rejects the gfx942 flash mask-limit and sliced-K ring cohorts.
• Stopgap: 64-bit KV offsets with the sliced-K ring.
  • Develop (perf(rocke): build one K/V resource per page on the gfx942 2D i64 path #14349) ported the gfx942 2D per-page i64 K/V loads. The ring's K slice loader still builds one i32 offset, so the C++ IR would differ from Python's and overflow at run time.
  • Until that is ported, the builder rejects use_i64_kv_addr with an active ring (NOTIMPL), and the 2D gate rejects those problems as "not ported to C++ yet". The gap is recorded in known_gaps.md §9.
  • A sweep of every i64 problem Python's 2D gate admits, run at the top of the stack: 810 specs, of which 558 are identical and the 252 ring specs are rejected.

Stack

1. feat(hip-kernel-provider): bring the rocKE C++ unified-attention engine to parity with Python #14522 Stopgaps: stop C++ silently misbuilding specs
2. feat(hip-kernel-provider): port rocKE unified-attention selectors, routing and support gates to C++ #14523 Selectors, routing and support gates
3. feat(hip-kernel-provider): assemble rocKE unified-attention specs from a problem in C++ and run selector parity in CTest #14524 Spec from problem (gfx942 2D, 3D), scalar fallback, selector parity in CTest
4. feat(hip-kernel-provider): port the rocKE gfx950 tiled-2D V double-buffer, D256 fast path and softmax interleave to C++ #14525 gfx950 tiled 2D: V double-buffer, D256 fast path, softmax interleave
5. feat(hip-kernel-provider): port the rocKE gfx942 tiled-2D mask-limit, sliced-K ring, phase split and HBM-direct K/V to C++ #14526 gfx942 tiled 2D: mask-limit, sliced-K ring, phase split, HBM-direct K/V ← this PR
6. feat(hip-kernel-provider): port the rocKE gfx942 4-warp GQA attention builder to C++ #14527 gfx942 4-warp GQA builder
7. fix(hip-kernel-provider): rocKE unified-attention C++ parity follow-ups #14528 Follow-ups: gfx942 3D i64 gate, reject messages, emit_lse, CTest labels
8. feat(hip-kernel-provider): close the remaining rocKE unified-attention C++ builder gaps #14758 i64 sliced-K ring, shared arch gate, 3D kv_layout, LDS resolver, gfx950 Q re-read and depth-3 ring, parity coverage
9. feat(hip-kernel-provider): port the rocKE unified-attention mask band and split-KV LSE to C++ #14759 Mask band (top-left / right-bound) and split-KV LSE

Each PR is based on the one before it. Review from 1 upward; each diff shows only its own slice.

Test plan

• gfx942 tiled-2D parity emitter configs 11–30 (develop's i64 configs stay 7–10):
  • 11–16: the production D64 fp16 and D128 bf16 flash specs, per-iteration mask-limit with partial rows, and hoist with ALiBi / QQ bias / sliding window.
  • 17–21: the production D64 fp16 ring, the depth-2 ring, both LdsSeq maps, iglp, and the ring with softcap and a sliding window.
  • 22–30: each phase-split / HBM-direct knob on the paths it changes.
• test_attention_routing_cpp_parity.py and attention_unified_routing_test.cpp: the newly admitted cohorts, and the i64 + ring reject reason
• test_attention_binding_spec_fields.py: the new fields are read rather than rejected
• End-to-end check against Python, run at the top of the stack, from the C++ gate through the C++ spec to the C++ build: 696 mask-limit / ring specs, 0 mismatches, 0 build failures
• CTest, pybind build: 39/39
• Byte-identity gate GREEN at llvm20, llvm22, llvm23 and llvm24
• Golden: gfx942_attention_tiled_2d 11–30 blessed; no other attention entry moves
• pre-commit's black (25.12.0) and clang-format (18) hooks, run by hand on every file this PR changes: clean (pre-commit itself is not installed locally)
• GPU numeric lanes not run for this PR; the byte-identity gate is the parity check

🤖 Generated with Claude Code

[CongMa13]

Two actionable findings are left inline.

[CongMa13]

[P2] K-direct mode aliases Q into the one-element K LDS stub

With use_k_hbm_direct=true and use_q_direct_global=false, Q_ALIAS_K can still become true because its size check uses the full logical K buffer size. This branch allocates K_lds as [K_BUFS, 1, 1], so the Q gather then indexes the one-element stub as [BLOCK_M, HD]. The added parity configurations 29 and 30 both exercise this combination. Please keep Q from aliasing the K stub, either by forcing direct Q loads or allocating a full Q LDS buffer.

[CongMa13]

[P2] Direct HBM readers do not honor i64 KV addressing

This direct-K path obtains one descriptor offset and passes it to buffer_load_vN, but does not use the i64 split-offset path. The direct-V helper does the same before its global_load. A spec can currently combine either direct-HBM flag with use_i64_kv_addr=true; for caches beyond the i32 offset range these loads can address the wrong pages. Please reject those combinations until both readers handle split 64-bit offsets.

[cgmillette]

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

PR #14526 — port the rocKE gfx942 tiled-2D mask-limit, sliced-K ring, phase split and HBM-direct K/V to C++ (amd-khushbu, AICK-2712) — 🟠 First-round: not merge-ready — rebase needed — conflicts with develop (its own golden hunk is a plain union); the khbm Q-alias bug is ported from Python

Prior item	Status
CongMa13 (10-28) ctx_prologue.cpp:535 [P2]: with use_k_hbm_direct and no direct Q, Q aliases the one-element K LDS stub (cfg 29, 30)	🟠 OPEN, valid — no reply. Real, and a faithful port of the same Python code: attention_tiled_2d.py:1740 sizes Q_ALIAS_K against the full K_TOTAL_BYTES, :1748 allocates the [K_BUFS, 1, 1] stub, and :2094 points Q_lds at it, as ctx_prologue.cpp:522, :539 and :707 do. The head's Python emitter reproduces the blessed llvm20 hashes of cfg 29 and 30. In that IR each Q row is stored as <8 x bfloat> through getelementptr [1 x [1 x [1 x bfloat]]] …, i32 0, i32 0, row, col (cfg 30: [2 x [1 x [1 x half]]]). So (row, col) collapses to element row + col, and Q rows overwrite each other before the Q gather reads them back the same way (ran: probe). Byte identity passes only because both engines emit this IR. Python's tuner already holds gfx942 use_k_hbm_direct out as a known-wrong knob ("produced wrong outputs", dispatch/attention/axes.py:423); this alias is a likely cause. No selector sets the knob and tuning holds it out, so only hand-built specs, opt-in probes and these two configs reach it. Neither #14527 nor #14528 touches it (read). → Fix it in Python and mirror it: force direct Q under khbm (Q_DIRECT_GLOBAL = K_SLICED_ACTIVE or use_q_direct_global or use_k_hbm_direct at attention_tiled_2d.py:1739 and ctx_prologue.cpp:521, as the ring already does for its smaller K slab; no LDS needed), then re-bless cfg 29/30. If the Python fix lands separately, set use_q_direct_global=True on cfg 29/30 meanwhile so the golden does not bless a corrupt kernel
CongMa13 (10-28) kv_body_qk_softmax.cpp:246 [P2]: the direct-HBM K and V readers ignore use_i64_kv_addr; reject those combinations	🟡 OPEN, mirrors Python's design — no reply. Both readers build one i32 offset (:235, kv_descriptors_loads.cpp:258), but so do their Python twins (_read_k8_mfma_operand, _v_hbm_elem). Python lists both flags as i32-only loaders under i64 (_I32_ONLY_KV_LOADERS, attention_unified.py:3055). _spec_kv_addr_limit caps an i64 khbm spec at 0x7FFF0000 bytes and a vhbm spec at 2^31 bytes, and _check_kv_addr_width raises at launch instead of reading wrong pages (:4264); cfv, cfvst and fp8-PV work the same way. So here C++ emits Python's IR. That is unlike the ring, where the IRs would differ, and a C++-only reject would refuse specs Python builds identically. #14527 and #14528 leave it as is (read). → Answer the thread with the launch cap. If a build-time reject is wanted, add it to Python's validator first and mirror it
Our #14523 (10-29) N2: the gfx942 transposed-x8 score reads the ALiBi slope only under mask-once	🟢 FIXED here fb4e2463 — the score now takes the hoisted, per-iteration or inline per-lane slope, branch for branch with Python (kv_body_qk_softmax.cpp:693), and cfg 14–16 cover the three branches (read)
Bots / CI	🟡 • At fb4e2463, TheRock Multi-Arch CI (run on the stack, not on develop) fails only Test hipkernelprovider, on all four shards. Linux gfx94X 1/2 and 2/2 and Windows gfx110X 1/2 and 2/2 fail rocke_differential_golden_ll and _canonical on the same 13 moe_fused_mega_fp8/target_intrinsics configs (4 in canonical IR) plus 1 new one. Develop carried exactly these at the stack base 0e1cfd3c and re-blessed them in #14501. None of this PR's 20 new gfx942_attention_tiled_2d entries fails at llvm20/22/23 or in canonical IR. • Windows 1/2 also fails the gfx1100–1103 census entries of TestGfx1151WmmaAttentionDecodePacks (descriptor set did not load); develop fixed that in #14553. Windows 2/2 also fails rocke_library_gdn_pytest: test_gdn_decode_fused.py::TestFusedEmission::test_lowers_and_compiles_without_scratch hits pytest's per-test timeout during a COMGR compile. It passes on the other three shards and timed out the same way on another branch's Windows run, and this PR touches no GDN code, so it is a Windows timing flake. • hipDNN superbuild, both ASAN runs, clang-tidy and Component CI pass. Develop's own Multi-Arch run at 9d729d32 passes, so a rebase gives a fresh run. GitHub's "mergeable" is against the stacked base #14525, not develop. • CI exercises none of the C++ port. Its differential golden runs --python-only, the pybind engine is never built (the PR's three Python test files need it), and the native rocke_attention_unified_routing ctest has no tier label, so it is absent from all four shard logs (#14528 adds the labels) (read: CI logs)
Our checks	🟠 • git merge-tree with develop 050acb2a (56 behind, 5 ahead): 3 conflicted files, 7 hunks. attention_unified_selectors.cpp (2 hunks) comes from #14523 vs #14150. prove_parity_binding.py (1 hunk) comes from #14524 (f9edddc0), which deletes the gfx950 selector-parity matrix that #14150 extends. Both resolve in those PRs. llvm_gfx_all.json (4 hunks, one per section) is this PR's own: develop's #14501 and the stack's #14522 blessed gfx942_attention_tiled_2d 7–10 with identical hashes, and this PR's keys 11–19 sort right after "10", so both sides insert at the same spot (ran: merge-tree). → Keep both sides: no hash disagrees, and 20–30 auto-merge. Once the stack rebases onto develop, #14522's copy of 7–10 drops out and the hunk applies cleanly. With no merged tree, the check script skipped the build, goldens, byte identity, tests and compile. • Drift: since the base, develop changed no gfx942 tiled-2D Python or C++ (#14349's i64 path is already in the base; #14150 is gfx950). Of this PR's files it touched only the golden (#14501's re-bless, #14415's llvm24 → llvm23 alias). Develop's CTest now also checks llvm24. With develop's Python, configs 11–30 emit the blessed llvm23 hashes at both llvm23 and llvm24 (20/20 each), so the new entries hold after the rebase (ran: probe). • Emitters: the head's Python emitter reproduces the blessed llvm20 hashes of cfg 11, 29 and 30, and the C and Python emitters set the same fields for configs 11–30 (ran: probe; static compare). We did not run C++ byte identity; the author reports it GREEN at llvm20/22/23. The PR's three Python test files need the pybind engine, and without it every case skips. • Parity (read): the port follows Python branch for branch and op for op in:   – the mask-limit threshold (row_ok folded into a −2^30 threshold, row offsets pre-subtracted, skip_mask honoured) and the per-lane invariant hoist;   – sinks seeding under transposed scalar state, the K single-buffer drain, and the split-cfvst V load/store placement;   – the sliced-K ring: slices 0 and 1 in the prologue, the drain-on-reuse (lgkmcnt(0) plus barrier), the partial vmcnt of one slice's calls, the full drain on the last slice, LdsSeq maps only at depth 3, and Python's two slice-geometry errors;   – the phase split (split clamped to the loop bounds, u-suffixed phase-1 phis, off with a sliding window or grouped-kv2);   – the direct-HBM readers, the khbm validator rules and the vhbm/khbm/cmps name tokens. This lifts #14522's mask-limit, hoist and ring NOTIMPLs. • i64 + ring stopgap (read): complete. The NOTIMPL sits in the builder's only entry, under K_SLICED_ACTIVE (ctx_prologue.cpp:472). The gate's test (attention_unified_gates.cpp:163) is exactly what the spec builder sets: the flash ring predicate and enable_i64_kv_addr. HBM-direct with i64 is not a C++-only gap (thread 2). • Rules 1–3 clean (0 tracker, 0 perf, 0 env); no links or tracker IDs added (ran)

Net-new: none.

Utility: 🟢 Positive — C++ now builds gfx942's production flash prefill specs; untested in CI.

---

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

First-round review: a local qwen3.6 model made the independent first pass, and the orchestrator (Claude Opus 5.5) verified it. The PR conflicts with develop, so there is no merged tree yet; the build, byte identity, goldens and tests wait for the rebase. The orchestrator simulated the merge of every stack PR with git merge-tree and compared the differential golden across the stack and develop. It read the ported mask-limit, hoist, ring, phase-split and HBM-direct code against the Python twins. It emitted the PR's parity configs with the Python emitter at llvm20, llvm23 and llvm24, and inspected the LDS addressing of configs 29 and 30. It checked #14527 and #14528 for fixes and read the four failing CI logs. The orchestrator owns this merge-worthiness call. This review flags whether a closer look is needed; it is not an approval or a rejection.

[cgmillette]

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

PR #14526 — port the rocKE gfx942 tiled-2D mask-limit, sliced-K ring, phase split and HBM-direct K/V to C++ (amd-khushbu, AICK-2712) — 🟠 Follow-up: rebase fixed the conflicts; the khbm Q-alias is still open

Prior item	Status
CongMa13 (10-28) ctx_prologue.cpp:535 [P2]: with use_k_hbm_direct and no direct Q, Q aliases the one-element K LDS stub (cfg 29, 30)	🟠 OPEN, valid — unchanged, no reply. The commit is the same patch as at fb4e2463: Q_DIRECT_GLOBAL is still set only for the ring (ctx_prologue.cpp:521), Q_ALIAS_K still compares Q against the full K_TOTAL_BYTES (:467, :522), and khbm still allocates the [K_BUFS, 1, 1] stub (:535). On the merge, cfg 29 and 30 lower byte-identically in both engines, and both store Q and gather it back through getelementptr [K_BUFS x [1 x [1 x T]]] (56 and 28 such GEPs). Row and column indices therefore add up to one small offset, and Q rows overwrite each other. With use_q_direct_global=True, both engines emit none of these GEPs (ran: probe). Python's twin (attention_tiled_2d.py:1739) is identical on develop, and the tuner still holds gfx942 use_k_hbm_direct out as a known-wrong knob (axes.py:424). No selector sets the knob, so only hand-built specs reach it, but the golden still blesses both kernels (read) → fix it in Python and mirror it (Q_DIRECT_GLOBAL = K_SLICED_ACTIVE or use_q_direct_global or use_k_hbm_direct), then re-bless cfg 29/30; until then, set use_q_direct_global=True on cfg 29/30
CongMa13 (10-28) kv_body_qk_softmax.cpp:246 [P2]: the direct-HBM K and V readers ignore use_i64_kv_addr; reject those combinations	🟡 OPEN, mirrors Python's design — unchanged, no reply. Both readers still form one i32 offset (kv_body_qk_softmax.cpp:246, kv_descriptors_loads.cpp:269), as their Python twins do. On the merge, four hand-built i64 specs (cfg 29 with khbm and vhbm, cfg 30 with khbm, cfg 26 and 27 with vhbm) lower byte-identically in both engines. Python caps them at launch: _spec_kv_addr_limit allows 0x7FFF0000 bytes under khbm and 2^31 under vhbm (attention_unified.py:3056), and _check_kv_addr_width raises past that (ran: probe). Neither engine's spec-from-problem sets either flag (C++ defaults at attention_tiled_2d.cpp:134), and only the pybind binding and the parity tests call the C++ builder (read). So C++ builds exactly what Python builds here. The one i64 combination whose IR would differ, the ring, is refused: cfg 17 with i64 builds in Python, and the C++ builder raises the stopgap NOTIMPL (ctx_prologue.cpp:479) (ran) → answer the thread with the launch cap; a build-time reject belongs in Python's validator first, then mirrored
Our #14523 (10-29) N2: the gfx942 transposed-x8 score read the ALiBi slope only under mask-once	🟢 FIXED fb4e2463, unchanged — the score takes the hoisted, per-iteration or inline slope as Python does (kv_body_qk_softmax.cpp:693). cfg 14–16 cover the three branches and are byte-identical at llvm20/22/23 (ran: byte identity)
Our 10-29 (13:23Z) header: rebase needed — 3 files, 7 hunks (attention_unified_selectors.cpp from #14523, prove_parity_binding.py from #14524, and this PR's golden)	🟢 FIXED by the stack's rebase onto develop 4a3fc208 (18:19Z). This PR now sits on #14525's current head 8c8e4429, and its commit has the same patch-id as fb4e2463. The merge with develop df91d568 (24 behind, 5 ahead) has no conflicts: the rebased #14522 no longer re-adds golden cfg 7–10, and this PR's 11–30 merge in cleanly. None of the 24 develop commits since the stack's base touches the unified-attention kernels (ran: merge, range-diff)
Bots / CI	🟢 • At head f87527fa: TheRock Multi-Arch CI passes after one re-run. Its first attempt failed only the Windows gfx110X hipkernelprovider shard 2/2: rocke_pytest hit pytest's per-test timeout in tests/runtime/test_hip_device_properties_layout.py while Windows was starting a compile subprocess. This PR touches no runtime code, and the re-run of that shard passed. All four shards (Linux gfx94X 1/2 and 2/2, Windows gfx110X 1/2 and 2/2) now pass, and both Linux logs show rocke_golden_static and both differential goldens passing; gfx950 and gfx125X ran no test shard. • The 4 failed checks (the Multi-Arch CI, Multi-Arch ASAN, clang-tidy and hipDNN Superbuild summaries) and all 10 cancelled ones belong to five duplicate runs that were cancelled within 20 s of the 18:19Z push. Their twin runs pass: both hipDNN superbuilds, ASAN, clang-tidy and Component CI. • CI still exercises none of the C++ port. The differential golden runs --python-only, the pybind engine is never built (the PR's three Python test files need it), and the native routing ctest is absent from both Linux shard logs; #14528 adds the tier labels. • GitHub reports MERGEABLE (CLEAN) into #14525's branch, with auto-merge off and no review decision. No thread, reply or review has been posted since our first round (read: gh, job logs)
Our checks	🟢 • Range-diff from fb4e2463: #14522–#14524 were rebased, and #14525 and this PR's commit are unchanged. Every file that differs from our reviewed head, such as attention_unified.py and the two scalar-attention test files, comes from develop (#13889, #14150) or from the rebased lower PRs (ran). • Merge with develop df91d568 (0 conflicts): the native build is OK, and ctest passes 37/37 (1 skipped), the routing ctest with the i64 + ring reject included. Byte identity is GREEN at llvm20/22/23 for 9 families; gfx942_attention_tiled_2d covers all 31 configs, this PR's 11–30 included. The representative golden passes ×3. The differential golden is OK at llvm20–24 and in canonical IR, with 0 changed and 0 missing. This PR's 20 configs are blessed in all four sections, so the unblessed NEW configs belong to other families (ran). • Tests with the merge's rocke_engine: 165 passed, 1 skipped, 2 failed. Both fail on develop too: TestKvAddrLimit (#14651) and the graph-replay case, which needs torch (ran). • Compile (ROCm clang 24, -O3): no failures and no new rejections, and the 31 gfx942 2D configs do not spill. The one gfx950 2D spill is cfg 57, a lower-stack ALiBi config (28 B of scratch per lane at the 128-VGPR cap of waves_per_eu=4); develop's Python emits identical IR for it. gfx942 3D cfg 2 overflows LDS on develop too (ran). • Land order: after #14522–#14525. #14522 still needs an approval, and #14523 and #14525 have none. Three open PRs merge cleanly with develop but conflict with the stack: #14629 (cfvst i64 path) in 7 of this PR's files (conflict_free_v.cpp 6 hunks, kv_body_qk_softmax.cpp 4, both emitters, kv_body_pv.cpp, kv_descriptors_loads.cpp, the golden), the approved #14616 in 6, and the approved #14542 only in lower-stack files and the golden, though it changes the same Python kernel. Whichever lands second must re-port and re-run byte identity (ran: merge-tree). • rule_scan is clean, and no links or tracker IDs were added (ran)

Net-new: none.

Utility: 🟢 Positive — C++ builds gfx942's production flash specs byte-identically; untested in CI.

---

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

Follow-up review: a local qwen3.6 model made the independent first pass, and the orchestrator (Claude Opus 5.5) verified it. The checks started with a range-diff of the rebased stack against the head we reviewed. On the merge with develop df91d568, they covered the native build and ctest, byte identity at three flavors, both goldens, the PR's test files with the merge's rocke_engine, and a ROCm clang 24 compile. A probe lowered the specs behind both threads (khbm with Q aliasing, khbm and vhbm with i64, and the ring with i64) in both engines and checked Python's launch caps. Develop's Python IR for the spilling gfx950 config was compared with the merge's and compiled. Merge-trees were run against the open PRs that touch the same files. CI at head, including the log of the failed Windows shard, and the threads since our first round were read. The orchestrator owns this merge-worthiness call. This review flags whether a closer look is needed; it is not an approval or a rejection.

[amd-khushbu]

This stack is now combined into one PR, #14522, so it can be reviewed and merged as one change. GitHub shows this PR as merged because its commit is now on #14522's branch; nothing has gone into develop. This PR's change is one commit there, unchanged. Open review comments from here are listed in #14522's description and will be answered there.