[Kernel][Feature][Dialect][DSL][Perf][Doc][Test] Add gfx120x RDNA4 wave32 ops

[dimitri91209]

Summary

gfx120x (RDNA4, wave32, WMMA) kernels for flash attention, GEMM, quantization, SwiGLU, RoPE, and AdaLN. The host does not copy Q, K, or V to round K or the head dimension up to a tile. Lanes past the real shape are zero.

Versus main (43ec5da9): 106 files, +38582 / −27. Tip 21a11efb. The quant paged helper is _gfx120x_quant_paged. Every positive head is padded to a multiple of 32. A head whose KV prefetch does not fit in 64 KiB is split into LDS chunks. The value head can differ from the query head, including a ragged paged cache on bf16, fp16, fp8, and int8. Quant paged and varlen paged zero-extend a value head that is not a multiple of 32 and crop the output. e5m2 QK and PV both use the e5m2 matrix product. Quant QK-norm keeps each normalized pair in registers. Logit softcap is tanh(score * scale / softcap) * softcap on bf16, fp16, fp8, and int8, including varlen and paged. Dropout probability 0 leaves that kernel unchanged. A positive probability uses an explicit Philox seed. The accuracy check is one step of the output format at that output's size. Near 1.0 the step is 0.00098 for fp16, 0.0078 for bf16, 0.125 for e4m3, and 0.25 for e5m2. The step doubles when the value doubles. Int8 is half a code, the max absolute value divided by 254. Main's flash-attention sentences stay. gfx120x notes are added after them. The function that contains @flyc.jit asserts get_rocm_arch().startswith("gfx120") before that definition. There is no gfx120x_arch.py. On gfx120x, flydsl_flash_attn_func returns into flash_attn_gfx120x_host.flydsl_flash_attn_func. That host stores attn_mask in bias and launches dense, varlen, paged, fp8, and int8. The value head, softcap, and dropout probability are cache keys. Paged decode, MLA, linear attention, and quant QK-norm are separate gfx120x modules. Decode takes the same arguments as main, including a sliding window. MLA reads the latent cache through the WMMA paged kernel. Linear attention is the prefix recurrence. Quant QK-norm is RMSNorm, tail RoPE, then fp8 or int8.

#1223 stays a draft and stays targeted at main. Shared compiler bugfixes that are not gfx120x kernels are on #1240. This branch does not carry that PR's CDNA5 tensor2tdm null check.

There is no packed-int4 flash attention. iu4_gemm is the int4 GEMM, and the atom is K=32 only (wmma_i32_16x16x32_iu4). Attention dtypes are bf16, fp16, fp8 e4m3, fp8 e5m2, and int8.

fx.Int64 is the index and offset type. fx.Index is deprecated, same as main. The disk cache stores compiled IR only. Source IR is not persisted.

Still copied on the host, on purpose: the nmajor wrapper around the shared f16 GEMM, which requires a tile multiple, and a direct call to the in-register SwiGLU kernel when K or the FFN width is below 16. The product SwiGLU host does not send those shapes there.

Bug fixes

File: kernels/attention/flash_attn_gfx120x_host.py (_pad_narrow_value_cache)
Bug: Quant paged and varlen paged launched a value head that is not a multiple of 32 at the query tile, so heads 40, 48, and 80 read past V.
Fix: The value cache is zero-extended to the WMMA stride and the output is cropped back to the caller width. The uniform bf16 paged path uses the same helper.
Notes: e4m3 and e5m2 pass those heads at one step of the fp8 format. int8 no longer overflows.

File: kernels/attention/flash_attn_gfx120x_host.py (apply_sliding_window)
Bug: A causal window was measured from q_row. A decode query (Sq=1) kept the start of the cache.
Fix: The band origin is q_row + (Sk - Sq) when the mask is causal, and q_row otherwise.
Notes: Equal lengths are unchanged. A decode window (left=4, right=0) matches the bottom-right mask at cosine 1.0 and does not match the old key range.

