TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx1151/attention

gfx1151 WMMA FMHA-forward: an optimization case study

TL;DR (executive summary)

A native WMMA flash-attention-forward kernel for gfx1151 (Strix Halo, Radeon 8060S, RDNA3.5, wave32), tuned with the runbook loop across four kernel bodies and ~14 levers, every variant correctness-gated before timing.

Bottom line: the best this DSL achieves on this part is ~11 TF (~18% of the 59 TF f16 WMMA peak), from the single-wave gather kernels — fmha_pipelined (software-pipelined) at D64, fmha_singlewave fuse_k at D128. The kernel is issue-bound and register-bound, not FLOP- or bandwidth-bound (WMMA is ~1% of issued instructions; the 192-VGPR cap spills at D128). Every structural rewrite that the textbook CDNA design prescribes — multi-wave, cooperative LDS K/V staging, register-blocked tiles — loses here (fmha_multiwave ~6.3 TF, fmha_blockn ~8–9 TF, fmha_regblocked ~6 TF), because on this large-cache APU the single-wave global gather is already cache-resident and the cheapest way to feed the matrix unit. The one ~2× win is algorithmic: the causal early-exit.

vs PyTorch SDPA (MATH fallback — the only SDPA backend enabled on gfx1151; FLASH/EFFICIENT are runtime-disabled):

metric	this kernel	notes
correctness vs torch SDPA	max_abs 3e-5 (full D64/D128/GQA), 4.9e-4 (causal)	inside the 2e-2 f16-WMMA-accum gate
throughput vs SDPA (25-shape avg)	13.3×	SDPA MATH does the full quadratic
causal speedup vs SDPA	up to 29.5×	our causal early-exit skips masked tiles; SDPA MATH does not
best absolute	11.39 TF = 19.3% of peak	B2 GQA Sq1024 D128
25-shape avg	17.0% of peak	uniformly issue-bound; 0/25 reach 25%

Key takeaways surfaced from the body:

• Name the roofline first — the original "200 TF" target is ~3.4× the 59 TF physical peak; unreachable by any kernel.
• Count the matmul fraction — WMMA at ~1% of instructions ⇒ issue-bound; chasing FLOPs is futile until that ratio rises (needs gfx12 larger-K WMMA).
• CDNA's bandwidth-optimal LDS-staging design is pessimal here — the APU's cache feeds the matrix unit for free, so staging only adds barriers/ds_loads.
• Believe the measurement — ±0.6 TF thermal drift on this box; several 6% "wins" evaporated on back-to-back repeat. Decisions made A/B in one thermal window.

See the iteration ledger, Strix Halo nuances, gfx950 (CDNA) algorithmic differences, and CK Tile lessons sections for the durable takeaways.

New to flash attention? ALGORITHM.md derives the winning kernel from the math up — the attention spec, the online-softmax recurrence (with a correctness proof), the causal early-exit, and how every step maps onto what one wave32 does per K-loop iteration. Read it first if you want to understand what the kernel computes before reading how it was optimized.

---

A worked application of the optimization runbook to the native gfx1151 WMMA flash-attention forward kernel (kernels.gfx1151.wmma_fmha_fwd):

• Part 1 follows the runbook loop on a single lever (V-LDS staging) and records a result worth keeping precisely because it is the opposite of the intuition: the "optimization" is a regression, and the correct decision is to revert it.
• Part 2 is a multi-lever campaign (fmha_singlewave.py + tune.py, plus the fmha_pipelined/fmha_blockn/fmha_multiwave/fmha_regblocked siblings and the combo.py cartesian sweep) that pushes toward peak TFLOP/s, names the hardware roofline honestly (the ~59 TF f16 WMMA peak — a "200 TFLOP/s" target is physically unreachable on this part), and diagnoses why the kernel is stuck at ~17% of peak: it is issue- and register-spill-bound, not FLOP-bound.
• Part 3 surveys the kept winners across 25 shapes and lands the campaign's one ~2× win (the causal early-exit).

The synthesis sections at the end — the iteration ledger, Strix Halo nuances, algorithmic differences from a gfx950 (CDNA) MFMA FMHA, and lessons ported from CK Tile — are the durable takeaways.

Reproduce everything here with:

PYTHONPATH=python python3 -m builders.gfx1151.attention.bench_v_staging \
    --seqlen-q 512 --seqlen-k 512 --head-size 128 --heads 8 --batch 4

(Run on the Strix Halo box: Radeon 8060S, RDNA3.5, wave32, WMMA wmma_f32_16x16x16_f16.)

---

The kernel under test

One wave32 owns 16 Q rows for a (q_tile, head, batch). Per K-tile it runs QK^T (WMMA) → online softmax → PV (WMMA), carrying the running max m, sum l, and the <8 x f32> PV accumulator as scf.for iter-args. BLOCK_M = BLOCK_K = 16.

The lever lives in the PV matmul's B-operand. WMMA computes A @ B^T, so PV = P @ V needs V in (d × k) layout: for this lane's d-column d_col, the B-fragment is V[k, d_col] for k = 0..15 — an inherently column-strided gather of V (stride = head_size).

---

The Loop

1. Hypothesis

The baseline gathers the PV V-operand straight from global memory with a per-(d, k) scalar global_load — n_dk * 16 scalar loads per lane per K-tile (128 at head_size=128). Staging each K-tile's V rows into a 16 × head_size LDS tile once (with wide vector loads), then reading the B-operand from LDS, should cut global-memory traffic and speed the kernel up.

