vulkan : WMMA selected-key attention for QSA block selection

[Nathanw1014]

Stacked on #119 (and #115 under it); only the last commit is new here. This makes the selected-key attention from #119 faster on Vulkan. The change is limited to the selected-key dispatch and its pipelines.

Two pieces:

• A small prep pass cleans each selection row once (drops invalid and repeated cells, keeps order, stores the count) instead of every K/V head and split redoing it. The scalar kernel reads these rows, so its sort goes away and its shared memory at GQA 12 drops from 63 KB to 47 KB.
• A WMMA kernel: one wave per (query, KV head), with the GQA heads as the tile columns, and K/V loaded straight into the matrix registers. It needs 64-lane subgroups and 16-byte aligned K/V; anything else falls back to the scalar kernel, and GGML_VK_SEL_CM=0 forces the fallback.

On a Strix Halo (gfx1151, RADV) with Flash-Next, the attention op at 12k depth goes from about 68 ms to about 8 ms, and pp2048 goes from 568 to 711 t/s at 12k and from 538 to 679 t/s at 32k. Depth 0 doesn't use this path and is unchanged. Decode at 32k is slightly faster (21.8 to 22.4 t/s).

Perplexity moves within the error bars. The scalar fallback gives identical output to #119. The WMMA kernel sums in a different order, which gives a KLD of about 0.027 against #119 on Flash-Next. A near-exact variant I tried (hi/lo split) gave 0.026, so that looks like this model's floor for any change in summation order rather than precision loss. New test cases cover complete, partial and empty rows, GQA 4/8/12, and 32k/64k cells; scrambling the output order makes them fail.

Written with Claude, tested by me.

Assisted-by: Claude Opus 5.5

[gaetan-puleo]

Independent verification on a second Strix Halo (strix-halo-002, gfx1151, RADV)

Built test-backend-ops at head f1b0c44d1 in the kyuz0 RADV container.

Correctness

Full QSA suite, -b Vulkan0:

op	result
QSA_PREFILL	26/26
QSA_DECODE	92/92
QSA_EXPAND	18/18
QSA_PREFILL_MASKLESS	13/13 (NLOOP=1 and NLOOP=3)
QSA_DECODE_MASKLESS	33/33 (NLOOP=1 and NLOOP=3)

All pass. The maskless tests are the strong ones here: their err() override builds a host FP64 oracle over the listed cells and returns max(gpu, cpu) NMSE against the 5e-4 threshold. The WMMA kernel comes in at ~4.1e-8 to 4.6e-8 (e.g. q=9 keys=65536 → 4.35e-8), versus ~5e-14 to 1.8e-13 for the scalar kernel. So the f16 Q*f16-exp path is a couple orders worse than the scalar path but still four orders under the oracle tolerance — the KLD you measure is the model's summation-order floor, not a precision bug.

The WMMA path is really used

I temporarily printed use_cm in ggml_vk_flash_attn_sel (reverted afterwards): it is 1 for all 46 maskless cases (both op filters), and 0 for all of them with GGML_VK_SEL_CM=0. So the new flash_attn_sel_cm_f16 pipeline is created on gfx1151 and selected, and the tests above genuinely exercise it rather than silently falling back to the scalar kernel.

Review

The two pieces look right to me:

• Prep pass. The open-addressing table with atomicCompSwap / atomicMin correctly keeps the first occurrence regardless of thread order, the Hillis-Steele scan preserves the original order over the contiguous per-thread ranges, and padding with key_sh[first_sh] is always a surviving cell (the lowest valid position can never be a duplicate). The termination argument (at most n_sel <= TABLE distinct cells, so a probe either finds its key or an empty slot) holds. I also checked the "read a whole key block" edge: kb_begin + 15 can only reach into the row when list_stride == n_sel, and there both n_sel and chunk are multiples of 32, so the final chunk always ends at least 32 entries early — no overrun.
• WMMA kernel. I walked the fragment assumptions against the host comment (A: lane L%16 = key, element i = dim; B: lane L%16 = head, element i = dim; acc: lane L%16 = head, rows L/16 + 4i): Kᵀ fill, the Qᵀ fill, the column-shared max/sum via xor 16/xor 32, the Pᵀ rebuild from own/par with the grp parity, and the Vᵀ 8-tile permutation (dim = (t/8)*128 + 8*r + t%8) that the store undoes all line up. Heads >= GQA are zeroed in q_sh and discarded on store, and because cooperative-matmul columns are independent they cannot contaminate the live heads. The output store is race-free (the four grp lanes of a column write disjoint dim sets), the split-K partial layout matches flash_attn_split_k_reduce, and an empty split writes M = -FLT_MAX to avoid exp(-inf - -inf).

