[dpeabody-amd]

Motivation

attention_dense selects its LDS pitch from the DMA row-packing factor:

ROWS_PER_INSTR = 128 // D          # 1 at D=128, 2 at D=64
LDROW = D + PAD          if ROWS_PER_INSTR == 1 else D
VROW  = D + _LDS_PAD_V   if ROWS_PER_INSTR == 1 else D

At D=64 both pads collapse to zero, so neither LDS bank-conflict lever is active
on that head size. The module docstring already recorded this as known — "its
reads take some LDS bank conflict (V-pad is a future D=64 perf lever)" — because
D=64 was added for coverage under the constraint that it cost the D=128 schedule
nothing.

Why it matters, in bank terms. On gfx950 bank = (byte_addr / 4) % 64 over 64
banks. A wave64 ds_read_b128 moves 1024 B while LDS delivers 64 banks x 4 B =
256 B/cycle, so 4 lanes per bank is the conflict-free floor, not merely
"better". Aggregating all 64 lanes of one QK read (n=8 bf16 = 4 dwords each,
krow = nsub*32 + lane%32, col = ks*16 + (lane//32)*8):

K LDS layout	banks touched	lanes per bank
D=128 unpadded [BN, 128]	8 / 64	32
D=128 shipped [BN, 136]	64 / 64	4 (floor)
D=64 shipped [BN, 64]	16 / 64	16
D=64 group-padded [BN/2, 128+8]	64 / 64	4 (floor)

So D=64 today sits far closer to the unpadded D=128 case than to the shipped one.

Technical Details

A per-row pad is not possible at D=64: one async_buffer_load_lds moves 64 lanes
x 2 bf16 = 128 contiguous elements, which is exactly one D=128 row but two D=64
rows, so a padded row pitch would place the second row at the wrong offset. The
pad therefore goes between DMA row-groups — the group stays contiguous for the
DMA, while the group pitch restores the bank spread.

• New lds_k_group_pad spec field (default 8; 0 reproduces the old unpadded
  layout for A/B). __post_init__ rejects a non-multiple-of-8 pad, because
  smem_load_vN stamps align 16 on the n=8 read unconditionally, so an
  8-byte-aligned pitch would keep the ds_read_b128 while breaking its
  alignment contract silently. It also joins kernel_name() on the packed path
  — it changes the emitted layout, so it has to be part of the kernel identity
  or two layouts share a symbol name and a launcher-cache entry. The
  head_size=128 name is unchanged.
• K LDS becomes
  [NBUF, BN // ROWS_PER_INSTR, ROWS_PER_INSTR * D + spec.lds_k_group_pad]
  when ROWS_PER_INSTR > 1; unchanged [NBUF, BN, D + PAD] otherwise.
• _async_load strides its destination by group rather than by row. The sixth
  parameter is renamed ldrow_bytes -> group_bytes; the two are equal when
  ROWS_PER_INSTR == 1. V passes the plain unpadded ROWS_PER_INSTR * D pitch,
  so its layout is bit-identically addressed (V_GROUP_BYTES * grp reduces to the
  previous VROW * 2 * row0).
• do_qk maps krow to (krow // K_GROUP, (krow % K_GROUP) * D + col). This is
  exact because nsub*32 is a multiple of K_GROUP, so the group/sub-row split
  commutes with the nsub offset. Both terms are hoisted above the KV loop.
• Mirrored in the default and the persistent builder, with an assert that the
  group split divides evenly so a future head size fails loudly.
• Why 8 and not 16. The whole-wave model above cannot separate 8, 16 and 24 —
  all three reach the 4-deep floor. A 16-lane phase (16 x 4 dwords = one full
  64-bank sweep) does separate them: pad 8 gives 16 distinct start banks, pad 16
  repeats each twice. A sweep over {0, 8, 16, 24, 32} confirms 8 and 16 are
  indistinguishable in practice while 0 and 32 are clearly worse (32 falls back to
  32 banks at 8 deep). 8 is therefore chosen as the model-optimal and cheaper of
  the two equivalents.
• New library/tests/test_attention_dense_d64_lds.py (CPU-only, 63 cases): the
  pad is inert at D=128 (identical IR and identical LDS, both grid variants),
  live at D=64, grows LDS by exactly one pad per K row-group across block_n
  32/64/128 on both builders, and a misaligned pad is rejected. Every invariant
  runs on both allowed dtypes: fp16 takes the same packed group-pad path as
  bf16, and fp16 at D=64 is already a recorded golden case, so bf16 alone would
  not establish it.
• Also fixed, in the same file: _DENSE_LAUNCHER_CACHE is keyed on
  kernel_name(), which omits batch — but batch is baked into the kernel
  (buffer extents, persistent work-item count). Two specs differing only in batch
  collided, so the second silently reused the first binary: B=1 then B=2 produced
  wrong output with only 1/B of the tensor written, and B=2 then B=1 raised
  hipError(700) illegal memory access. The key is now (kernel_name(), batch);
  keying here rather than widening kernel_name() leaves the emitted IR
  untouched. Pre-existing, reproduces on unpatched develop.
• Exposed on the dense prefill benchmark
  (library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py)
  as --lds-k-group-pad, forwarded at both spec sites, so the pad can be swept
  directly rather than being reachable only from a unit test. Each sweep row
  now records the tuning knobs it was measured at, so the JSON/CSV report is
  self-describing instead of relying on kernel_name() to carry them.
• The same cache-key bug exists in that script, and is fixed here too. Its
  private _LAUNCHER_CACHE keyed on kernel_name() alone, so it collided on
  batch and waves_per_eu. That one fires on the shipped default run
  rather than only in principle: the varlen sweep drives [2048] * 4 (B=4) and
  [256, 512, 768, 1024, 1536, 2048] (B=6) at the same max_s, which produce
  an identical kernel_name() but different lowered IR, so on --mode all the
  B=6 config silently reused the B=4 binary. The key is now
  (kernel_name(), batch, waves_per_eu). Pre-existing, reproduces on unpatched
  develop.

V is deliberately left alone. The same 128-element DMA granularity blocks a
per-row V pad, and a group-padded V would need a pad-aware pv32_v_load_paired
— SSA-emitting platform/ helper code, which pulls in the C++ mirror and a
byte-identity re-bless. ROCKE_DENSE_VPAD is currently inert at D=64 (VROW is
64 at both 0 and 32), so conflict-free transposed-V remains an untapped D=64
lever, tracked as follow-up rather than bundled here.

Risk Assessment

• D=128 codegen does not change, by construction. Every new operation sits
  behind ROWS_PER_INSTR > 1, and the builder is side-effecting, so a guarded
  branch emits nothing. At ROWS_PER_INSTR == 1 the derived values
  (K_GROUP=1, K_ROWS_LDS=BN, LDROW=D+PAD, K_GROUP_BYTES=LDROW*2, V_GROUP_BYTES=1*VROW*2) are each identical to the prior expression. Verified
  empirically: every D=128 shape the kernel serves in the tracked cohort lowers
  to identical LLVM IR before and after, under both llvm20 and llvm22;
  only the D=64 cases move.

• Numerically neutral at D=64. A 19-variant matrix (causal / non-causal,
  bf16 / fp16, MHA and GQA 4/8/16, persistent and default grids, block_n
  32/64/128, sliding window, ragged, lazy-rescale off, waves_per_eu=1) gives
  max_abs against an fp32 oracle identical to the pre-change baseline to every
  digit. varlen is not in that matrix: it forbids persistent and shares the
  default builder's K path verbatim, so it is covered structurally rather than
  numerically.

• Static resources at D=64, block_n=64 (llvm-readelf on the HSACO, no
  launch — throttle-independent):

  builder	pad	VGPR	AGPR	SGPR	spill	LDS	waves/CU	limiter
  default	0	141	0	46	0	32768 B	12	VGPR
  default	8	145	0	46	0	33792 B	12	VGPR
  persistent	0	197	0	46	0	32768 B	8	VGPR
  persistent	8	197	0	46	0	33792 B	8	VGPR

  Zero spill throughout. The persistent builder's register count is unchanged
  (the group/sub-row terms are loop-invariant and hoisted); the default builder
  takes +4 VGPR with no occupancy change. LDS grows 1 KB per workgroup against
  160 KB/CU, and the kernel is VGPR-limited rather than LDS-limited at D=64, so
  the extra kilobyte does not move waves/CU.

• Python-only; no C++ mirror required. This is a library/kernels/ layout
  change built from existing IRBuilder primitives — no new IR op, no
  platform/ helper or emitter touched, so the C++ lowerer needs no change.
  The perf(rocke/attention): recover fp16 gfx942 D128 prefill via depth-2 sliced-K ring #9662 / feat(rocke/attention): enable gfx942 bf16 D128 prefill ring for fp16 parity #9198 precedent that these attention kernels are not registered
  in the platform byte-identity gate no longer holds: develop registered them
  after this branch was cut. At this PR's base the representative golden held
  zero attention_dense entries across two llvm flavors; it now holds 165
  across three, having gained llvm23. The D=64 cases are re-blessed here
  accordingly (see Test Result); the C++ lowerer still needs no change.

• Blast radius is bounded to an opt-in path. dispatch/attention.py matches
  the attention_dense candidate only on an explicit
  algorithm="attention_dense", never under auto.

• No interaction with D=256. __post_init__ admits only head sizes 64 and
  128; D=256 is a different kernel family.

ASIC Coverage

gfx950 only. attention_dense is gfx950-only by construction
(supports_attention_dense rejects every other arch), and the change is further
gated behind ROWS_PER_INSTR > 1, i.e. head size 64. Validated on a physical
gfx950 node.

Test Plan

• Byte-identity A/B over every shape the dense kernel serves in the tracked
  cohort, pristine vs patched, for llvm20 and llvm22.
• D=64 numeric vs an fp32 oracle across the 19-variant matrix, baseline vs
  patched, compared digit-for-digit.
• K group-pad sweep {0, 8, 16, 24, 32}, forward and reverse (ABBA) to cancel
  ordering drift, plus a higher-rep 8-vs-16 tiebreak.
• Geometry A/B at D=64: block_n {32, 64, 128} x waves_per_eu {1, 2, 3} x
  persistent {on, off}.
• LDS bank-conflict counters (SQ_LDS_BANK_CONFLICT, SQ_LDS_IDX_ACTIVE) at the
  shipped vs padded layout.
• Static occupancy probe (llvm-readelf) at D=64 for both builders, per pad.
• D=128 null control: the same binary timed with the pad knob at 0 and 8, to
  bound same-session A/B repeatability.
• test_attention_dense_d64_lds.py, on bf16 and fp16; black --check;
  ruff check (no --fix —
  the emitter-code autofix hazard platform/AGENTS.md describes applies equally
  to library kernels, and this file has pre-existing F841s whose assignments
  have op-emission side effects).

Test Result

• Byte-identity: PASS. All D=128 cohort shapes identical on both llvm
  flavors; only the D=64 cases changed; no build errors.
• Numeric: PASS, 19/19, every max_abs identical to the pre-change baseline.
• LDS bank conflicts drop substantially at the padded layout. Note the
  whole-wave model predicts a 4x reduction and the measured drop is larger; the
  counters are whole-dispatch and cover both the K ds_read_b128 and the
  still-unpadded transposed-V reads, so the residual conflict is consistent with
  V — the lever tracked as follow-up.
• D=64 throughput improves materially at the tracked configurations,
  monotonically with sequence length. Quantitative results are in the protected
  internal page (AICK-1828).
• Same-session repeatability, from the D=128 null control where the pad knob
  is inert, is a few tenths of a percent. Absolute throughput on this part swings
  far more across sessions, so only same-session ratios were used.
• Static resources: zero spill, occupancy unchanged (table above).
• Geometry: no change warranted. The shipped bn=64, waves_per_eu=2, persistent stays optimal. block_n=128 is now feasible at D=64 (it is
  VGPR-blocked at D=128) but does not pay for its doubled LDS.
• New test: 63/63 pass. The 30 dtype-dependent cases now run on both bf16
  and fp16; the 3 spec-validation cases stay bf16, since that path never reads
  the dtype. black --check clean. ruff check reports the same
  pre-existing findings as an unpatched build and none from this change.
• Benchmark wiring: verified as far as a CPU-only environment allows. The
  script is GPU-only, so neither CI nor this environment can run a sweep. Checked
  by importing it against a stubbed torch: --lds-k-group-pad is present in
  --help, and both the success and failure record paths carry the knob fields.
  The varlen cache collision was confirmed directly rather than inferred, by
  building the B=4 and B=6 specs and comparing them: identical kernel_name(),
  different lowered-IR SHA.
• Platform golden re-blessed for D=64, and validated before writing. The
  CI-wired platform/tests/golden/rocke_representative_ir_sha256.json now
  covers attention_dense, so the intended D=64 IR change trips it. Exactly
  three cases move -- fp16_h64_sq512, persist_h64_sq512,
  ragged_h64_sq500 -- on each of llvm20, llvm22 and llvm23, for 27
  changed fields and nothing else. To confirm the blessing environment was
  trustworthy rather than merely convenient, the harness was first run against
  unpatched develop, where it reproduced the committed golden with zero drift
  on all three flavors; each re-blessed hash then matched the value CI
  reported independently.
• Every D=128 attention_dense entry in that gate is unchanged, on all
  three flavors. That is an independent confirmation of the D=128 byte-identity
  claim above, and it extends it to llvm23, which was not available when the
  original A/B was run.
• Not run: test_attention_dense_golden.py::test_attention_dense_cpp_python_byte_identity
  requires a built rocke_engine, which was unavailable in this environment. No
  new IR op is introduced, so no C++ change is expected; a reviewer with the
  engine built should confirm.

Notes: pre-existing, not introduced here

library/tests/golden/attention_dense_ir_sha256.json does not match a pristine
build of develop for either llvm flavor (0 of 18 cases), which is consistent
with the fixture being deliberately not CI-wired — its own docstring says so, and
drift is therefore unpoliced. This PR does not re-bless it: doing so would
mix an unrelated whole-fixture refresh into a D=64 change, and would land on an
already-red baseline. The byte-identity evidence above is the pristine-vs-patched
diff rather than the committed fixture. Refreshing the fixture (and deciding
whether to CI-wire it) is better handled on its own.

Submission Checklist

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

Changes are behind the feature flag HIPKERNELPROVIDER_ENABLE_ROCKE which is defaulted to OFF.

JIRA ID : AICK-1828

[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.

[yraparti]

The D=64 approach looks solid. Padding between DMA row-groups preserves the packed async-LDS load path while restoring K read bank spreading, and the alignment validation is particularly important for the ds_read_b128 path. A few things I'd like to see for confirmation:

Perf data comparing kpad={0,8,16,24,32} to justify the default value of 8.
Coverage demonstrating no regressions in persistent+D64 variants.
Follow-up evaluation of launcher cache uniqueness. I agree with @AviralGoelAMD comment that waves_per_eu appears codegen-affecting and may need to participate in the cache key as well.

Otherwise the design direction looks correct and maintains D=128 codegen stability while addressing the D=64 banking issue

[dpeabody-amd]

@yraparti thanks, taking your three in order.