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

gfx942 unified attention — parity + benchmark harness

A torch-reference parity + latency harness for the gfx942 (CDNA3 / MI300X) unified-attention SDPA-forward kernels: attention_tiled_2d for prefill and attention_tiled_3d for the split-KV decode path, both in kernels.gfx942.attention_tiled_2d / ..._tiled_3d. Unlike the gfx950 attention harness it has no Triton/AITER dependency — the oracle is an fp32 torch reference — so it runs on any box with torch + a gfx942 GPU.

New to flash attention? ALGORITHM.md derives what these kernels compute from the math up — the online-softmax recurrence, the paged-KV layout the harness feeds them, the causal masking, and the CDNA3-specific transposed-x8 (register-P^\top) matmul that the D128-fp16 flash regime uses. Read it first if you want the why before the how to run.

What it shows

• A correct gfx942 tiled SDPA-forward over a canonical inline shape set (default / fmha / creative groups): fp16/bf16, head_size 64 / 128, MHA + GQA, and decode / prefill regimes — each gated against an fp32 paged reference.
• The two shipped D128-fp16 flash geometries side by side: the provider's analytic-default wide4 (workgroup 256, num_warps=4) and the L4 (workgroup 64) contrast, selectable by environment variable.
• A CK-vs-Torch latency comparison in two timing modes — eager (host launch overhead + kernel) and CUDA-graph (kernel-only, host overhead removed) — so a perf gap can be attributed to kernel time vs launch overhead.

Prerequisites

• A gfx942 (CDNA3 / MI300X) GPU and a ROCm PyTorch build (torch.cuda must resolve to the HIP device).
• The repo's python/ on PYTHONPATH (commands below set it).
• No numpy required: the harness uses plain Python lists for the block tables so it runs on the minimal container torch build.

File map

file	role
parity_unified_attention.py	per-shape parity + latency harness; defines inline scenario groups (default / fmha / creative), builds the gfx942 spec explicitly (wide4 vs L4), and exposes the HIPDNN_GFX942_* lever knobs via --scenario
final_shapes_check.py	the definitive correctness + perf check over every shape via the production dispatcher (run_unified_attention_torch), timed eager + graph against PyTorch's flash SDPA
__init__.py	package marker + module docstring

Running

Both scripts need torch + a gfx942 GPU. From the repo's composablekernel/ directory (so python/ is the package root):

# 1) The definitive correctness + perf check over every shape, via the
#    production dispatcher, timed eager + graph vs PyTorch flash SDPA:
PYTHONPATH=python python python/rocke/examples/gfx942/attention/final_shapes_check.py

# 2) The explicit parity + latency harness (builds the spec by hand; default
#    runs all inline scenario groups):
PYTHONPATH=python python python/rocke/examples/gfx942/attention/parity_unified_attention.py

final_shapes_check.py flags (all verified against argparse)

flag	default	meaning
--groups G [G ...]	all non-*_disabled groups	which shape groups to run (e.g. --groups decode, or --groups d256_disabled to include the off-by-default d256 set)
--warmup N	10	warm-up launches before timing
--iters N	50	inner timed launches per measurement
--reps N	10	timed measurements per cell, reduced per --reduce
--reduce {median,mean,trimmed}	trimmed	how to reduce the --reps measurements
--trim F	0.2	fraction stripped from each end for --reduce trimmed (must be in [0, 0.5))

It prints, per shape, the eager and graph CK times, the eager and graph Torch times, the two CK/Torch ratios (a ratio < 1 means CK is faster), the SDPA backend used (flash or default), and PASS/FAIL. It exits non-zero on any correctness failure.

parity_unified_attention.py flags

flag	default	meaning
--scenario S (repeatable)	all shapes	a group (default / fmha / creative / all) or an exact shape name; multiple flags union together
--attempts N	30	timed launches
--warmup N	10	warm-up launches
--tol F	2e-2 fp16 / 4e-2 bf16	absolute-tolerance override
--report PATH	none	write a JSON results report
--debug-mismatch N	0	on failure, print the N worst mismatch samples

Note on --scenario: the selector accepts the group names default, fmha, creative, and all. --scenario all (or no flag) runs all three groups; any other value is matched as an exact shape name, e.g. --scenario decode_d128_q1k1024.

# Force the L4 (WG=64) flash geometry instead of the default wide4, on one shape:
HIPDNN_GFX942_FLASH_WIDE=0 PYTHONPATH=python python \
    python/rocke/examples/gfx942/attention/parity_unified_attention.py \
    --scenario fp16_h32kv8_b1_s2048x2048_d128

The shape set (inline scenarios)

Shapes are defined inline in parity_unified_attention.py across three groups; every shape is causal (the kernel always applies a causal mask), so non-causal prefill shapes are not included.

group	what it covers
default	core production shapes — decode (seqlen_q == 1) and prefill, fp16/bf16, D64/D128, MHA + GQA, varied batch sizes
fmha	integration-test net shapes — GQA ratios, asymmetric seq-lens, tiny seqs, cross-attention, large-batch decode, and select D256 shapes
creative	exploratory sweep — long-KV decode (up to 64 k), chunked prefill, non-standard GQA ratios, D64/D256 bf16, large-batch variants

Arch notes

