[morousg]

Phase 4 of the RDNA4 port: FlashAttention on gfx12 WMMA, bf16 and FP8 QK^T. Tested and benchmarked on the RX 9060 (gfx1200, ROCm 10.1).

Stacked on #394, #395 and #396: this branch contains their commits. The phase 4 change is the last commit, 7f1d2ea.

1. Spike (spikes/rdna4/wmma_gfx12.hip, not built by CMake)

• Builtins: __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12 (A/B: 8 × bf16 as short8, C/D: float8) and __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12 (A/B: 8 e4m3 bytes as int2), the operand types rocWMMA uses for gfx1200/1201.

• Lane mapping, validated against a CPU matmul (bf16 max error 2.4e-7, fp8 exact). Lane l holds:

  • A: row l%16, k = 8*(l/16) .. +7;
  • B: column l%16, the same k split;
  • C/D: column l%16, rows 8*(l/16) .. +7.

  This differs from gfx11 and from NVIDIA's m16n8k16.

• Raw WMMA throughput, all waves busy, median of 20 runs: bf16 ≈ 40 TFLOPS, fp8 ≈ 82 TFLOPS. The number of independent accumulator chains (1 to 16) makes little difference. HIP reports 14 "CUs" (WGPs); rocminfo shows 28 CUs at 2.4 GHz.

2. Kernel (attention/flash_attention_wmma.h, HIP only)

FlashAttentionWmmaDPP, the counterpart of FlashAttentionMmaDPP: FA-2 forward with online softmax in fp32 registers.

• Transposed formulation: each wave owns 16 queries and computes Sᵀ = K·Qᵀ and Oᵀ += Vᵀ·Pᵀ. The accumulator layout of Sᵀ (lane = query, 8 keys) is the B layout of Pᵀ, so P feeds the second WMMA straight from registers with no LDS round trip. Row max and row sum are lane-local, with one xor 16 exchange each.
• Staging: K is staged row-major and V transposed in LDS, both padded against bank conflicts. Loads go through the prologue IOps, like FlashAttentionDPP: one 16-byte load for raw bf16, an 8-byte load plus dequantization for int8/fp8 caches, and an element-wise read for any other prologue or fused chain. Q goes through its IOp once per block. The epilogue chain runs in registers before the single write.
• One tile ahead in registers (bf16 only): measured +1–7% for bf16 and −12–20% for the fp8 path, so the fp8 path doesn't use it. __launch_bounds__ removed the spills (up to 142 VGPRs) the 192-VGPR default cap caused at d128.
• FP8 QKᵀ: an Fp8TokenDequantRead K selects it, mirroring the CUDA FP8_QK design. The K tile holds the raw e4m3 cache bytes. Q is quantized per query in-kernel (scale = max|q| / 448), and the scores are rescaled by kScale·qScale. V and P stay bf16.
• Defaults: BLOCK_KV 32 with 8 waves (128 queries per block), the fastest of BLOCK_KV {32, 64} × {4, 8} waves on every measured shape.
• API (new, no existing signature changed):
  • executeFlashAttentionWmma<HEAD_DIM, BLOCK_KV, NUM_WAVES>(q, k, v, o, bh, seqQ, seqK, causal, stream, scale, epilogue) throws on non-gfx12 devices.
  • executeFlashAttentionAuto<HEAD_DIM>(...) picks the WMMA kernel on gfx12 and FlashAttentionDPP elsewhere. Note that the WMMA path rounds Q, K, V and P to bf16, while the scalar kernel stays fp32 throughout.
  • On non-gfx12 targets the builtins compile to a trap, so a fat binary with other arches still builds.

Tests (utest_flash_attention_wmma.h, new __ONLY_HIP__ marker: _hip only)

On non-gfx12 GPUs the test prints SKIPPED. It has 17 cases against the double-precision oracle, with HEAD_DIM 32/64/128, dense/causal, ragged and cross-attention shapes, and both tile configs. Max errors on the RX 9060:

Path	Max error	Tolerance
raw bf16	≤ 1.5e-3	WMMA_BF16_TOL 1e-2
fp32 prologue (rounded to bf16)	2.3e-3	WMMA_BF16_TOL 1e-2
int8 KV	5e-4	WMMA_BF16_TOL 1e-2
fp8 QKᵀ	≤ 9e-3	WMMA_FP8_TOL 3e-2
fused epilogue vs plain run	1.2e-7	1e-6
executeFlashAttentionAuto vs executeFlashAttentionWmma	0 (bit-identical)	0

The tolerances are documented in the test. For bf16: P has 8 significant bits and outputs are convex combinations of V in [-1, 1], so the error stays below 2⁻⁸ = 3.9e-3, plus input rounding. For fp8, Q has 3 mantissa bits.

Mutation check: reading the wrong half of the Vᵀ fragment fails 12/14 cases.