This is the obvious move: fewer, wider global loads + on-chip reuse.

2. One lever

WmmaFmhaFwdSpec.v_lds_stage toggles exactly this and nothing else:

• v0_vgather (v_lds_stage=False): per-(d, k) scalar global gather.
• v1_vlds (v_lds_stage=True): each lane vector-loads its k-row's full head_size slice (8-wide global_load) into V_lds, then the PV B-operand is a strided LDS read. Costs one extra 16 × head_size × 2 B LDS tile and shares the existing P-staging barrier.

3. Measure (correctness gate, then time)

Every variant is gated against a numpy dense-attention reference before it is timed — no speed number is reported for an incorrect kernel. Both variants are bit-for-bit equal to each other and within 3e-5 of the reference.

shape (B Sq Sk D, Hq=Hk=8)	variant	max_abs	µs/iter	TFLOP/s
4 512 512 128	v0_vgather	3.05e-05	384.8	11.16
4 512 512 128	v1_vlds	3.05e-05	669.1	6.42
4 1024 1024 128	v0_vgather	3.05e-05	1517.1	11.32
4 1024 1024 128	v1_vlds	3.05e-05	2537.3	6.77
4 512 512 64	v0_vgather	3.05e-05	248.2	8.65
4 512 512 64	v1_vlds	3.05e-05	442.3	4.86
4 512 512 128 causal	v0_vgather	4.88e-04	444.1	4.84
4 512 512 128 causal	v1_vlds	4.88e-04	685.3	3.13

v1_vlds is a consistent 1.5–1.8× regression across every shape.

4. Inspect the ISA (explain the number)

Static instruction mix from llvm-objdump -d over the (fully unrolled) body, head_size=128:

variant	global_load	global_store	ds_load	ds_store	wmma
v0_vgather	160	64	4	8	16
v1_vlds	48	64	132	24	16

The hypothesis was right about the global loads — v1_vlds cuts them 160 → 48 (3.3×). It was wrong about everything that matters:

1. The strided access pattern is preserved, just relocated. The PV B-operand is fundamentally V[k, d_col] for fixed d_col across k — a column gather. Moving V into LDS does not make that pattern contiguous; it becomes 132 scalar (n=1), column-strided ds_loads with LDS bank pressure. LDS bought nothing the gather didn't already have.
2. The barrier has nothing to hide behind. This kernel runs one wave32 per CTA, so there is no second wave to overlap with the s_barrier that the staged path forces between the V store and the V read. The latency is exposed in full.
3. The baseline loads are cache-resident. Across the 16 lanes of the wave, each k-row V[k, :] is read once, distributed by column — so the gather's global traffic is already near-optimal and the L1/L2 services the rest. The "expensive" scalar gather is cheap in practice.

5. Keep / revert decision

Revert. v_lds_stage now defaults to False (the measured winner) in WmmaFmhaFwdSpec, mfma_attention_fwd_inner_body, and the WMMA inner body. The toggle is kept so this A/B stays reproducible.

---

Generalizable lessons

• Fewer global loads ≠ faster. Counting global_load in the ISA is a hypothesis, not a measurement. Trading global loads for LDS only wins when LDS changes the access pattern (gather → contiguous/broadcast) or enables reuse that the cache wasn't already giving you.
• LDS staging needs occupancy to pay for its barrier. With one wave per CTA the staging barrier is pure exposed latency. A staging optimization here only becomes viable after a structural change that puts multiple waves per CTA (a separate, larger lever).
• The cache is doing real work. On RDNA3.5 a column-distributed scalar gather of a small reused tile stays resident; "obviously bad" memory patterns can be fine, and the only way to know is to measure.

---

Part 2: the multi-lever campaign (fmha_singlewave.py + tune.py)

The single-lever study above reverts cleanly, but "one lever lost" is not "the kernel is optimal" — a lever can be dead on its own yet alive in combination, or masked by a different bottleneck. So we built a heavily-parameterized vehicle (fmha_singlewave.py, driven by tune.py) and swept levers individually and in combination, every variant GPU-measured against the numpy reference.

PYTHONPATH=python python3 -m builders.gfx1151.attention.tune \
    --bm 1 2 --pmode lds --vmode gather lds_t --qpreload 0 1

First, the hardware ceiling (why "200 TFLOP/s" is unreachable here)

The headline target was 200+ TFLOP/s. That is ~3.4× the physical peak of this part, so it cannot be reached by any kernel — stating it plainly up front so the rest of the numbers have an honest frame:

• f16 WMMA peak, Radeon 8060S (RDNA3.5, 40 CU, ~2.9 GHz): 40 CU × 512 FP16 FLOP/clk × 2.9e9 ≈ 59 TFLOP/s. That is the wall.
• Empirical reality check: a naive one-wave-per-16×16-tile WMMA GEMM (instances.gfx1151.wmma_gemm) sustains only 4.3 TF at 4096³ — it is memory-bound, re-reading A/B from global every K-step.
• This attention kernel sits at ~10.3 TF — i.e. it already beats the naive GEMM (better on-chip reuse) and is within noise of the production wmma_fmha_fwd (~11 TF). It is at ~17% of the 59 TF peak.

So the realistic envelope for this kernel family is tens of TFLOP/s, and the honest objective is "approach the 59 TF roofline," not "hit 200."

Why we're stuck at ~10 TF: it's issue- and register-bound, not FLOP-bound

Static instruction mix of the unrolled K-tile body (tune.py ISA counter + msgpack resource decode), head_size=128:

metric	value	reading
total instr / K-tile body	1398	the issue budget
wmma	16	~1% of issued instructions
global_load	160	128 of them are the PV V-gather
VGPR / wave	192 (the cap)	+ 26 spills to scratch
waves / CTA	1	no second wave to hide latency

Two structural walls fall out of this:

1. WMMA is ~1% of the instruction stream. Even with infinite occupancy the kernel is bottlenecked on VALU/memory issue, not matmul. You cannot approach the FLOP roofline while 99% of issued instructions are not matmul.
2. The register file is overcommitted. The PV accumulator alone is n_dk × <8×f32> = 64 VGPR at head_size=128; with K-frags and softmax temps the allocator pegs at the 192-VGPR cap and spills 26. Confirmed by the head_size=64 case (accumulator halves to 32 VGPR): vgpr=142, spill=0 — no spill, cleaner code. Causal head_size=128 spills worst (41).

Levers swept (all GPU-measured, B4 Sq512 Sk512 D128 Hq8 unless noted)

lever	result	verdict
Vectorized acc-rescale — rebuild alpha as a <8×f32> and rescale each acc[d] with one vector_mul instead of c_frag scalar extract/mul/insert	9.5 → ~9.9 TF	KEEP
Q-resident preload — hoist the BM·n_dk Q-frags out of the K-loop (Q is loop-invariant)	neutral (±noise); VGPR already at the 192 cap, and Q is L1-resident so dynamic loads were already cheap	default off
V-LDS staging, transposed (v_mode=lds_t) — revisited the rejected V-staging idea with the new occupancy understanding: stage V transposed so the B-operand read is contiguous	8.06 TF (regression); the transpose itself costs 136 ds_store scatters — staging just relocates the strided access, same lesson as Part 1, now reconfirmed against the register/occupancy hypothesis	revert
M-amplification (bm_tiles 2, 4) — one wave owns 2/4 Q-tiles to amortize the shared K/V load across more matmuls	consistent regression (D128 bm2 = 5.4 TF vs bm1 9.5; D64 bm4 = 2.9 TF) — bigger tiles shrink the grid and cut occupancy on a kernel that is already latency-bound	revert
K-frag fusion (fuse_k) — load each K-frag inside the QK matmul (1 live frag vs n_dk) instead of hoisting all n_dk K-frags before the loop	the campaign's win at D128: spills 26 → 10, 9.5 → 10.5 TF (512), 10.4 → 11.1 TF (1024). But it hurts at D64 (8.8 → 7.4 TF) where there is no spill to relieve and the extra reloads are pure added latency. So it is auto-enabled only when the kernel spills (fuse_k=None → head_size>=128)	KEEP (conditional)
BLOCK_N widening (fmha_blockn.py, bn_tiles 2/4/8) — consume bn_tiles 16-key subtiles per K-loop step so each iteration issues bn_tiles× more WMMA, directly attacking the "WMMA is ~1% of instructions" diagnosis	the WMMA fraction does rise (1.1% → 2.3% at bn8) but it loses at every width: D128 10.8 → 8.6/9.0/8.7 TF (spills explode 10 → 44 → 68 → 130 holding NK score+P tiles live), and even at D64 with spill headroom (bn2 stays spill=0) it still regresses 8.6 → 8.3 → 7.9 → 7.1 — the extra per-subtile V-gathers and P-transpose ds_loads outweigh the loop-overhead saved	revert
Software pipelining (fmha_pipelined.py, ck_tile qr_ks_vs) — hoist the next K-tile's QK (independent of the current tile's softmax) and carry the pre-computed score through the scf.for iter-args, so the matrix unit runs QK of tile i+1 while the VALU does the softmax of tile i	the mirror image of fuse_k: a win at D64 (no spill) — 8.6 → 10.7 TF (512), 11.1 TF (1024) — but a regression at D128 where carrying the extra <8×f32> score + the next-QK temps pushes spills 10 → 57 and drops it to 8.6 TF. The overlap only pays when there is register headroom to hold the carried state	KEEP (conditional, D64)

Three changes are kept. Vectorized rescale is a small unconditional win. The other two are complementary, register-headroom-gated levers that each attack the named bottleneck from the opposite side:

• fuse_k wins when the kernel spills (D128): it trades recomputed K-loads for fewer live VGPRs, relieving the spill (26 → 10) — so it is gated on head_size>=128.
• Software pipelining wins when there is register headroom (D64): it spends spare VGPRs on a carried next-tile score to overlap QK with softmax — but at D128 there is no headroom, the carry causes spilling (10 → 57), and it loses. So the per-shape best is D64 → fmha_pipelined, D128 → fmha_singlewave fuse_k, and both top out at ~11 TF.

The sched_group_barrier interleave hints (PipelinedCfg.sched) are noise-level here (±0.5 TF) — the spill, not instruction scheduling, is the wall. Everything else is reverted to its losing-but-reproducible toggle.

Two more levers, both negative (ds_bpermute P-transpose; lever recombination)

After software pipelining landed as the D64 winner, two follow-ups were tried.

ds_bpermute register P-transpose (PipelinedCfg.p_xpose="shuffle") — port ck_tile's gfx11 PermuteWarpGemmCToA (permlanex16 + v_perm_b32) to remap the QK-score C-distribution into the PV A-operand in registers, eliminating the LDS round-trip + s_barrier. The ISA confirms the LDS traffic vanishes (ds_load/ ds_store 12 → 0), but it is a structural dead-end on this DSL and a regression (D64 10.2 → 8.9, D128 8.9 → 8.3):