File: kernels/attention/flash_attn_gfx120x_host.py, flash_attn_gfx120x.py, flash_attn_fp8_gfx120x.py
Bug: Ragged varlen ALiBi was the top-left of a bias built at the max lengths, so a shorter sequence used |i + max_k - max_q - j|.
Fix: A 2D ALiBi is copied from the bottom-right slice of each sequence. A per-head ALiBi tensor sets bias_bottom_right, and the kernel adds (max_q - sq, max_k - sk).
Notes: A per-head user mask stays top-left. Scalar ALiBi plus a 2D user mask is expanded separately, then added. Lengths 8/20 against 24/12 match a per-sequence fold at cosine 0.99999.

File: kernels/attention/flash_attn_gfx120x_host.py (add_alibi_scores), flash_attn_gfx120x.py, flash_attn_fp8_gfx120x.py
Bug: A ragged varlen call with both a per-head user mask and per-head ALiBi had one bias tensor. The mask is top-left and ALiBi is bottom-right, so the ALiBi term used the max lengths.
Fix: That mix passes the slopes into the kernel. add_alibi_scores adds -slope * |i + sk - sq - j|, scaled like the user bias, and adds 0 on a pad column. Equal lengths, pure ALiBi, and a scalar slope plus a 2D mask still fold into one bias.
Notes: Lengths 8/20 against 24/12 match a per-sequence fold. fp8 and int8 use 1/(sm_scale * q_descale * k_descale). The AOT export passes the slopes pointer.

File: kernels/attention/flash_attn_gfx120x_host.py (flydsl_flash_attn_quant_paged_func)
Bug: Quant paged attention that is not packed-varlen kept a shared ALiBi tensor when seqlen_k differed across the batch, so the kernel read column 0 instead of max_k - sk.
Fix: The host packs Q and calls the varlen-paged path, the same handoff the float paged host already uses. Uniform lengths stay on the dense paged path.
Notes: fp8 e4m3 and int8, seqlen_k 32 and 16, match a per-batch bottom-right fold and a bottom-right causal mask. Output stays bf16.

File: kernels/attention/flash_attn_gfx120x.py, flash_attn_fp8_gfx120x.py
Bug: Bias and ALiBi were added to unscaled WMMA scores. The contract is softmax(qk * sm_scale + bias).
Fix: bf16 divides the bias by sm_scale. fp8 and int8 divide by sm_scale * q_descale * k_descale.
Notes: Coverage bias, per-head bias, and ALiBi sections pass. Causal alignment stays bottom-right. It is not PyTorch is_causal when Sq != Sk.

File: kernels/attention/flash_attn_fp8_gfx120x.py, flash_attn_gfx120x_host.py
Bug: A quant split narrower than the KV tile counted those columns twice. A dense split whose head was shorter than the tile wrote the tile stride into the output. The split workspace was uninitialized. A quant varlen shared bias record was too small, and a per-head mask was squeezed away.
Fix: Split columns outside [kv_lo, kv_hi) are masked. The combine writes a tile buffer, then copies the real head. Workspace m is -inf and l and o are zero. The varlen bias record covers the packed rows. [H, Sq, Sk] with H > 1 stays per-head.
Notes: Split-K versus the same kernel stays at cosine ≥ 0.99. Odd head 80 is in the coverage test.

File: kernels/gemm/rdna4_int8_linear_fused.py
Bug: Fused int8 activation quant clamped to [-127, 127] and rounded with ±0.5, so -128 never appeared and ties rounded away from zero.
Fix: roundeven, then clamp to [-128, 127].
Notes: -127.5 and -128 store -128. 2.5 stores 2. -2.5 stores -2. Matches the torch reference exactly.

