Repository navigation
vulkan : selected-key flash attention for QSA block selection - #119
Conversation
|
Tested this on a Strix Halo (gfx1151, Vulkan), on current master with #115: the maskless QSA tests pass (6/6 and 25/25) and Flash-Next perplexity is in line with the masked path. Looks good to me. It needs a rebase after #111 too, since the Vulkan backend is split into several files now. I have a rebased copy here if it saves you the work. I'm also working on a speed-up on top of it (loading K and V straight into the matrix registers). I'll open it as a separate PR stacked on this one once it's tested. Assisted-by: Claude Opus 5.5 |
…ernels run Complete-block selection builds a maskless flash-attention op whose visibility lives only in the selected-cell list (src[5]). Only the selected-key kernels read that list, and they exist only in the HIP backend (RDNA3.5); the gate did not check the device. ggml-cuda aborts on the op when no QSA kernel takes it. The Vulkan backend accepted it and ran its dense kernel, which ignores src[5], so the QSA layers attended to the whole K/V view, unmasked, as soon as the context exceeded indexer_top_k + ratio - 1: - other sequences' cells: with several conversations in one unified cache, a conversation's output depended on another conversation's content (a 2048-token prompt, greedy, changed at the 2nd token next to related text, where it had been 6.6 nats ahead; unchanged next to unrelated text) - no sparsity: every cell instead of the top-k blocks test-backend-ops shows it directly: QSA_PREFILL_MASKLESS fails on Vulkan (0/6, ERR 0.48 to 0.98 against the FP64 reference). - The memory decides once per model whether its QSA layers all run on ROCm (LLAMA_QSA_SELECTED_KEY=0/1 overrides that), and qwen4exp_use_block_selection() requires it. Everywhere else the masked top-k path is used, which the dense kernels compute correctly (QSA_PREFILL and QSA_DECODE pass on Vulkan except their poison=1 cases, which fail on master as well). - Vulkan no longer claims flash attention with a selection list and no mask, so the op can never be computed wrongly there. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Complete-block selection builds a maskless FLASH_ATTN_EXT whose visibility is only in the selection list (src[5]). Vulkan had no kernel that reads the list: halo-box#115 makes it decline the op, and Vulkan then takes the masked per-cell top-k path, which is correct but slow at depth. flash_attn_sel.comp computes the op: one workgroup per query x K/V head, one 64-lane subgroup per query head of that K/V head (GQA 4/8/12, head 256). It sorts the query's list in shared memory and skips invalid and duplicate entries, as the CPU path does, then gathers K and V in blocks of 32 cells. Two lanes score one key, so the online softmax needs two subgroup reductions per block. With fewer than 64 workgroups the rows are split, and flash_attn_split_k_reduce combines the parts. supports_op takes a maskless selected-key op exactly when this kernel does. llama_memory_hybrid_idx asks each QSA layer's device with a probe op instead of checking for ROCm by name. HIP stays accepted; CUDA is not asked, because its supports_op ignores the list. Radeon 890M (gfx1150), Qwen3.8-Flash-Next UD-IQ4_XS: - test-backend-ops: QSA_PREFILL_MASKLESS 6/6 and QSA_DECODE_MASKLESS 25/25 (master: 0/6 and 0/25), NMSE <= 6.9e-13 against the FP64 reference - pp2048 at depth 12000 / 32000: 140.9 / 134.1 t/s (masked path, halo-box#115: 139.0 / 88.7) Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2029c38 to
3bf5356
Compare
|
Thanks for testing on gfx1151. Good to see 6/6 and 25/25 there too. Rebased onto 307c50d: the pipeline array went to On 307c50d (gfx1150), pp2048 at 32k / 64k is 149.8 / 144.1 t/s (#115: 93.3 / 60.8; master's wrong path: 113.5 / 80.6). One thing that may matter for the speed-up: on this path two compute buffers (~810 MiB and ~175 MiB) are reallocated every prompt ubatch, growing by 1-2 MiB each time. The reserve at load is probably sized for the masked graph. Results don't change, but you'll see the churn in a memory logger. Assisted-by: Claude Opus 5.5 |
|
Reproduced the buffer churn on gfx1151, on both Vulkan and HIP (ROCm 7.14). The load-time reserve only builds the maskless graph, which has no KQ mask. While the context is short (n_kv up to 2048 on Flash-Next), ubatches take the dense shortcut instead, so the buffers grow by the mask size until they've seen it. Reserving that dense graph at load as well fixes it: no runtime reallocations in llama-server on either backend, about 8 MiB extra at ub2048, and identical perplexity. On Vulkan, perplexity also runs about 2% faster. HIP perplexity still grows the device buffer once (about 38 MiB) with or without the fix; I haven't tracked that one down yet. It's llama-side only and applies cleanly on this PR: hb/qsa-reserve-dense. Happy for you to pull it in here, or I can open it separately. Assisted-by: Claude Opus 5.5 |
|
Thanks, that's a clear diagnosis. I tried
I'd keep your fix as its own PR, since it's llama-side and covers HIP too. Stacking it on this one is fine. Assisted-by: Claude Opus 5.5 |
On master, Flash-Next takes the selected-key path on Vulkan (halo-box#119), where the QSA block selection is a GGML_OP_TOP_K over the KV cells. ggml_vk_topk sizes its double buffer in the shared scratch (prealloc_x) to the exact need, so the scratch was still reallocated as the context grew: 172, 173, 174, 176 MiB during one 4096-token prompt next to 174k resident cells. Size it with ggml_vk_scratch_size() like the other context-sized scratches; below 64 MiB (any ordinary top-k) nothing changes. Its passes bind dbl_buf_size subranges, not the whole buffer. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On master, Flash-Next takes the selected-key path on Vulkan (halo-box#119), where the QSA block selection is a GGML_OP_TOP_K over the KV cells. ggml_vk_topk sizes its double buffer in the shared scratch (prealloc_x) to the exact need, so the scratch was still reallocated as the context grew: 172, 173, 174, 176 MiB during one 4096-token prompt next to 174k resident cells. Size it with ggml_vk_scratch_size() like the other context-sized scratches; below 64 MiB (any ordinary top-k) nothing changes. Its passes bind dbl_buf_size subranges, not the whole buffer. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 participants
Overview
Depends on #115. This branch contains it; the change here is the last commit.
#115 stops Vulkan from running the maskless selected-key op with its dense kernel, which ignored the selection list
and attended to every cell. Vulkan then falls back to the masked per-cell top-k path. That path is correct but slow
at depth: it costs 18% of prompt speed at 32k cells and 25% at 64k (numbers in #115).
This PR adds the missing kernel. Vulkan can then take the maskless block-selection path, the same one the HIP
selected-key kernels take on RDNA3.5, and attend only to the cells in the list.
flash_attn_sel.comp:row is read once per query. There are pipelines for GQA 4, 8 and 12, and head size 256 (shared memory limits
the GQA).
-1) and duplicate entries. This iswhat the CPU path does and what the FP64 reference in test-backend-ops expects.
dot product (half the dimensions each), so the online softmax needs two subgroup reductions per block instead
of one per key. Accumulation is in f32.
flash_attn_split_k_reducecombines them. An empty chunk writesM = -FLT_MAXinstead of-inf, so the reducenever computes
exp(-inf - -inf).supports_opaccepts a maskless selected-key op exactly when this kernel takes it, and still declines everythingelse. So the dense kernel can never run the op.
llama_memory_hybrid_idx) no longer checks for the ROCm backend by name. It asks each QSA layer'sdevice with a probe op built from the layer's shapes, as llama.cpp already does for weight buffer types. HIP stays
accepted as before. CUDA is not asked, because its
supports_opdoes not look at the selection list, and on adevice without the QSA kernels ggml-cuda aborts on the op (qwen4exp: GGML_ABORT ("maskless selected-key attention was not taken by a QSA kernel") on KV-cache exhaustion under --kv-unified with concurrent slots #85).
Measurements
test-backend-ops -b Vulkan0:QSA_PREFILL_MASKLESSQSA_DECODE_MASKLESSQSA_PREFILL,QSA_DECODE(masked)poison=1cases fail on master too)TOPK_QSA,FLASH_ATTN_EXTThe maskless cases compare against a host FP64 reference over the listed cells. The largest NMSE this kernel gives
is 6.9e-13 (12643af); the CPU backend gives 1e-7 to 5e-5 on the same cases. Queries with at most 31 rows go through the split
path (1 query: 3e-14).
llama-bench -m MODEL -p 2048 -d 0,12000,32000,64000 -b 2048 -ub 512 -n 0 -r 4 -ngl 99 -fa on -ctk f16 -ctv f16 -lm none -lzm on-direct -o jsonlMaster: (307c50d; wrong past 2051 cells, see #115)
#115: (72f0f12, the masked per-cell path)
After: (3bf5356)
With the kernel, prompt speed hardly depends on depth: 158 t/s at 12k and 144 t/s at 64k cells, because every query
attends to 2051 cells however long the cache is. At 64k this is 2.4x #115 and 1.8x master's dense (and wrong) path.
Decode at depth,
llama-bench -p 0 -n 128 -d 0,12000,32000,64000 -r 2, run in the order #115, #117, this PR, thisPR, #117, #115 (4 samples each):
Decode at depth is limited by the indexer, which scores the whole cache for each token. #117 (the incremental QSA
cache) is what speeds that up; the two combine.
Device memory near a full context:
llama-server -c 262144 -np 4 -kvu, MTP n=1,GGML_VK_MEMORY_LOGGER=1, 174kcells resident, slot 0 reads a 4096-token prompt:
The maskless path needs about 1.5 GiB less device memory than the per-cell path. Two compute buffers still grow with
the cache and are reallocated every prompt ubatch. The buffers reserved at load are most likely sized for the masked
graph, not for this one. That is left for a follow-up; it does not change results.
Correctness (307c50d):
-c 65536 -np 4 -kvu, no MTP: a 2048-token prompt generates 48 greedy tokens alone, next to three16k-token conversations X (related text) and next to three Y (unrelated text). All three give the same 48 tokens.
Master, on the same path, depends on the other conversations' content (qwen4exp : use maskless block selection only where the selected-key kernels run #115 has the details).
llama-perplexity -c 4096on text the model cannot have memorised (the setup is in qwen4exp : use maskless block selection only where the selected-key kernels run #115): 4.6348 +/-0.1076; qwen4exp : use maskless block selection only where the selected-key kernels run #115: 4.6444 +/- 0.1079; master: 1.1422 +/- 0.0071, because its dense kernel sees later tokens. At
-c 2048, below the block-selection threshold, all three give 5.2977.-ub 512and-ub 256gives the same 32 tokens. The max top-1 log-probdifference is 0.0011 (qwen4exp : use maskless block selection only where the selected-key kernels run #115's masked path: 0.0010).
take the same path below 2052 cells. Repeated to 12.7k-53.4k tokens, block selection and qwen4exp : use maskless block selection only where the selected-key kernels run #115's per-cell top-k
select nearly the same cells: the top token after prefill and all 64 greedy tokens are the same on all four
prompts, and the next-token distributions after prefill differ by KL 0.002 (prose, 12.7k), 0.024 (code, 22.8k),
0.015 (structured, 27.1k) and 0.051 (numeric, 53.4k). On 12643af the numeric prompt, which ends where the model
is nearly split between two tokens, flipped its top token on batch shape alone; at that cache size the kernel
matched the FP64 reference to NMSE 7e-13 (temporary
test-backend-opscases with 32768, 53504 and 65536 cells,not part of this PR).
llama-perplexity, 16384 context, 2 chunks of Python standard-library source: 1.0677 +/- 0.0047; qwen4exp : use maskless block selection only where the selected-key kernels run #115: 1.0664+/- 0.0046.
the masked path (comment below).
Requirements
fork's qwen4exp QSA attention. Our measurements come from Strix Point (gfx1150), not Strix Halo.
ran every measurement. Every commit carries
Assisted-by: Claude Opus 5.5.subgroup sizes other than 64: the pipelines are only created with subgroup size 64.