[BlivionIaG]

Purpose

Follow-up 1 from the exploration plan in #28: the D=128 prefill kernels.

• Based on the current rdna_extras (74315f4), which has perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 and perf(rocm): FA-RDNA2 decode kernel for spec-decode verify and mixed batches #29 merged and reworked. It is independent of perf(rocm): FA-RDNA2 decode, scores on every wave and 4 barriers per tile #33 and merges cleanly with it.
• Draft: not run on gfx1030 yet. It was verified offline (see Test Result). The Test Plan is for the agent with V620 access.

Problem

D=128 prefill (Qwen3/Llama-class models) ran the _short kernel below 4096 keys. Longer sequences used the per-head general or split-K kernels. The _short kernel (BR=32 × BC=32, 256 threads) has three costs:

1. Serial softmax. The online softmax runs row by row, in one wave (32 rows, each with two dependent 5-step ds_bpermute reductions), while the other 7 waves wait at the barrier.
2. O in LDS. sO is read-modify-written every tile (rescale plus P·V accumulate).
3. LDS footprint. About 45 KB per CTA, so only about 2 CTAs fit per WGP.

The D=256 GQA prefill kernel already fixed all three for D=256 in #28: O in registers, an in-register softmax with 16-lane shuffles, and one barrier fewer. It is the validated production D=256 path: 1.63-1.74× on unchunked prefill in #28's gfx1030 A/B.

Change

1. fa_prefill_paged_varlen_gqa_kernel is templated on HEAD_DIM (128/256).
   • D=128 keeps the same mapping: 16 query rows × 16 keys per tile, one score per thread, and 8 O dims per lane in registers.
   • LDS is about 14 KB per CTA (sQ 4 KB, sK/sV 8.5 KB, sP 1 KB).
   • All static_asserts hold unchanged.
2. Odd GQA groups and MHA.
   • A HEADS_PER_CTA=1, BR_STEP=16 instantiation keeps 16 rows per CTA from a single head.
   • The op now dispatches <D, 2, 8> for even groups and <D, 1, 16> otherwise.
   • The op name and schema are unchanged.