File: kernels/gemm/rdna4_scaled_mm_fp8.py, rdna4_scaled_mm_fp8_fused.py, rdna4_w8a16_linear.py, rdna4_int8_linear.py, rdna4_int8_linear_fused.py
Bug: K=0 with M>0 and N>0 still launched, and an out-of-range lane reloaded byte 0 of an empty allocation.
Fix: Return zeros of shape [M, N]. Where the host has a bias, add it after. Scales are still checked.
Notes: Covered for scaled-mm, its fused and multi paths, W8A16, int8 linear, and both fused int8 paths.

File: kernels/quant/rdna4_asym_w4a8.py
Bug: Asymmetric mode used (max(|v|) - min(v)) / 15, so a larger negative peak widened the scale and collapsed the codes. Codes and the correction were chosen from the f32 ratio, then s_rel was cast to e4m3.
Fix: Mode 2 uses (max(v) - min(v)) / 15. When s_rel is e4m3, that cast is rounded first and the codes and correction use the stored value.
Notes: Modes 0 and 1 still use abs-max. For [-1, -0.2, 0.1] the codes are [0, 11, 15], including after the e4m3 round-trip.

File: kernels/quant/rdna4_svdquant_w4a4.py
Bug: Int4 codes were roundeven(x / scale_f32) while the stored scale was bf16 or fp16 of that value. The linear then cast the scale again.
Fix: Round the scale through the stored dtype before the reciprocal. svdquant_w4a4_linear quantizes in wscales.dtype and does not cast afterward.
Notes: -0.2138671875 with scale 1/7 is code -2 on the stored bf16 scale and would be -1 on the f32 scale. The kernel stores -2.

File: kernels/quant/rdna4_quantize_int8_tensorwise.py
Bug: A caller-supplied scale was replaced by its absolute value, so a negative scale no longer flipped the codes.
Fix: inv is the reciprocal of the signed scale. The magnitude is still clamped at 1e-30.
Notes: scale=-0.5 matches round(x / -0.5). Abs-max quant is unchanged.

File: kernels/quant/rdna4_stoch_fp8.py
Bug: Stochastic rounding pre-rounded through fp16.
Fix: The value stays f32 (fx.Float32(raw)).
Notes: test_stoch_fp8_uses_f32_not_fp16 passed. The cast still stores no scale.

File: kernels/quant/rdna4_int8_convrot.py
Bug: The int8 Hadamard quant stored only the logical K and dropped the rest of the orthogonal group.
Fix: The kernel loads logical_k columns and zero-fills the rest. It stores the full Hadamard width. The host returns [..., k_pad].
Notes: test_int8_linear_convrot_odd_k_stays_in_kernel passed. The FWHT helper still returns the caller's K.

File: kernels/gemm/rdna4_w8a16_linear.py, rdna4_scaled_mm_fp8_fused.py, rdna4_int8_linear_fused.py
Bug: A partial N tile could load a scale past N. A fused LoRA row could be negative when the column was out of range.
Fix: The scale column is clamped to 0 when skip_bounds is false. The LoRA row is 0 when the column is out of range.
Notes: The existing W8A16 and fused GEMM tests passed in the full suite.

File: tests/python/examples/aot_gfx120x_example.py, kernels/attention/flash_attn_gfx120x_splitk.py
Bug: The AOT pack skipped MXFP GEMM and quant, SiLU, FP8 dequant, split-K, ALiBi, the bool mask, AWQ dequant, and int8 ConvRot dequant. AdaLN and RMS+RoPE were given tensor ranks the kernels do not index.
Fix: --all exports those kernels. AdaLN scale and shift are [rows, N]. RMS+RoPE frequencies are [rows, pairs, 2, 2]. Split-K exposes its jit launcher for export. The host wrapper is unchanged.
Notes: The broaden export compiled every one of those objects.

File: kernels/common/gfx120x_autotune_tables.py
Bug: The block-size table was a size gate left from an old measurement. The hosts already mask lanes past the real row.
Fix: The file is gone. AdaLN, RoPE, RMS+RoPE, int8 rowwise quant, and asymmetric W4A8 use 256 threads unless the caller passes a block.
Notes: Flash attention BLOCK_M and the scaled FP8 tile pick stay in their own hosts.

