Replace the per-MFMA V transpose with CDNA4's transposing LDS read

[demandal25]

Summary

The in-tree ROCm prefill kernel spends 35-41% of its wall clock on one in-register transpose that CDNA4 can do in a single instruction. Replacing it with ds_read_b64_tr_b16 is 1.71x on gfx950 (2.813 -> 1.648 ms at MHA 32/32, head_dim 128, seq 4096). gfx942 is untouched — its generated ISA is byte-identical to base.

This started as an attempt to widen the MFMA tile shape from 16x16x16. That premise did not survive measurement and was dropped; see Why not the MFMA shape below.

What changed

• include/flashinfer/rocm/mma.h — adds load_transposed_fragment, wrapping __builtin_amdgcn_ds_read_tr16_b64_v4f16, plus kHasTransposingLdsRead, the single constant everything else keys off so the seed and the load cannot drift apart.
• include/flashinfer/rocm/attention/permuted_smem.cuh — load_matrix_m16n16_trans picks the hardware read on CDNA4 and keeps transpose_mma_tile on CDNA3. trans_frag_row/trans_frag_col now live here, beside the consumer that requires them.
• include/flashinfer/rocm/attention/prefill.cuh — the three V read seeds (single, ragged, paged) and v_col_idx use those helpers instead of the open-coded lane%16 / lane/16.
• tests/rocm/test_batch_prefill_kernels.py — adds test_batch_prefill_matches_independent_reference (ragged and paged, head_dim 64/128/256, against naive_attention), and warms head_dim 64 in the JIT fixture.

Architecture / design notes

compute_qk reads K with a plain load_fragment; compute_sfm_v reads V through load_matrix_m16n16_trans, which appends transpose_mma_tile — six cross-lane shuffles, once per PV MFMA. Both phases issue the same 16,777,216 MFMA instructions, yet PV costs 9x more on gfx942 and 15x on gfx950. All six shuffles lower to ds_bpermute_b32 (2264 in the base kernel ISA; zero ds_swizzle_b32, zero v_mov_b32_dpp), so none of them are cheap.

The builtin is not a drop-in at the same address. It redistributes as out(16g+4a+b)[j] = in(16g+4j+a)[b], so at the load_fragment seed lane 0 receives {0,64,128,192} where the software path gives {0,16,32,48}. Solving that for the required seed gives row = 4*(lane/16) + (lane%16)/4, col = lane%4 — a different index expression into the same swizzled tile, so no LDS layout changes. Verified bit-identical to load_fragment + transpose_mma_tile across all 64 lanes before any kernel code was touched.

Because the two arches want different per-lane addresses, seeding the obvious way stays correct on CDNA3 and is silently wrong on CDNA4. That is why the seeds moved onto smem_t next to the load, and why one constant gates all three sites.

The read is 16-bit only. KernelTraits::DTypeKVSmem already guarantees that for every KV dtype — since #368 an fp8 cache is dequantized on the way into LDS and static_assert(sizeof(DTypeKVSmem) == 2) holds it there — so no fp8 guard is needed.

Scope. The S-fragment transpose_mma_tile in compute_sfm_v is unchanged and still issues ds_bpermute_b32; S lives in registers, never LDS, so the hardware read cannot reach it. This removes the V-side transposes only.

Why not the MFMA shape

The investigation began from a measurement that 16x16x16 runs at 27% of peak with 3.6x available. Re-measuring found the probe divided the fastest rep's time by a different rep's clock, and ran blocks that did not co-reside (its derived reference clock read 49-78 MHz against a true 100 MHz). Corrected: 16x16x16 sustains 46%, so the shape is worth 2.1x on gfx950 and ~1.1x on gfx942, where the two 16-bit shapes are equal per cycle (16 vs 32 cycles for 2x the flops). Against that, the kernel measured 4-5% of peak — roughly 10x below the ceiling of the instruction it already issues — so the shape was not the constraint.

Benchmark results

tmp/probes/prefill_standalone.hip, MHA 32/32, head_dim 128, seq 4096, non-causal, 20 iterations, three interleaved reps against base at this branch's merge base.

gfx950 / MI350X

build	mean	TFLOP/s
base	2.815 / 2.813 / 2.805 ms	97.7
this PR	1.648 / 1.648 / 1.659 ms	166.8
speedup	1.71x

gfx942 / MI300X — no change by construction. The device ISA is byte-identical to base (43515 lines each, diff reports no difference), so this is proven rather than measured-equal; an interleaved wall-clock A/B on a shared node agreed to within run-to-run noise.

Test plan

Backend is fa2 explicitly throughout — auto routes to AITER and would exercise none of this.

gfx950 / MI350X — the only arch where the new code path compiles:

• test_single_prefill_kernels.py -k fa2 — 1632 passed, 288 skipped (independent naive_attention oracle)
• test_batch_prefill_kernels.py -k independent_reference — 24 passed (new; ragged + paged, head_dim 64/128/256)
• in-tree fp8 (-k "fp8 and not aiter") — 162 passed, 67 skipped
• test_batch_prefill_bf16_custom_mask.py — bf16 KV through the transposing read
• test_logits_cap.py 330, POD/cascade/shared-prefix 622, batch prefill 3492, ragged-tail/sliding-window/block-sparse
• test_force_cta_tile_q.py — 6 passed, covering forced CTA_TILE_Q 16/64/128

gfx942 / MI300X — the CDNA4 branch is not compiled here, so a green says only that CDNA3 is unregressed; the ISA diff above is the real evidence.

• test_single_prefill_kernels.py -k fa2 — 1632 passed, 288 skipped
• ragged-tail subset — 144 passed
• new independent-reference test — passed

A/B on the new test — with the CDNA4 seed deliberately wrong the max abs error is 1.41-2.29 against 0.002 correct, so it detects a V-layout defect by ~150x and the 1e-2 tolerance sits ~5x above the fp16 noise floor.

• pre-commit run -a

[Copilot]

🔵 Needs a closer look

The architecture-specific LDS instruction and lane mapping require final validation on supported CDNA hardware.

Pull request overview

Optimizes ROCm prefill on CDNA4 by replacing V-fragment software transposition with a transposing LDS read while preserving the CDNA3 path.

Changes:

• Adds the gfx950 transposing LDS-read intrinsic wrapper.
• Updates V-fragment addressing and loading across all prefill variants.
• Adds independent-reference coverage for ragged and paged batch prefill.

File summaries

File	Description
include/flashinfer/rocm/mma.h	Adds architecture-gated transposing LDS load support.
include/flashinfer/rocm/attention/permuted_smem.cuh	Selects hardware or software transposition and centralizes address mapping.
include/flashinfer/rocm/attention/prefill.cuh	Applies architecture-aware V-fragment seeds to all prefill paths.
tests/rocm/test_batch_prefill_kernels.py	Adds FP32-reference tests and head-dimension 64 JIT warmup.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[demandal25]

Review 5217298933 generated 0 comments; its only note asks for hardware validation of the lane mapping. Validated on MI350X: ds_read_b64_tr_b16 under the new seed is bit-identical to load_fragment + transpose_mma_tile over all 64 lanes (0/256 mismatches), and the gfx950 suites are in the test plan. gfx942 keeps the software path — its device ISA is byte-identical to base, 43515 lines either way. Loop closed.