perf(gqa): fused v5 prefill for sinks and sliding windows, at 44% fewer instructions per KV tile

[BoarQing]

Summary

Two commits against the GQA prefill path, both on gqa_flash_prefill_v5 (the d==64 warp-private kernel):

1. perf(gqa): fused flash prefill for attention sinks and sliding windows — teaches v5 attention sinks, smooth softmax and sliding windows so gpt-oss no longer falls back to the decomposed path for its 12 sliding layers and 24 sink layers. HAS_WINDOW is a template parameter rather than a runtime test because the kernel runs at the VGPR ceiling: carrying the window's per-row state live in the full-attention instantiation pushed M_TILES=2 from 120 to 180 bytes/lane of scratch and cost that path 21% at deep KV.

2. perf(gqa): cut v5 prefill KV-loop instruction count 44% — the loop was bound by instruction issue, not occupancy, bandwidth or the matrix units: 1,734 instructions per KV tile to retire 32 WMMA, so the matrix units saw 2% of the stream. Three changes take that to 967 on the 99.6% of tiles that need no causal mask: native exp2 in the softmax, the softmax denominator's cross-lane reduction deferred out of the loop into the epilogue, and a mask-free/masked tile split via a tag-dispatched tile body.

Measured

Interleaved A/B on gpt_oss-sink sq=512 past=8192, each arm a separate build, 3 rounds:

arm	vs parent
native exp2 alone	-0.5%
deferred reduction alone	-8.9%
both	-12.4%
all three	-21.4%

exp2 pays nothing on its own because the reduction chain it sits behind is what the loop waits on; it only becomes visible once that chain is gone. The tile split is likewise not separable from the restructure it needs — the same lambda without the split costs 32% on every d=64 full-attention case, because the split is what lets the compiler keep a tile's live values in registers (in-loop scratch traffic 52 ops → 6).

In-model, 16k-token gpt-oss-20b prefill, RGP capture:

	before	after
v5 full-attention dispatch mean	5,989.7 us	5,207.7 us
share of capture window	25.3%	23.0%
TTFT	9,653 ms	9,172 ms (-5.0%)

Correctness

test_gqa_prefill covers 26 cases against an fp32 CPU reference: all four sink/smooth/both/none modes, windows at 1/100/128/4096, past at 0/512/1000/8192, d at 64 and 128, and the two d=128 cases that assert v5 declines so the runtime falls back rather than silently ignoring a sink or a window. All 26 pass, worst relL2 5.74e-4, unchanged by commit 2 — the reordered summation in the deferred reduction is not observable at this tolerance.

Across the 24 cases the suite also times: 21 faster, 1 tie, no regressions. d=128 routes to v7 and is unaffected (1.00x, 1.01x).

The new mask-free instantiation is exercised by every full-attention case with past > 0, and by the interior q-tiles of the past = 0 cases.

Notes for review

• The three A/B switches used to select these changes are not in the diff; they were collapsed to their enabled path before this PR. The collapse is provably free: the device image is byte-identical to the switched build across all 443,062 lines of gfx1151 ISA, differing only in the __hip_cuid symbol, which is a hash of the translation unit and moves when a comment does.
• No new diagnostics, env vars or debug surface. The ablate bitmask and the hip_gqa_flash_prefill_v*_set_ablate exports visible in this file are pre-existing on main from [feat]:add opt for normal GQA #438; these commits neither add to them nor remove them. Two getenv readers I had used to drive them during development were taken back out before review, and the device image is byte-identical either way, since they were host-side only.
• The RGP one-shot capture fence used to take the in-model numbers is deliberately not in this PR; it is unrelated profiling infrastructure and belongs in its own.

What is deliberately left on the table

Re-profiling after the fact says the phase order inverted: the value GEMM is now the largest phase at 36.5% (was 19.8%) and softmax has fallen to 24.2% (was 39.1%). The loop is now only 28-41% issue-bound, so instruction count has largely stopped being the lever. Two candidates remain, both aimed at the LDS pipe and both sized from the ISA but not built: storing V transposed in LDS to turn 128 ds_load_u16 into 16 ds_load_b128 (worth -5.4% by instruction count, 3.6-7.1% by an LDS-pipe model, against an unquantified bank-conflict risk), and DPP row_xmask for the remaining 64 ds_bpermute (verified available on gfx1151, but it does not fold into the consuming v_max, so it is instruction-neutral and worth 0% or 4.6-9.1% depending on which LDS throughput assumption holds). Together roughly 2% of prefill at the optimistic end.

Test plan

• test_gqa_prefill — 26/26 pass, fp32 reference, worst relL2 5.74e-4
• test_gqa_prefill timing sweep, 24 cases, 20 and 40 iters, no regressions
• Byte-level ISA equivalence of the shipped source against the measured build
• In-model 16k gpt-oss-20b prefill, RGP capture, TTFT -5.0%
• CI on this branch

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	4.8668E-07	64	0
GroupQueryAttention_seq256	25.2366	2621440	0
MatMulNBits_o_seq128	259.905	368640	0
QMoE_seq128	34.957	368640	0

Threshold: 0.01 | Run: 3632 - Commit: 9813372

[github-actions]

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.50	6.19	361	3	1245
GroupQueryAttention_seq128	4397.78	1.67642	10	6	311
matmul_down_seq128	519.55	2.33	75	3	352

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.54	45.32	365	3	15590

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.54	9.83	363	3	15761

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	186.8	80.4	1.33	13.53
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	320.1	40.4	1.22	6.43

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	141	39.4

Run: 3632 - Commit: 9813372

[BoarQing]

http://xsjnts.xilinx.com/ROCm/compare?base_url=http%3A%2F%2Fxsjvitisaijenkins%3A8080%2Fjob%2Fdaily_hip_ep_oga%2F1633&base_job=&base=&curr_url=http%3A%2F%2Fxsjvitisaijenkins%3A8080%2Fjob%2Fdaily_hip_ep_oga%2F1679%2F&curr_job=&curr=

One last room