File: tests/kernels/oracles.py
Bug: The numeric references were split across nine modules.
Fix: They live in tests/kernels/oracles.py. The tests import that module.
Notes: The reference math is unchanged.

File: tests/kernels/test_gfx120x_odd_shapes.py
Bug: Product tests mostly used tile multiples of 64 and 128.
Fix: One sweep calls the hosts with no block override at other widths: AdaLN 96/160/384, RoPE and RMS+RoPE head 80 and 96, int8 rowwise K 100/300/700, scaled FP8 and e5m2 off the 64 tile, int8 13×27×41, MXFP8 17×23×96, attention seq 17/97/129 and head 72/80/96, and SwiGLU 7×24×40.
Notes: Each section is compared to the same reference the tile tests use. Cosine stayed at or above 0.98. The sweep does not assert a size-based block pick.

File: kernels/common/gfx120x_swiglu.py
Bug: swiglu_chunk stored packed row-major from data_ptr() into out, including when out was a non-contiguous slice.
Fix: A non-contiguous out is written through a contiguous temporary and copied back on the same stream.
Notes: buf[:, :H] matches silu(gate) * up. A contiguous out is still launched directly.

Tests

On the kernel commit, scripts/run_tests.sh was 3831 passed, 2074 skipped, 33 deselected, 1 failed. The failure is tests/unit/test_max_blocks_per_cu.py: the installed build-fly package does not have CompiledFunction.max_blocks_per_cu, which this tree's source does. FileCheck passed, including wmma_gfx120x.mlir and swmmac_gfx120x.mlir. After the arch check, routing guards, public entries, scaled FP8, and RoPE passed (26 tests). 7d961b1b keeps that behavior. 16f2607b passed the routing guards, the interface route, a dense causal cross launch, and test_gfx120x_attention_covers_implemented_ops (11 passed). That coverage test goes through the interface, including sliding window, varlen, paged, fp8, and int8. The full suite was not re-run. tests/kernels/test_gfx120x_attn_rest.py passed (9 tests). The wide latent decode covers bf16, fp16, fp8 e4m3, fp8 e5m2, and int8 at query width 576 and value width 512. Quant QK-norm also passes head widths 1, 40, and 128. The follow-up tests passed (30 tests) on the previous tip: register linear attention, quant QK-norm, one-partition MLA splits, CSR block tables and mixed-dtype decode, and an unaligned paged value head. This commit pads quant and varlen paged value heads that are not a multiple of 32. The e5m2 short test checks one step of that format and passed. An e5m2 grid of 70 shapes passed. Quant QK-norm passed 49 shapes, including heads 1, 17, 40, 65, 128, 256, and 512, and an empty row. Repeat calls are about 0.055 ms. The full suite was not re-run after this commit. Examples 01, 02, and 05 pass on this GPU. Examples 03, 04, and 06 abort because they emit CDNA MFMA instructions. Examples 07–22 are not in this tree. Tests skip unless the process arch starts with gfx120, the same way the gfx950 suites skip.

The earlier bug fixes are inside that suite. The attention cases, including the ragged mixed ALiBi and the ragged quant paged cases, are in tests/kernels/test_flash_attn_gfx120x_coverage.py, which names the section if it fails. The value-head pad, the e5m2 matrix product, and the register QK-norm are the checks named above.

[dimitri91209]

If anyone is looking at this, honestly just don't.

It'll get there eventually, the attention stack is close. Like really really lose to being a drop in for Rdna4 for pretty much everything that's industry standard. Jit should be able to build kernels by going through the attention host with args passed from the program host and build pretty much anything one needs for Rdna4.

Currently waiting on a true16 pr on LLVM before any of that works for qwen though. In order to avoid a lds round trip for attention.