• The D128-fp16 flash regime is the CDNA3 trick. For head_size = 128 fp16 the harness routes to the transposed-x8 path (use_mfma_32x32x8 + use_transposed_qk_32x32): it computes S^\top = K Q^\top via the gfx942-legal mfma_f32_32x32x8_f16 atom so that P^\top stays register-resident as the PV B-operand — no P_lds round-trip. See ALGORITHM.md §6.2.
• use_mfma_32x32x8 is fp16-only (gfx942 has no bf16 32x32x8 atom), so D128 bf16 uses the narrow 16x16x16 path, not the flash regime.
• wide4 is the provider's analytic default, not the spec's. A bare spec with no flash knobs lands on L4 (WG=64); parity_unified_attention.py sets num_warps=4 explicitly to reproduce the shipped peak. HIPDNN_GFX942_FLASH_WIDE selects between them: unset / 4 → wide4, 0 → L4, 2 → wide2 (WG=64).
• The generic LDS support gate over-counts the flash path. supports_tiled_2d models the narrow footprint (P_lds present, K double-buffered, full Acc_lds), which would wrongly reject the shipped wide4 D128 config (whose P^\top is register-resident). The harness applies the gate only to the narrow path and lets the spec's __post_init__ + comgr arbitrate the flash configs, mirroring the provider.
• Graph replay is overhead-driven. The dispatcher's _recommend_graph_replay graphs decode (max_seqlen_q == 1) and short prefill (<= 768) — where host launch overhead is a large fraction — and ungraphs long prefill (kernel-bound). Toggle with HIPDNN_GFX942_2D_GRAPH / HIPDNN_GFX942_3D_GRAPH (both default-on).

Lever environment variables

parity_unified_attention.py exposes the spec knobs through environment variables (each is read by a small _*_enabled() / _*_setting() helper and echoed at startup):

variable	effect
HIPDNN_GFX942_FLASH_WIDE	flash-tile width for D128 fp16: unset/4 = wide4 (WG=256), 0 = L4 (WG=64), 2 = wide2 (WG=64)
HIPDNN_GFX942_NUM_WARPS	override num_warps in the flash regime
HIPDNN_GFX942_WAVES_PER_EU	set the waves_per_eu occupancy hint
HIPDNN_GFX942_CFV_STORE	opt into the experimental conflict-free V store path
HIPDNN_GFX942_CFV_STORE_SPLIT / HIPDNN_GFX942_CFV_CK_VLDS	sub-knobs of the conflict-free V store (default-on)
HIPDNN_GFX942_CFV	the legacy gather-fill conflict-free V diagnostic path
HIPDNN_GFX942_K_SLICED_RING	the experimental sliced-K ring for the conflict-free-V-store path
HIPDNN_GFX942_K_LDSSEQ	the CK Tile LdsSeq sliced-K variant
HIPDNN_GFX942_IGLP	the iglp_opt(1) scheduling hint
HIPDNN_GFX942_KV_CACHE_POLICY	the KV-cache load policy (default stream)
HIPDNN_GFX942_Q_DIRECT	load Q directly from global (skip the Q LDS stage)
HIPDNN_GFX942_GLOBAL_LOAD_LDS_K	use direct global→LDS loads for K
HIPDNN_GFX942_Q_MAJOR_GRID	transpose the launch grid to query-major

Troubleshooting

• head_size = 256 is off by default. There is no tiled d256 path on gfx942, so the dispatcher falls back to the scalar kernel, which is much slower than flash, fails the tolerance on some shapes, and is slow enough at S2048 to stall graph capture (it can look like a hang). The d256_disabled group is skipped unless you pass --groups d256_disabled (and expect failures / slowness). A proper tiled d256 implementation is needed to revisit it.
• Flash-ineligible shapes fall back to Torch's default SDP backend. In final_shapes_check.py, non-square causal shapes (seqlen_q != seqlen_k) and d256 are rejected by AOTriton flash, so those rows are timed against Torch's default SDP backend (marked default) and excluded from the flash win-rate. For non-square shapes the fp32 reference uses a bottom-right-aligned causal mask while Torch's is_causal=True is top-left-aligned, so those rows reflect slightly different masking and are reported for context only.
• Full sweeps are slow. final_shapes_check.py runs the full set × 4 timing cells × --reps × --iters, plus the fp32 reference on large (up to S8192) shapes — tens of minutes on a single MI300X. Scope it with --groups, --iters, and --reps. Cells whose spread exceeds ~10% are flagged noisy.
• SKIP (unsupported on gfx942). parity_unified_attention.py raises NotImplementedError (and skips) when a narrow-path shape fails the supports_tiled_2d LDS gate; flash-path shapes bypass the gate and are arbitrated by comgr instead.

How the harness maps a dense SDPA problem onto the paged kernel

Each batch element becomes one sequence with (query_len, kv_len) = (seqlen_q, seqlen_k); the per-sequence block table is a contiguous, non-overlapping run of block_size = 64-token cache blocks, so the KV working set is genuinely per-sequence-distinct. Inputs are filled uniform(-0.1, 0.1) to keep the softmax accumulation in a numerically friendly range. The decode group uses a length-1 query attending to all keys (is_causal=False, identical to causal at q == 1); prefill is causal. The launch grid ((kv_heads, total_num_q_blocks, 1), or query-major under HIPDNN_GFX942_Q_MAJOR_GRID) and block = (64 * num_warps, 1, 1) are recomputed exactly as the production dispatcher does, so the example exercises the same build

• launch plumbing as the provider.

Performance

This example is a correctness-and-latency harness, not a published benchmark: absolute timings depend on the GPU, ROCm/torch versions, clocks, and thermal state of the box, and are not benchmarked into this folder. Run final_shapes_check.py to produce the CK-vs-Torch eager/graph table for your own machine; the script prints, and gates on, the correctness result for every shape.