[amangupta-tether]

Overview

Vulkan prefill flash attention for Strix Halo (gfx1151, RADV). This is the attention-at-depth follow-up to #393, rebased onto temp-11018 now that #393 is merged. It is one squashed commit (9e18341cd). Changes from review:

• The multi-row kernel zeroes V at the score path's mask threshold.
• The V^T scratch goes through the device-local capacity gate.
• The perf-logger query pool is sized for the FA sub-op marks.
• The contiguize pass defaults on for RADV RDNA3 only.

Prefill at depth is limited by FLASH_ATTN_EXT on Qwen3.8 27B and Flash-Next, and this stack speeds it up. All of it comes from halo-box/strix-llama.cpp, except one upstream fix:

• upstream vulkan: fix stale prealloc_y reuse across flash attention and soft_max ggml-org/llama.cpp#29591: fixes stale prealloc_y reuse across flash attention and soft_max, which the contiguize pass below relies on.
• contiguize strided f16 K/V for prefill FA. The KV-cache view reaching FA has head-interleaved rows, and the cm1 coopMatLoads run about 2x slower on that stride. A strided copy into the existing dequant-once scratch (vulkan : dequant q8_0 KV once in coopmat1 ggml-org/llama.cpp#25494) makes K/V per-head contiguous, together with the scratch's device-local capacity gate. It is on by default for RADV RDNA3, where it was measured; GGML_VK_FA_KV_CONTIG=1 turns it on for every device and =0 turns it off.
• The contiguize pass is billed on its own perf-logger line.
• cm1 FA: hoist the P-fragment load out of the hsv_tile loop, and store Psh query-major so the GEMM2 A load vectorizes.
• wave32 pin for cm1 FA on RDNA3, applied when hsv <= 128 and at least 32 rows (GGML_VK_FA_WAVE32=0 turns it off). It does not engage at Qwen3.8's head size of 256.
• multi-row coopmat1 FA with a V^T scratch (flash_attn_cm1_mr.comp, dequant_f16_transpose_vt.comp). One workgroup of eight wave32 subgroups covers 64 query rows (H heads x 64/H tokens sharing a KV head). Q stays in registers and every 16-key K/V tile serves all 64 rows; cm1 pins Br at 16. V comes from a per-head V^T copy, so a P x V fragment is one contiguous 32-byte run per lane. It is on by default for head size 256 on RADV RDNA3. GGML_VK_FA_MR=0 turns off the multi-row kernel, and GGML_VK_FA_VT=0 turns off V^T (and with it the multi-row kernel).

The squashed commit credits Nathan Wilson (the strix-llama.cpp work) and François-Xavier Gsell (upstream ggml-org#29591) as co-authors and lists the strix commits it ports. Two strix FA commits are left out: NaN-safe V staging (1b899d039), which needs the selected-key attention path that fabric does not have, and ggml-org#29988, which is for the scalar FA MMQ path that fabric disables.

Additional information

Results

GMKtec EVO-X2: Ryzen AI MAX+ 395 / Radeon 8060S (gfx1151), Mesa 26.0.8 RADV. llama-bench -ngl 999 -fa 1 -b 2048 -p 2048 -n 0; Flash-Next adds -lzm on. d0 and d8192 use -ub 2048 -r 2. d32768 uses -ub 512 -r 1, because -ub 2048 at 32K trips the amdgpu compute-ring timeout on Vulkan with and without this PR. Before and after are clean builds, interleaved twice.

Vulkan pp2048 (t/s), #393 vs this PR:

model	d0	d8192	d32768
Qwen3.8 27B Q4_K_M	306.5 -> 309.1	244.0 -> 302.1 (+24%)	106.7 -> 266.4 (2.50x)
Qwen3.8 27B UD-Q8_K_L	325.3 -> 329.4	257.9 -> 321.3 (+25%)	100.9 -> 285.8 (2.83x)
Qwen3.8 Flash-Next UD-IQ4_XS	674.8 -> 687.9	536.3 -> 605.6 (+13%)	316.8 -> 380.2 (+20%)

For reference, strix-llama.cpp deb5e2f at d8192 (-ub 2048): 27B Q4_K_M 291, 27B Q8 315, Flash-Next 637. Both 27B models now beat it. Flash-Next still trails, because strix runs its full-attention layers as QSA selected-key attention, which is not part of this PR.

Where the gain comes from: this PR with one switch at a time (pp2048, t/s):

configuration	27B Q4_K_M d8192	d32768	Flash-Next d8192	d32768
this PR (default)	301.4	266.7	606.0	384.0
GGML_VK_FA_MR=0 (cm1 + V^T)	286.8	226.2	572.7	357.0
GGML_VK_FA_VT=0 (cm1, no V^T, no multi-row)	283.3	224.2	559.7	346.1
GGML_VK_FA_KV_CONTIG=0 (multi-row on the strided view)	298.0	265.3	607.3	388.0
GGML_VK_FA_KV_CONTIG=0 GGML_VK_FA_VT=0	242.2	126.2	551.9	330.3
#393	244.0	106.7	536.3	316.8

The contiguize pass does most of the work for the cm1 kernel. The multi-row kernel adds 5% at d8192 and 18% at d32768 on top, and it does nearly as well without the contiguize pass because it reads the strided view efficiently. With every switch off, the cm1 shader changes (P-fragment hoist, Psh layout) remain, and they account for the 107 -> 126 at d32768.

Fabric-specific adaptations

• Flag bits: fabric's FA Flags already use 16 = USE_SPARSE and 32 = turbo-quant QJL_FULL_PROJ. The new bits are 64 (V_TRANSPOSED) and 128..1024 (multi-row, host-side selection and schedule knobs), with the constants in ggml-vulkan.cpp. V^T and multi-row never run on the sparse gather path.
• Fabric's TurboQuant V path (quantized V, scalar P x V loop) also reads Psh and now indexes the query-major layout. Without this, 88 TBQ/PQ FLASH_ATTN_EXT cases fail.
• FA_MASK_DEAD comes over from the NaN-safe commit, because the multi-row kernel's mask test uses it. The fork's Xe FA decode code that sat in the same hunks is dropped.
• The sub-op label keys on k_quant && v_quant, since fabric has no kv_needs_dequant.

Testing

On gfx1151, Vulkan0, from a clean build: FLASH_ATTN_EXT 5725/5725 (41 cases go through the multi-row kernel), MUL_MAT 2123/2123 and SOFT_MAX 212/212. The head-size 128/256 prefill FA cases also pass with each of GGML_VK_FA_MR=0, GGML_VK_FA_VT=0, GGML_VK_FA_WAVE32=0 and GGML_VK_FA_KV_CONTIG=0, and all prefill FA cases pass with GGML_VK_FA_KV_CONTIG=1. The new tests cover multi-row shapes (GQA 2/4/6/8/12, H = 1/2, ragged token blocks, no mask, sinks, ALiBi + softcap, q8_0 K/V, KV-cache views and a per-batch mask the kernel must decline), plus perf cases at the Qwen3.8 27B prefill shape.

Not tested: other Vulkan drivers and vendors. The contiguize pass, V^T and the multi-row kernel only default on for RADV RDNA3. The cm1 shader changes (P-fragment hoist, Psh layout) apply on every coopmat1 device, but were only run on RADV.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. Claude Code (Claude Opus 5.5) did the cherry-picks and conflict resolution, found and fixed the TurboQuant Psh regression, built and ran test-backend-ops and the benchmarks on the EVO-X2, and drafted this description. The commit carries Co-Authored-By trailers.

🤖 Generated with Claude Code

[yingying0906]

Two code findings in the Vulkan flash attention path.

[gianni-cor]

GGML_VK_FA_KV_CONTIG defaults on across coopmat1 devices, although the PR says it was benchmarked only on RADV. NVIDIA CI establishes correctness, not performance. would be better to have non-RADV performance evidence or a narrower default gate.

[gianni-cor]

The perf logger’s query pool can overflow. It allocates n_nodes + 100 slots, but every eligible FA node can emit two extra sub-operation timestamps. A graph with 50 such nodes needs 1 + 50 + 100 = 151 timestamps but receives 150 slots. The helper drops further sub-op marks, but the normal node timestamp then indexes the vectors and Vulkan query pool without a bounds check.

[amangupta-tether]

@gianni-cor thanks, both addressed. I also rebased onto temp-11018 now that #393 is squash-merged; the diff is unchanged.

K/V contiguize default (fc7846d8b). Agreed, there is no non-RADV performance evidence. The pass now defaults on only for RADV RDNA3, the same gate as the V^T scratch and the multi-row kernel. Other devices keep the strided cm1 loads, as on temp-11018. GGML_VK_FA_KV_CONTIG=1 turns it on for any device, and =0 turns it off.

Perf-logger query pool (cd0c14439). Confirmed: with 50+ FA nodes, the per-node timestamp after a dropped sub-op mark wrote past the pool and its vectors. The pool is now sized from the graph: one start timestamp, one per node, and two per FLASH_ATTN_EXT node (the only op that emits sub-op marks), plus the old 100 slots of headroom. Every timestamp write asserts the bound, and the sub-op helper asserts rather than silently dropping a mark.

Tested on gfx1151 (clean Vulkan build): FLASH_ATTN_EXT 5725/5725 with the new defaults, and the prefill FA cases pass with GGML_VK_FA_KV_CONTIG=1 and =0. A perf-logger run of Qwen3.8 27B pp2048 at d8192 still bills FA_KV_CONTIGUIZE and FA_V_TRANSPOSE on their own lines, and pp2048 at d8192 is unchanged at 334 t/s. That model has 16 FA nodes per graph, so the 50+ node case itself was not run; the bound covers it by construction.

[amangupta-tether]

unrelated gpu-webapple-apple failure. Merging