Throughput, kernel vs scalar FlashAttentionDPP (TFLOPS = 4·bh·s²·d, halved for causal; bf16 inputs for both; median of 20; spikes/rdna4/fa_wmma_bench.hip)

shape	scalar	WMMA bf16 (default tile)	WMMA FP8 QKᵀ (BKV64/W8)	bf16 speedup
d64 bh16 s1024 dense	0.18	11.5	9.4	64×
d64 bh16 s4096 dense	0.19	13.8	11.9	72×
d128 bh16 s1024 dense	0.15	9.5	8.6	63×
d128 bh8 s4096 dense	0.16	10.7	9.7	67×
d64 bh16 s1024 causal	0.18	8.3	6.6	47×
d64 bh16 s4096 causal	0.19	12.3	10.3	66×
d128 bh16 s1024 causal	0.15	7.6	6.6	50×
d128 bh8 s4096 causal	0.16	9.1	8.5	57×

• The bf16 kernel reaches about 35% of the raw WMMA peak.
• FP8 QKᵀ is 0.85–0.95× bf16, not faster. The kernel isn't math-bound at this point, so the fp8 WMMA's 2× rate doesn't show, while the in-kernel Q quantization, the per-key scales and the V dequantization add work. For now the fp8 cache's benefit is half the KV bytes and memory, not forward speed.
• Remaining levers, not done: double-buffered LDS to drop one barrier per tile, cheaper softmax bookkeeping (exp2 per score is quarter-rate), and a packed fp8 → bf16 V conversion.

Verified on this machine

• Full HIP suite: 52/52, Phase 3's 51 plus utest_flash_attention_wmma_hip. Same unique warnings.
• nvcc 13.4 + g++-13 (sm_86): the full build has 0 failures and 0 warnings, no WMMA target is generated, and the 49 _cpp tests pass with the same list. The new header is __HIPCC__-only, so no CUDA-visible code changed apart from the test-discovery cmake.

🤖 Generated with Claude Code

https://claude.ai/code/session_014d2rht2WMSopuodhtXYZzx

[Copilot]

🟡 Changes recommended

Device dispatch, FP8 indexing, compilation guards and numerical validation still need correction.

6 open issues

• Use stored K height for FP8 scale-plane stride · New
• Cache device capabilities per device, not per template · New
• Guard unsupported MMA kernels at compile time · New
• Guard FP8 code when FP8 support is unavailable · New
• Reject NaN differences in oracle error reduction · New
• Reject non-finite values in epilogue and Auto comparisons · New

What changed in this PR

Adds RDNA4 WMMA FlashAttention to FKL, alongside the stacked HIP attention port and supporting tests.

Changes:

• Adds bf16 and FP8 WMMA kernels with automatic backend selection.
• Ports scalar attention and decode to HIP and fixes decode loads.
• Expands validation, benchmarks, HIP architecture configuration and test discovery.

File	Description
utests/​algorithm/​attention/​utest_softmax.h	Makes softmax tests GPU-portable.
utests/​algorithm/​attention/​utest_flash_decode.h	Adds decode and packed-load coverage.
utests/​algorithm/​attention/​utest_flash_attention.h	Adds HIP execution and FP8 checks.
utests/​algorithm/​attention/​utest_flash_attention_wmma.h	Tests WMMA precision, epilogues and dispatch.
utests/​algorithm/​attention/​utest_flash_attention_mma.h	Adds CUDA MMA validation.
utests/​algorithm/​attention/​utest_attention_wave.h	Tests wave-local exchanges.
utests/​algorithm/​attention/​attention_test_common.h	Shares oracle and buffer helpers.
spikes/​rdna4/​wmma_gfx12.hip	Validates WMMA mapping and throughput.
spikes/​rdna4/​fa_wmma_bench.hip	Benchmarks WMMA against scalar attention.
include/​fused_kernel/​algorithms/​attention/​softmax.h	Adds HIP launches and host conversions.
include/​fused_kernel/​algorithms/​attention/​flash_decode.h	Adds HIP support and load fallbacks.
include/​fused_kernel/​algorithms/​attention/​flash_attention.h	Shares low-precision readers and enables HIP.
include/​fused_kernel/​algorithms/​attention/​flash_attention_wmma.h	Implements WMMA attention and dispatch.
include/​fused_kernel/​algorithms/​attention/​flash_attention_mma.h	Relocates shared readers and traits.
include/​fused_kernel/​algorithms/​attention/​attention_wave.h	Adds portable wave32 helpers.
cmake/​tests/​discover_tests.cmake	Supports GPU-only and HIP-only markers.
cmake/​libs/​hip/​hip.cmake	Aligns HIP target architectures.
cmake/​libs/​hip/​hip_init.cmake	Applies architecture overrides before initialization.
cmake/​libs/​hip/​archs.cmake	Uses resolved HIP architectures.
AGENTS.md	Documents HIP configuration and attention support.

_{Review effort: Balanced}