• ck_tile's permute only exchanges lane ↔ lane^16 (2 lanes) and feeds the result into a WMMA whose A-operand distribution is co-designed to consume that cheap permute (MakeABlockTileDistribution). The DSL's mma op has a fixed standard a_map where lane l needs row l — data gathered from 16 source lanes. permlanex16 provably cannot gather across 16 lanes, so the cheap permute cannot produce the layout the DSL's WMMA requires. (The shuffle output scores 1.1–1.3e-2 — wrong, just coincidentally close because post-softmax probabilities are similar in magnitude; it slips under the 2e-2 gate but is not correct.)

• A correct full-row gather needs 8–16 ds_bpermutes + selects, and ds_bpermute uses the same LDS hardware unit on RDNA — it removes the barrier, not the LDS access. On an issue-bound kernel, trading 12 ds_load/ ds_store for ~113 bit-twiddle instructions (+spills) is a guaranteed loss.

  Lesson: a register-shuffle transpose ported from a co-designed C++ kernel needs the consuming matmul's operand distribution to be customizable too; a DSL with a fixed WMMA a_map can't express it, and the generic substitute (ds_bpermute) is the same LDS engine with more instructions.

Recombining old levers with the winner — with software pipelining as the D64 base, fuse_k and sched were swept in combination (the standing "revisit rejected levers when the bottleneck model changes" instruction). One run showed a tempting fuse_k+sched = 10.93 TF synergy, but a repeat collapsed it into the noise: the whole D64 shape drifts ±0.6 TF run-to-run (thermal/clock), and on re-measure all four fuse_k × sched configs cluster within ±0.1 TF. No robust combination win — fuse_k/sched are noise-level at D64; the only result above the noise floor is software pipelining itself (8.6 → ~10–11 TF, a

1 TF lift). Lesson: measure twice before believing a 6% combination effect on a box with 6% thermal variance.

The structural lever, actually built: multi-wave (fmha_multiwave.py)

The hypothesis from every micro-lever was the same: the only thing that breaks the two walls is structural — put multiple wave32s in one workgroup, cooperatively stage K/V into LDS once per CTA (shared across all waves), and let the second wave hide the barrier latency that killed every single-wave staging attempt. So we built it: fmha_multiwave.py (MultiWaveCfg, driven by mw_tune.py). n_waves wave32s per CTA, each owns one 16-row Q-tile; K_lds ([16][hs]) and the transposed V_lds_t ([hs][16]) are loaded cooperatively by all threads, both WMMA B-operands become contiguous LDS reads, and the P-transpose uses a per-wave LDS slab with an intra-wave s_waitcnt (no cross-wave barrier).

PYTHONPATH=python python3 -m builders.gfx1151.attention.mw_tune --waves 2 4

It is correct (max_abs 3.05e-05) but slower: w4 = 6.3 TF vs the single-wave 10.5 TF. This is the decisive result of the campaign, not a disappointment: it confirms the kernel is not memory-bound. Cooperative LDS staging only pays when global traffic is the bottleneck; here the single-wave gather was already cache-resident, so moving K/V through LDS adds barriers and ds_load/ds_store overhead that the cache was absorbing for free — the same lesson as Part 1, now proven at the structural scale that was supposed to rescue it. The win would require also register-blocking larger output tiles per wave to raise the WMMA:overhead ratio (production CK GEMM design), a substantially larger rewrite; even then the ceiling is ~25–35 TF (40–60% of the 59 TF peak), not 200.

The full production rewrite, actually built: fmha_regblocked.py

So we built that substantially larger rewrite — the one design the campaign kept naming as the only remaining path. fmha_regblocked.py (RegBlockedCfg, driven by prod_tune.py) does both structural levers at once, the two that each lost alone:

• multi-wave + shared K/V in LDS (the fmha_multiwave lever) — amortize the staging barrier and global K/V traffic across num_warps query-row tiles;
• register-blocked larger output tiles per wave (the fmha_blockn lever) — each wave owns m_repeat M-atoms × n_repeat = block_n/16 N-atoms, issuing 2·m_repeat·n_repeat·n_dk WMMAs per K-tile (m1 n4 D128 = 64 vs single-wave 16) to raise the WMMA:overhead ratio.

The bet: the density rise pays for the staging barrier, and the extra waves give the barrier (and the spills) latency to hide behind. Warps partition along M only (warp_n=1), so softmax stays intra-warp (wave_reduce_* over 16 lanes

• an in-register combine across the n_repeat score fragments); K is staged [n][d], V transposed [d][n], P→A via the per-wave LDS slab (not ds_bpermute).

PYTHONPATH=python python3 -m builders.gfx1151.attention.prod_tune \
  --head-size 128 --num-warps 2 4 --m-repeat 1 2 --block-n 32 64

It is correct in every mode (max_abs 1.5e-5 D128, 3.05e-5 D64, 4.9e-4 causal) but a consistent ~2× REGRESSION: ~6 TF across the entire D64 and D128 sweep (num_warps∈{2,4} × m_repeat∈{1,2} × block_n∈{32,64}), vs the single-wave wall (~11 TF D64 fmha_pipelined, ~11 TF D128 fmha_singlewave fuse_k). This is the decisive close of the structural campaign — all three structural variants lose to the single-wave gather kernels:

variant	what it does	best TF	vs single-wave
single-wave gather	no LDS K/V, no blocking	~11	baseline
fmha_blockn	register-block, global gather	~8–9	regression
fmha_multiwave	multi-wave, LDS K/V, no blocking	~6.3	regression
fmha_regblocked	multi-wave + LDS K/V + register-block	~6	bigger regression

Register blocking did behave as designed once the B-operand reuse was made real. The first cut reloaded each K/V LDS B-operand inside the m_repeat loop (and Q inside the n_repeat loop), so one LDS read fed only one matmul — m_repeat=2 was slower than m_repeat=1. Hoisting the loads so one K_lds read feeds all m_repeat QK matmuls (and one V_lds_t read all the PV matmuls) flipped it: D64 w2 m2 n32 rose 5.49 → 6.19 TF over m1, WMMA density doubled (16 → 32), spill stayed 0. The lever works — it just isn't enough.

Each structural addition still makes it worse, monotonically, for one reason: the kernel is issue-bound and every scheme adds instructions. The ISA confirms it — even with proper reuse the WMMA fraction stays ~1–2% (wmma=32–64 of instr=1700–4400), because the cooperative LDS staging + the per-(mr,nr) P-transpose round-trips add instructions in proportion to the WMMA density they buy, leaving the ratio flat. Spills were even relieved (D128 spill=0 here vs single-wave spill=10/11) and it still lost — proving the wall was never spills or memory traffic but raw instruction issue. The production-CK LDS-staging design was built for an HBM-bandwidth-bound CDNA part with dense MFMA; on this large-cache, issue-bound RDNA3.5 APU those assumptions invert — the single-wave global gather is cache-resident and the cheapest possible way to feed the matrix unit, so any LDS staging is a net loss no matter how many waves amortize the barrier. (The B-operand LDS reads still use the vec_concat-of-two-<8×f16> pattern — two ds_read_b128s instead of 16 scalar ds_loads — so the loss is not an artifact of naive LDS reads; the structural overhead is intrinsic.)

Conclusion of the campaign: the single-wave gather kernels (fmha_pipelined at D64, fmha_singlewave fuse_k at D128) at ~11 TF (~18% of the 59 TF peak) are the best this DSL achieves on gfx1151. The remaining ~5× gap to peak is not reachable by any structural reorganization expressible here; it would need RDNA4/gfx12 intrinsics (larger-K WMMA, ds_read_tr, async-LDS — see below).

Different WMMA intrinsics? A dead-end on this silicon

A natural question: would a different WMMA atom relieve the register/issue walls? Investigated against core/arch/target.py (_MMA_FRAGMENT_INFO) and core/isa/backend.py (_RDNA_WMMA, MemoryCapabilities for gfx1151):

• f16-accumulate WMMA (vs the f32 accumulator) would halve the PV accumulator VGPRs in principle — but on RDNA3.5 the f16-output WMMA still writes a <8 x float>-shaped C fragment (8 VGPR); packed-f16 C output is a gfx12/RDNA4 feature, not available here. No register saving.
• Larger-K WMMA (16x16x32) would cut the QK/PV instruction count per K-tile (better WMMA:overhead ratio) — gfx12 only.
• ds_read_tr transpose-load would remove the P-transpose shuffle/LDS round-trip — MemoryCapabilities.has_ds_read_tr=False on gfx1151; also gfx12.
• has_async_lds=False too, so no async-copy overlap for the staging path.

The only WMMA atom on gfx1151 is wmma_f32_16x16x16_f16 (_bf16 is the same shape). Different intrinsics are a dead-end lever on this part — the relevant ones are all RDNA4/gfx12. The realistic path remains the structural rewrite above, capped by the 59 TF roofline.

Part 3: the 25-shape survey + the causal early-exit win (survey.py)

The micro/structural campaign above all measured one shape (B4 Sq512 D128). To see where the kernels actually land across the workload space, survey.py runs 25 shapes — D64/D128 × seqlen {512…4096} × batch/heads/GQA/causal — keeping the faster of fmha_pipelined and fmha_singlewave fuse_k per shape, and timing PyTorch SDPA (only the MATH fallback exists on this APU; FLASH/EFFICIENT are runtime-disabled on gfx1151) as the reference.

The survey exposed one large, un-swept inefficiency: causal shapes ran at half the efficiency of non-causal (~7.5–8.4% of peak vs ~15–19%). The cause was not a micro-lever — the causal kernels looped every K-tile and masked the upper triangle, doing ~2× the matmul work the math requires. The fix is the standard flash-attention causal early-exit: a CTA owning query rows [group_row0, group_row0+BM·16-1] only needs K-tiles whose lowest key index is ≤ its max query position, i.e. kt < group_row0/16 + BM. Clamp the K-loop bound to that (one cmp_lt + select, gated on mask_mode=="causal"):

loop_stop = seqlen_k // 16
if causal:
    causal_stop = group_row0 // 16 + BM        # last needed tile + 1
    loop_stop = select(causal_stop < loop_stop, causal_stop, loop_stop)

Added to both fmha_pipelined.py and fmha_singlewave.py. fmha_pipelined's software pipeline stays correct: its prologue computes tile-0 QK (always needed, loop_stop≥1) and its next-tile prefetch is already clamped against loop_stop, so the carried score_next on the final iteration is computed-but-never-consumed as before.

Result — causal throughput ~doubled, in line with non-causal:

causal shape (Sq=Sk)	before	after
B4 H8 S1024 D64	~4.7 TF	8.92 TF
B4 H8 S2048 D64	~5.0 TF	9.34 TF
B4 H8 S1024 D128	~5.0 TF	8.44 TF
B4 H8 S2048 D128	~5.2 TF	8.89 TF