Points worth considering, none blocking:

• The cm kernel hard-codes RADV's gfx11 WMMA register layout, which the KHR cooperative-matrix spec does not guarantee. The gate checks vendor/arch/driver but no driver version, so a future RADV lowering change would silently produce wrong numbers rather than falling back. A version guard, or a one-time tiny self-check at pipeline creation, would make this fail safe.
• ggml_vk_flash_attn now takes the selected-key branch whenever src[5] != nullptr && mask == nullptr and GGML_ASSERT(done). That is safe under the scheduler because supports_op uses the same predicate, but it aborts rather than declining if the op ever reaches the backend without that check first.
• qsa_selected_key_supported probes a representative op per QSA layer at memory-construction time. It uses the same supports_op, so it should track the runtime op, but it is a shape approximation — if a layer's real op is declined for a reason the probe doesn't reproduce, the maskless graph would be built anyway. Given the mctx gate only forms the graph when the probe passes, this is the one place I'd keep an eye on.

I did not reproduce the end-to-end prefill t/s or the KLD; only the op-level tests and the oracle.

CI

windows (both) and ubuntu pass; ubuntu-arm64 passes. ubuntu-llvmpipe fails the same three tests as master b79812d59 (test-recurrent-state-rollback, test-recurrent-state-rollback-qwen35moe, test-save-load-state, all Subprocess aborted), and metal is still queued. Nothing here is from this PR.

LGTM. The WMMA kernel is exercised on gfx1151, matches the FP64 oracle by ~4e-8, and the scalar fallback is preserved.

[Nathanw1014]

Thanks for the oracle numbers and for confirming the WMMA pipeline is really selected. On your three points:

• Layout guarantee. The pipelines are created only on RADV + RDNA3 (driver_id == eMesaRadv), so other drivers already use the scalar kernel. A future RADV change to the WMMA lowering is the remaining risk. The maskless op tests with the FP64 oracle would catch it, and GGML_VK_SEL_CM=0 is the escape hatch. A one-time self-check at device init would make it fail safe without a test run. That would serve this kernel and the chunked GDN scan in vulkan : chunked gated delta net prefill #128, so I'd rather do it as one follow-up PR covering both than add it here.
• GGML_ASSERT(done). supports_op and the dispatch both go through ggml_vk_flash_attn_sel_pipeline, and ggml_vk_flash_attn_sel returns false only when that is null, so they cannot disagree. Aborting is deliberate: there is no correct fallback on Vulkan, because the dense kernel would attend to every cell (that is Vulkan: wrong results with flash attention + f16 KV cache on Qwen3.8-Flash-Next (perplexity 1.05 vs 1.55 reference), reproduced on b1004 and b1005 #130).
• Per-layer probe. Agreed, that is the place to watch. If the real op were ever declined for a reason the probe does not model, Vulkan refuses it and the op falls back to the CPU, which computes it correctly; it would be slow, not wrong. (On Vulkan: wrong results with flash attention + f16 KV cache on Qwen3.8-Flash-Next (perplexity 1.05 vs 1.55 reference), reproduced on b1004 and b1005 #130 I forced block selection on a build without the Vulkan selected-key kernel and got the correct PPL that way.)

[dzannotti]

Verified on Strix Halo (gfx1151, RADV), current master ad219f329 vs master + this PR. Qwen3.8-Flash-Next GSQ-RCO-Q2_0, -fa 1 -b 2048 -ub 512, ABCCBA.

test	master	PR	Δ
pp2048	596.2	596.0	±0
pp2048 @ d12000	458.9	544.3	+18.6%
pp2048 @ d32000	433.4	521.9	+20.4%
tg64 @ d32000	25.7	26.6	+3.6%

Correctness:

• QSA_PREFILL_MASKLESS 13/13 and QSA_DECODE_MASKLESS 33/33, with WMMA and with GGML_VK_SEL_CM=0.
• QSA_PREFILL, QSA_DECODE and QSA_EXPAND pass. FLASH_ATTN_EXT 5344/5344, same as master.
• KLD against master at 8k context: 0.0011 (PPL ratio 1.0002). Against the HIP reference, master's Vulkan is at 0.00244 and this PR at 0.00249, so no measurable loss.
• The scalar fallback gives master's output exactly (KLD 0).

Two small follow-ups: put the prep pipeline (~34 KB shared memory) under the same shared-memory check as the scalar kernel, and track the RADV-layout self-check you mentioned, since the WMMA fragment layout isn't version-guarded.

Comment generated via automated review of clanker.