[BlivionIaG]

Purpose

Retarget of draft #33 onto current rdna_extras (1dc3bac4cd, which contains the v0.31.0 port at 8abf4c97 plus the later INT4 PLE merge). The per-head split-K decode kernels still score each KV tile on one part of the CTA. This carries the #33 schedule onto the v0.31 fa_rdna2 source:

• Scores on every wave. D=256 uses 8 lanes per key (32 dims, 3 xor shuffles). D=128 uses 4 lanes per key pair, with the odd-key uint4 rotation. Q stays in registers. Packed fdot2 only: no WMMA, no MFMA.
• One LDS exchange for the tile max. l accumulates per wave and is summed once per CTA.
• Double-buffered page lookups. A taken tile has 4 barriers (K/V loaded, maxima exchanged, P written, end of tile). A fully masked tile skips the P barrier on every thread. m_new is CTA-uniform, so that barrier is not divergent.
• P·V. Each thread owns 2 dims over half the tile's keys and reads them as half2. The two halves merge once per CTA.

The fp8 and int8 decode wrappers share these bodies. The GQA decode kernel is unchanged. Every fa_rdna2 launch still passes the current stream.

Not a duplicate: this is #33 itself, moved off the archived rdna_extra/v0.28.0 base (a6ab43cf). No other open PR rewrites these per-head decode bodies.

AI assistance: re-derived with Grok 4.7 in Cursor. A human must review every changed line and run the GPU checks below before this leaves draft.

What changed vs the old #33 diff

• Re-derived against the v0.31 fa_rdna2.cu rather than replaying 40931b3d. Those kernel bodies match the pre-port source aside from clang-format. The port's behavioral delta kept here is at::cuda::getCurrentCUDAStream() on every launch.
• Same schedule as perf(rocm): FA-RDNA2 decode, scores on every wave and 4 barriers per tile #33: all-wave scores, double-buffered lookups, 4 barriers per taken tile, Q not staged in LDS. Host dynamic smem drops the Q tile and the block-reduce broadcast slot (THREADS/32 floats, not + 1).
• k_end <= 0 does not enter the P·V loop. The old for (k = 0; k < k_end) already did not run when k_end was negative.
• Not replayed: bench_results/2026-10-0*, the rocprof scripts, and the "no merge (QSA path)" retest commit. Those numbers were measured on the old tree. This PR does not claim tok/s.
• The decode benchmark script is kept. It still calls fa_rdna2_decode_paged directly, so it measures the per-head kernel with VLLM_FA_RDNA2_GQA_DECODE unset.

Test Plan

Re-measure on the 8x V620 box, ROCm 7.14.0, this head against current rdna_extras. No numbers from the old base carry over.

1. Build _rocm_C for gfx1030 (wave32). Confirm fa_rdna2.cu has no new warnings, and check the per-head decode instantiations for spills and LDS occupancy.
2. Correctness. Outputs are not bit-identical to the previous kernels because the fp32 sum order changed; these tests use tolerances.

   pytest -v tests/kernels/attention/test_fa_rdna2_writer_layout.py \
             tests/kernels/attention/test_fa_rdna2_shape_sweep.py
   

   Decode cells: test_decode_writer_layout, test_decode_empty_rows_are_zero, test_custom_scale_into_out, test_decode_paged_sliding_window, test_decode_multi_token_queries, test_forward_split_decode.
3. Decode tok/s. Leave VLLM_FA_RDNA2_GQA_DECODE unset for the kernel benchmark.

   python benchmarks/kernels/benchmark_fa_rdna2_decode.py --batches 1 2 3 8 32
   python benchmarks/kernels/benchmark_fa_rdna2_decode.py --head-size 128 --heads 32/8 28/4 --block-size 16 --layout dense
   python benchmarks/kernels/benchmark_fa_rdna2_split_decode.py --batches 1 2
   

   In-model: confirm the launcher actually dispatches fa_rdna2_decode_paged before reading tok/s (the old retest found Flash-Next on the QSA path). Record decode tok/s for the cells that hit this kernel.
4. If it wins, porting the score layout to fa_decode_paged_splitk_gqa_kernel_256 is a follow-up. It is not in this PR.

Test Result

No GPU in this environment. No tok/s, spill count, or occupancy number is claimed.

CPU-side, on this head:

• Lane coverage: every D=128 key pair and every D=256 key covers each uint4 chunk once. Each 8-lane ds_read_b128 group has disjoint LDS banks.
• Numeric model of the schedule (wave32 xor-shuffles, online softmax, split-K combine) against a one-shot fp64 reference: 300 cases (D=128/256, tile edges, windows, splits 1/3/8, empty sequences). Worst absolute error 1.8e-15.
• LDS race check between barriers, including the double-buffered lookups, the sP/sW writers, and the sO/sK alias: 216 schedules, no conflicting accesses.
• The masked-tile predicate is CTA-uniform, so the inner P barrier is not divergent.
• Host clang++ -std=c++17 -fsyntax-only of both kernels, instantiated for fp16 and int8, with HIP and torch stubbed out. This is not a gfx1030 device build. There is no ROCm toolchain here, so the real _rocm_C build was not run.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results (CPU checks only; GPU re-measure is open)
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[BlivionIaG]

Branch updated (rebased onto current rdna_extras, incl. the bench/retest commits) and validation evidence recorded.

• ISA: all instantiations present, 0 SGPR/VGPR spills, wave32. Tests: 62/62 FA-RDNA2 (+ split-decode cases), max-diff vs fp64 within fp16 rounding.
• Kernel A/B: fa_rdna2_decode_paged -18%..-31% at the D=256 6/1 geometry; split-verify q=3 1.4-8.8x.
• In-model: Flash-Next shows parity because it dispatches QSA (Qwen4ExpQSAAttention) - no fa_rdna2_decode_paged at all, so this change cannot move it. The model that does dispatch it is the 27B AWQ; its clean A/B was blocked by a chassis PCI SERR and needs a re-run with VLLM_RDNA_AR=0 in a quiet session.
• Verdict: keep draft until the 27B A/B (or a low-variance measurement) shows an end-to-end win. The kernel change itself is sound.