This is the rare ~2× win in the whole campaign — and it is algorithmic, not a hardware-roofline play: we stopped issuing matmuls whose output is masked to zero. It does not move the issue-bound non-causal ceiling.

Survey aggregate (25 shapes, after the fix): best 11.39 TF (19.3% of peak, B2 GQA Sq1024 D128); avg 17.0% of peak (was 14.4% before the causal fix); 13.3× vs SDPA (causal shapes up to 29.5×, since SDPA's MATH path also does the full quadratic). Non-causal sits at 15–19% and causal now at 14–16% — uniformly ~17%, the issue-bound single-wave wall. 0/25 reach 25% of peak: that target needs the matmul:overhead ratio to rise, which on this part requires gfx12 larger-K WMMA, not another single-wave lever (see "Different WMMA intrinsics" above).

Iteration ledger (what was tried, and how much)

The campaign ran the runbook loop — hypothesis → one lever → GPU-measure against the numpy/torch reference → inspect ISA → keep/revert — across four kernel bodies and ~14 distinct levers, every variant correctness-gated before timing:

kernel	levers swept	outcome
fmha_singlewave	bm_tiles, p_mode, v_mode (gather/lds_t), q_preload, fuse_k, vectorized acc-rescale, vectorized P-read (vec_concat), epilogue inv_l hoist	D128 winner ~11 TF (fuse_k auto on D128)
fmha_pipelined	software pipelining (qr_ks_vs), sched hints, p_xpose (lds/shuffle)	D64 winner ~11 TF
fmha_blockn	bn_tiles 2/4/8 (BLOCK_N widening)	regression (spills)
fmha_multiwave / fmha_regblocked	n_waves, then num_warps × m_repeat × block_n (+ register-blocked B-operand reuse, double-buffer stub)	regression (~6 TF)

combo.py then ran the full cartesian product of every compatible single-wave lever across fmha_singlewave/fmha_pipelined/fmha_blockn per shape, to answer "did any combination beat what each lever scores alone?" — it did not: the per-shape best is always a single lever (D128 opt bm1 fuse_k 19.3%, D64 sp sched0 16.9%), and 0 combinations clear 20% robustly (best runs straddle it at 19–20%, the remainder thermal noise). survey.py then measured the kept winners across 25 shapes (avg 17% of peak, 13.3× vs SDPA).

Two micro-fusions from the final pass are kept because they cut static instructions with bit-identical output (the kernel is issue-bound, so instruction count is the currency): the vectorized P-transpose read (two ds_read_b128 + vec_concat instead of 16 scalar ds_load + vec_insert, 1407 → 1395 instr) and the epilogue inv_l hoist (rcp/select computed once per row instead of per d). Both are wall-clock-neutral within the ±0.6 TF thermal noise — expected, and consistent with the issue-bound model: removing a handful of the ~1400 body instructions is real but small.

Correctness is verified against PyTorch SDPA (MATH backend, the only one enabled on gfx1151), not just the numpy reference: on identical inputs the kernel matches torch to max_abs 3e-5 (non-causal D64/D128/GQA) and 4.9e-4 (causal) — well inside the 2e-2 f16-WMMA-accumulation-order gate.

Strix Halo nuances (why this part behaves differently)

gfx1151 is an APU, and almost every surprising result traces back to that:

• Unified LPDDR5X memory + a large last-level cache. K/V for a single (head, batch) flash-attention tile is small and reused across the wave's 16 lanes, so the "expensive" column-strided global gather stays cache-resident. This is the root cause of the whole campaign: LDS staging (Part 1, fmha_multiwave, fmha_regblocked) replaces cheap cached reads with serialized ds_loads + barriers and always loses. On a discrete HBM part the same staging is mandatory.
• wave32, one WMMA atom. The only matrix intrinsic is wmma_f32_16x16x16_f16 (_bf16 is the same shape) — K=16, f32-only C accumulator (<8×f32>, 8 VGPR even for f16 output; packed-f16 C is gfx12). No 16x16x32, no ds_read_tr, no async-LDS (MemoryCapabilities all False).
• 192-VGPR cap with no slack at D128. The PV accumulator alone is n_dk × <8×f32> = 64 VGPR at D128; with K-frags + softmax temps the allocator pegs the cap and spills. D64 halves the accumulator and runs spill-free — which is why the two best levers split by head_size (fuse_k relieves D128 spills; software pipelining spends D64's spare registers).
• ~59 TF f16 WMMA peak (40 CU × 512 FLOP/clk × ~2.9 GHz). The original "200 TF" target is ~3.4× the physical wall — unreachable by any kernel.
• Thermal drift ±0.6 TF run-to-run; the box heat-soaks during long build sweeps, lowering absolute numbers. Every keep/revert decision was made A/B back-to-back in one thermal window, never across sweeps — several apparent 6% "wins" evaporated on repeat.
• No FLASH/EFFICIENT SDPA backend on gfx1151; only the MATH fallback exists, so PyTorch is a correctness reference and a (quadratic) throughput floor, not a fair flash-vs-flash comparison.

Algorithmic differences from a gfx950 (CDNA) MFMA FMHA

The production CK / CK-Tile FMHA kernels target CDNA (gfx9xx, e.g. gfx950). The flash-attention math is identical (online softmax, causal early-exit, the qr_ks_vs pipeline), but the kernel structure inverts on several axes:

axis	gfx950 (CDNA, MFMA)	gfx1151 (RDNA3.5, WMMA)
wavefront	wave64	wave32 — softmax row-reduction butterfly spans 16 lanes, half the lane masks
matrix unit	v_mfma_*, many tile shapes (16x16x16, 16x16x32, 32x32x8) → high arithmetic intensity per instr	one 16x16x16 WMMA → low WMMA:overhead ratio, the issue-bound wall
operand semantics	MFMA A·B with native K-accumulation; large K tiles	WMMA A·B^T — forces the V column-gather / transpose dance the campaign fought
C fragment	can output packed f16	always <8×f32> — no accumulator VGPR relief
memory bound	HBM-bandwidth-bound → LDS staging + double-buffer + async-copy is the correct design	cache-resident / issue-bound → LDS staging is a net loss
transpose	ds_read_tr transpose-loads, async global→LDS	neither exists → P→A must round-trip LDS or ds_bpermute (both costly)
register file	large; deep pipelines fit	192-VGPR cap; even one-stage pipelining spills at D128

The headline inversion: the production LDS-staging FMHA design is bandwidth- optimal on CDNA and pessimal here. Porting its structure verbatim (fmha_multiwave/ fmha_regblocked) reproduces the CDNA dataflow faithfully and loses, because the bottleneck it was built to relieve (HBM bandwidth) is not this part's bottleneck (instruction issue, against a cache that already feeds the matrix unit for free).

Lessons ported from CK Tile

• qr_ks_vs software pipelining (hoist next-tile QK, overlap with current softmax) → fmha_pipelined. Ports cleanly and is the D64 win; at D128 the carried score state spills (10 → 57) and it loses. Lesson: a pipeline depth that's free on a large register file costs spills on a 192-VGPR part.
• PermuteWarpGemmCToA (register C→A transpose via permlanex16 + v_perm) → attempted as p_xpose="shuffle". Cannot port: CK Tile co-designs the consuming WMMA's operand distribution (MakeABlockTileDistribution) so a 2-lane permute suffices; the DSL's mma has a fixed a_map needing a 16-lane gather, which permlanex16 provably can't do, and the generic substitute ds_bpermute is the same LDS engine with more instructions. Lesson: a register-shuffle transpose needs the consuming matmul's layout to be customizable too.
• Cooperative LDS K/V staging + register blocking (MRepeat/NRepeat) → fmha_multiwave + fmha_blockn + fmha_regblocked. The register blocking is genuinely CK's density lever and works as designed (WMMA 16 → 32, spill 0); the LDS staging is the bandwidth lever and inverts here. Lesson: separate a porting source's levers by which bottleneck each attacks, and re-test each against the target part's actual bottleneck rather than adopting the bundle.
• Online-softmax + causal early-exit → the one unambiguous shared win (~2× on causal), because it's algorithmic (skip masked matmuls), not microarchitectural.

Generalizable lessons (campaign)

• Name the roofline before optimizing. A target above the hardware peak (200 vs 59 TF) is a spec bug, not a kernel bug — catch it with one line of arithmetic before burning iterations.
• Count the matmul fraction. If WMMA is 1% of issued instructions, you are issue-bound; chasing the FLOP roofline is futile until that ratio changes.
• Re-test rejected levers when the bottleneck model changes — but believe the measurement. We revisited V-staging under the register/occupancy lens (the user's standing instruction); the transposed variant is new, yet it still loses, for a now-better-understood reason (the transpose scatter, plus a barrier with no second wave to hide it). Revisiting was right; overriding the GPU number would not be.
• Spills are visible without a profiler. The AMDGPU msgpack note carries .vgpr_count / .vgpr_spill_count; decode it straight from the HSACO (tune._resource_counts) — no llvm-readelf needed.
• Survey before you micro-optimize. One-shape tuning hid that causal ran at half efficiency. The biggest single win in the campaign (~2× on causal) was an algorithmic skip of masked work, not a hardware lever — and it was only obvious once the 25-shape sweep put causal and non-causal side by side.

Reusing — and extending — the CK Tile helper layer for RDNA WMMA

All five kernels are built on the CK Tile helper layer (rocke.helpers): make_global_view + make_tile_window (with shift_by for the K-loop step) for Q/K/V/O addressing, make_lds_view + TileWindow for the P-transpose round-trip, the helpers.attention softmax primitives (wave_reduce_max/sum, apply_attention_mask), and a WMMA atom for the QK/PV matmuls.

The helper layer was built for CDNA / MFMA / wave64, so three pieces were added for RDNA3.5 WMMA / wave32. All three are additive — no existing MFMA or f32 path changed; the full test_rocke suite stays green:

gap	what was added	where
only MfmaAtom (wave64, MFMA) existed	WmmaAtom (+ wmma_atom, WMMA_F16/BF16_ATOMS): wave32, m=n=k=16, a/b_per_lane=16, c_per_lane=8 (<8×f32>), name="wmma_f32_16x16x16_f16". Lane-layout accessors delegate to the existing target.mma LayoutMaps so there is one source of truth; emit routes through b.mma(name,…)	helpers/atoms.py
WarpGrid was MFMA/wave64-only	a wave32 WMMA path accepting the WMMA atom	helpers/geometry.py
load_tile is f32-only (casts every element via load_vec_as_f32) — not a packed <16×f16> fragment, so not directly b.mma-able	load_wmma_fragment / store_wmma_acc: a packed fragment load/store built on TileWindow.load_vec(n=16) (raw packed vector → directly b.mma-able), matching the op's layout maps	helpers/distribution.py, helpers/tensor_view.py
StaticDistributedTensor stores a per-element f32 list (the right shape for elementwise/reduce ops, the wrong one for an issue-bound WMMA fragment)	WmmaTensor: a packed distributed tensor carrying one lane's fragment/accumulator as a single SSA vector, with tile-level load_wmma_tile / wmma_mma / store_wmma_tile wrappers and tile.scale (one v_mul) / tile.coord (off the verified layout map). Each is 1:1 over the underlying op, so the kernel body reads in tile terms with zero added instructions	helpers/distribution.py

Why the packed fragment path matters. This kernel is issue-bound — static instruction count is the perf currency — so the WMMA operand load must be a single packed vector-load feeding b.mma, and the accumulator must carry as a packed vector (the rescale is one v_mul, not eight scalar muls). Routing a WMMA operand through the f32 load_tile, or a fragment through the per-element StaticDistributedTensor, would insert per-element casts/repacks and regress TFLOP/s. So load_wmma_fragment/WmmaTensor stay packed (built on the raw load_vec, not load_vec_as_f32); all five kernels read in tile terms (load_wmma_tile → wmma_mma → acc.scale → store_wmma_tile) while lowering to the identical single load / b.mma / v_mul. Lesson: on an issue-bound kernel a convenience cast — or a per-element container — in a generic path is a silent regression; the fragment and the accumulator have to stay packed.

Files

• ALGORITHM.md — a from-the-math-up guide to the winning kernel for readers new to flash attention: the attention spec, the online-softmax recurrence with a correctness proof, the base-2 exp trick, the causal early-exit, and a step-by-step map from each line of the recurrence to the kernel body. Read before the case study below.
• bench_v_staging.py — Part 1 A/B harness (V-staging on the production spec): builds both variants, gates each on a numpy reference, times with HIP events (runtime.launcher.time_launches), disassembles for the memory-instruction mix, writes v_staging_perf.csv.
• v_staging_perf.csv — Part 1 results.
• fmha_singlewave.py — Part 2 vehicle: a self-contained, heavily-parameterized WMMA FMHA-fwd kernel (SingleWaveCfg: bm_tiles, p_mode, v_mode, prefetch_k, q_preload, fuse_k). Same WMMA contract/ABI as production; exposes the swept levers. fuse_k=None auto-enables K-frag fusion when head_size>=128. Keeps the vectorized acc-rescale, the vec_concat P-transpose read, and the epilogue inv_l hoist (the issue-bound instruction-cut micro-fusions).
• tune.py — Part 2 driver: verify_and_time(cfg, shape) gates on numpy, times with HIP events, and tallies both the disassembly instruction mix and the VGPR/SGPR/spill/LDS resource note. main() sweeps --bm/--pmode/--vmode/ --qpreload/--fusek and prints the running best.
• fmha_multiwave.py — the structural lever, built: a multi-wave WMMA FMHA-fwd kernel (MultiWaveCfg: n_waves). n_waves wave32s per CTA cooperatively stage K/V into LDS. Correct but slower (confirms not memory-bound) — see "The structural lever, actually built" above.
• mw_tune.py — driver for fmha_multiwave.py; --waves axis, same verify+time+ISA harness as tune.py.
• fmha_regblocked.py / prod_tune.py — the full production rewrite: a register-blocked multi-wave kernel (RegBlockedCfg: num_warps, m_repeat, block_n) combining the fmha_multiwave and fmha_blockn levers. Correct in every mode but a ~2× regression (~5 TF) — the decisive close of the structural campaign (see "The full production rewrite, actually built" above). Driver sweeps --num-warps/--m-repeat/--block-n.
• fmha_blockn.py / bn_tune.py — the WMMA-density lever: a BLOCK_N-widened single-wave kernel (BlockNCfg.bn_tiles) that processes bn_tiles 16-key subtiles per K-loop step. Correct but a regression at every width (see the levers table) — confirms the kernel is register/issue-bound, not loop-overhead-bound.
• fmha_pipelined.py / sp_tune.py — the software-pipelining lever (ck_tile qr_ks_vs): a single-wave kernel (PipelinedCfg) that hoists the next K-tile's QK and carries the score through the scf.for iter-args to overlap QK with the current tile's softmax. The genuine D64 win (8.6 → 11 TF) and a D128 regression (carry causes spilling) — the register-headroom mirror of fuse_k. Driver axes --sched 0 1 (sched_group_barrier hints, noise-level), --fusek auto 0 1, and --pxpose lds shuffle (the ds_bpermute register P-transpose; structural dead-end, see above). PipelinedCfg.p_xpose defaults to the measured winner "lds".
• combo.py — the full cartesian-product sweep: every compatible single-wave lever across fmha_singlewave/fmha_pipelined/fmha_blockn, per shape, each GPU-verified and timed; reports the best combination per shape and globally. Answered "did any combination beat each lever alone?" — no (best is always a single lever; 0 combinations robustly clear 20%). --quick for a 2-shape smoke.
• survey.py — Part 3: the 25-shape survey. Runs both best single-wave kernels per shape (keeps the faster verified one), times PyTorch SDPA (MATH fallback), reports TFLOP/s, % of the 59 TF peak, and speedup. The vehicle that surfaced the causal early-exit win (fmha_pipelined/fmha_singlewave mask_mode=="causal" clamp the K-loop to skip fully-masked tiles).