3. Routing.
   • RDNA_ATTN._forward_prefill sends every prefill to this kernel while VLLM_FA_RDNA2_GQA_MODE=subgroup, the default.
   • Any other value (e.g. off) restores the previous dispatch (_short / split-K / general), for A/B runs.
   • Unchanged: D=256 with an even group is bitwise identical.
   • Newly routed: D=128 at any group size, and D=256 with an odd group (none of the fork's D=256 models have one).
4. Benchmark. benchmarks/kernels/benchmark_fa_rdna2_prefill.py times gqa, short, varlen and splitk on prompts (512-8k), 4 × 1k, and chunks behind prefixes (1k @ 8k, 2k @ 16k). It prints the max output difference vs gqa. Every kernel is its own op, so one build covers the A/B.

Not in this PR:

• Retiring _short / varlen_128 / splitk_128. That waits on the measurements.
• A BC=32 variant with 2 keys per lane (fewer barriers per key).
• perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28's idea C: split-K for this kernel when a short chunk sits behind a long prefix, where there are few CTAs.

Not a duplicate: the only other open FA-RDNA2 PR is #33, which covers the decode kernels.

AI assistance: written with Claude Code. Per AGENTS.md, a human must review every changed line and run the GPU validation below before this leaves draft.

Test Plan

For the gfx1030 agent:

1. Build _rocm_C and check fa_rdna2.cu for new warnings.
   • Offline numbers: 85 VGPRs (D=128) and 87 (D=256), no spills.
   • The D=256 even-group instantiation is identical to rdna_extras.
2. Kernel tests.

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

   New or extended:
   • test_prefill_gqa_multiseq[*]: D=256 G=6/3, D=128 G=4/7/1 (MHA), dense and interleaved layouts.
   • gqa128 cases in test_prefill_chunked_offset and test_prefill_sliding_window.
   • test_forward_split_decode[128-*] now reaches this kernel through forward().
3. Kernel A/B. One build covers it:

   python benchmarks/kernels/benchmark_fa_rdna2_prefill.py                               # D=128: 32/8, 28/4 (G=7), 16/4; bs 16
   python benchmarks/kernels/benchmark_fa_rdna2_prefill.py --head-size 256 --heads 24/4 12/4 --block-size 784
   

   • Expect gqa to beat short below 4k keys, which is the point of this PR. The 2-4× estimate in perf(rocm): FA-RDNA2 causal tile skipping, in-place output, parallel GQA softmax #28 assumed the serial softmax is ≥50% of a _short tile; measure it.
   • Report where gqa loses, most likely short chunks behind long prefixes, where splitk has more CTAs. If it does lose there, route those shapes back to splitk in _forward_prefill: a max_seqlen_k / chunk-size rule, like e0112c5 did for split decode.
   • max|diff| vs gqa should be around 1e-3.
4. In-model on a D=128 model (e.g. Qwen3-8B or Llama-3-8B), with RDNA_ATTN and VLLM_FA_RDNA2_GQA_MODE=subgroup vs off:
   • TTFT for 1k/4k/16k prompts, and chunked-prefill throughput.
   • Greedy coherence.
   • gsm8k parity.
   • Qwen2.5-7B (G=7) exercises the one-head instantiation.
5. D=256 production (Flash-Next / serve_gfx1030_full.sh): spot-check that prefill TTFT is unchanged. The kernel for even groups is the same code.

Test Result

Offline (done):

• CPU emulation of the kernels, with the current rdna_extras as the baseline. The harness uses fibers with block barriers, wave32 shuffles, guard-paged buffers, and NaN-poisoned outputs and dynamic LDS. All 56 cases pass. The batches cover multi-sequence prefill, prefix chunks, single-token and zero-length requests, windows 1/17/100, custom scale, dense and hybrid layouts, and block sizes 13/16/32/784:
  • D=256, even group (6/1, 8/2): bit-identical to the current kernel in all 14 cases.
  • D=256, odd group and MHA (3/1, 4/4): within 8.9e-4 of an fp64 reference.
  • D=128 (8/2, 16/4, 7/1 with G=7, 4/4 MHA): within 7.4e-4 of the reference, and never worse than the current _short and general kernels on the same inputs. Those reach up to 1.1e-3.
  • Races: a randomized fiber schedule is bit-identical to the sequential run.
• gfx1030 device build (clang 18): no warnings. No spills or scratch. D=128 uses 85 VGPRs (occupancy 10); D=256 even groups are unchanged at 87 / 58 SGPRs / 10.
• Routing: checked on CPU by running RdnaAttentionImpl.forward with reference stub kernels on 47 batches.
  • In the default mode, D=128 and odd-group prefill go to gqa and outputs match the reference.
  • With VLLM_FA_RDNA2_GQA_MODE=off the previous dispatch returns.
• Lint: ruff and mypy report no new findings (the files' existing ones remain).
• Merge check: clean with rdna_extras and with perf(rocm): FA-RDNA2 decode, scores on every wave and 4 barriers per tile #33.
• Benchmark script: checked on CPU against stub ops.

Re-checked on rdna_extras 74315f4 (2026-10-01). The merge is clean. The base did not touch FA-RDNA2, so the sources are byte-identical to the version validated above.

Cost of the same work under the current _short kernel and this PR's gqa<128, 2, 8>: 32 query rows × 2 q-heads (G=2) behind a prefix. Executed counts come from the emulator; static counts come from the gfx1030 assembly of the main KV loop.

	_short (base)	gqa, D=128 (PR)
tile	32 rows × 32 keys, 1 head	16 rows (8 × 2 heads) × 16 keys
barriers per tile	4	4
shuffles per tile, busiest wave	320, all in one wave (the serial softmax)	8
shuffles per tile, all waves	320	64
per 32 × 32 (row, key) block: barriers / critical-path shuffles	4 / 320	16 / 32
one CTA walking 4608 keys: busiest-wave shuffles	46,080	2,300
LDS per CTA	~45 KB	~14 KB
main-loop LDS reads / instructions (static)	97 / 767, plus a 32-iteration softmax loop	43 / 483

The PR trades 4× more barriers per unit of work for a 10× shorter serial shuffle chain and about 3× less LDS. Those barriers are cheap because the work between them is balanced across waves. Only measurements can say whether the trade wins on every shape (Test Plan step 3).

Scope on the current launchers: none of them serves a D=128 model; Flash-Next and the 27B are D=256. For D=256 with an even group this PR is bit-identical, so it only changes D=128 models.

On gfx1030: pending, to be filled in from the plan above.

---

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 (offline only so far; GPU results pending)
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

🤖 Generated with Claude Code

https://claude.ai/code/session_018dwtAVjZdFLPVRucCFXBzY

[BlivionIaG]

Validated on gfx1030 (4× V620) and integrated into rdna_extras as 284fbad0a (kernel) + ef8f0a283 (validation artifacts). Re-authored to BlivionIaG <kev29lt@gmail.com> per fork policy (all commits in this repo are authored by the fork owner), so this PR is closed as integrated rather than merged via GitHub.

Results:

• Flash-Next is D=256 even-group ⇒ PR is bitwise-identical (cross-build max_abs_diff = 0.0) ⇒ no-regression proof, 8-cell matrix clean (TP4, FPP, prefix-caching, FA-RDNA2, MTP 0/2, c1/8, 1k/16k).
• D=128 kernel A/B: gqa is 4–10× faster than _short/varlen/splitk at every shape, incl. the flagged chunk-behind-prefix case (chunk 2k@16k: 126 ms vs 1341 ms splitk); max|diff| ≤ 9.8e-4.
• 72/72 FA-RDNA2 kernel tests pass; 0 spills, wave32, all four instantiations present.
• Routing toggle (subgroup vs off) is a no-op for D=256, as designed.

Summary: bench_results/2026-10-02_pr34-fa-prefill-d128/SUMMARY.md.