vulkan: sparse prefill flash attention for qwen4exp top-k masks - #10
Draft
LynxPDA wants to merge 8 commits into
Draft
vulkan: sparse prefill flash attention for qwen4exp top-k masks#10LynxPDA wants to merge 8 commits into
LynxPDA wants to merge 8 commits into
Conversation
LynxPDA
marked this pull request as draft
September 10, 2026 18:04
Port PR ggml-org#28105's sparse flash-attention compaction and wire it to the qwen4exp QSA prefill mask. The mask of a QSA layer is exactly the top-k selection intersected with causality, so only n_kv_max (= top-k width) cells per row are finite; the backend now compacts those positions per mask row and flash attention reads K/V/mask through the per-row index list instead of scanning the whole cache. - ggml_flash_attn_ext_set_sparse stores the per-row finite bound in op_params[5] (op_params[4] stays the fork's n_kv_raw); CUDA fattn passes the hint through for reference. - flash_attn_sparse_compact.comp builds the per-row index list with a deterministic subgroup-ballot scan (ascending position order, -1 padded). Upstream's atomic slot assignment is a race: the list order is the softmax accumulation order, so the run-to-run bits differ; the scan makes the sparse path bit-stable and identical between the cache on/off arms, which the A/B harness requires. - vulkan FA pipelines gain USE_SPARSE (bit 16) and the fork's DYNAMIC_KV moves to bit 32; cm2's sparse-only tensor-layout updates and gather offsets stay behind USE_SPARSE so the dense specialization keeps its codegen (an unguarded runtime KV select halved dense throughput on gfx1151). - The sparse gate follows the tiling contract of the shader: one index list and one mask row are resolved per DISPATCH TILE, so every row of a tile has to be the same query. That holds when the rows are the gqa heads of one token (gqa_ratio > 1, i.e. decode); large-N shapes run with gqa_ratio == 1 and correctly decline to dense. It also declines when the cache is under max(4096, min_ratio * n_kv_max) cells. FA_SPARSE_DISABLE reverts to dense for A/B. - Extend flash_attn_union/gather_union with a KV-head dimension and a batch offset, plus a grouped prefill driver (64-row groups, opt-in via GGML_VK_FA_TOPK_UNION_GQA): one compact set per group with the scratch reused per group. Inert by default; the per-row sparse path measured ahead of any shared-set compaction. That pp512 measurement was the broken configuration: it took the shared-tile path, which is fast and wrong. With the tiling fixed (see the following commit), the sparse path serves decode (gqa_ratio > 1) and prefill declines to dense. Micro model pp512 and A/B figures above therefore do not describe this commit as merged; re-measure before quoting them.
…roup The compaction shader tallied each chunk's per-subgroup finite counts in a shared array of 8, but the pipeline runs 1024 threads = 16 subgroups on wave-64 devices. Waves 8..15 wrote past the array: their counts were lost, the slot assignment shifted, and each mask row kept only ~930 of its ~2051 finite positions - a pseudo-random subset of the selection. Attention then read the wrong half of the selected cells: text stayed locally coherent but the model lost global structure (long-range analysis hallucinated non-existent issues). On the micro model at pf12288/c16384 the compacted list now matches the mask row exactly (2051/2051, ascending, in-bounds) and greedy decode tokens are identical to the dense-mask arm; residual logits drift is ~1e-5..1e-4 relative from online-softmax reblocking, same class as any summation reordering.
Write down how the "one index list and one mask row per dispatch tile" bug was localized, since the same shape of mistake is easy to re-introduce: a tile resolves one row for all its rows, so the gate has to encode the invariant that makes the rows interchangeable (gqa_ratio > 1), not a proxy for it (N >= 64). Includes the diagnostic method that found it -- sweep how much the query rows share and watch the error collapse monotonically -- and the two measurement traps hit on the way: a test that never set the sparse hint, so it measured dense while claiming to test sparse, and a generated shader header whose DEPENDS did not list the .comp sources, so shader edits were not compiled at all.
…/batch shapes The grouped union prefill path (per-group union of the query rows' top-k selections, then dense FA over the compact set) was correct only when the query-head count happened to equal the rows in a group. Three host-side shape mismatches hid behind that coincidence, all of them passing the group's dimensions where the shaders expect the destination tensor's: - the FA push constant ne1 carried the group's row count, but the shader uses ne1 as the head-to-head stride of dst ([HSV, n_head_q, n_batch, ns]): o_offset + iq2*HSV + row*ne1*HSV. It must be q->ne[2]. With nh != nb every head but the first was written to another head's rows, and at nb=128 the write ran past the tensor (DEVICE_LOST). - the group's slice of dst advanced by dst->nb[1], the head stride, instead of nb[2], the batch-row stride. - the split-K reduce was dispatched with x enumerating rows and with ne1/ne2 swapped against its own convention (x enumerates heads, z the rows of the split buffer), so every shape small enough to engage split_k was wrong. Found by bisecting the shape: the path failed for nh < nb and passed for nh == nb, which pointed at the host rather than at the shader. Coverage: the union gate is a measurement of the actual overlap, so the call that produces the estimate must itself decline, and test-backend-ops computes a case once - leaving the path with no deterministic coverage. GGML_VK_FA_UNION_FORCE=1 admits it without the estimate, for tests and for A/B runs; it can only cost a slow step, never correctness. The qwen4exp prefill cases with realistic adjacent-token overlap document the two variables they need. test-backend-ops on Vulkan: 13369/13369 with the union forced, with the gate alone (dense fallback) and by default.
…spatch tile's Record the grouped-union host-addressing bug as a recipe: the three shape mismatches (ne1 as the head stride, the group's dst slice stepping by nb[1], the split-K reduce's inverted convention), the pass/fail symmetry that located them, and the two traps that delayed it - a PASS verdict from a dense fallback when the gate is a measurement, and a loose -p regex claiming a verdict for a case that never ran.
It was opt-in (GGML_VK_FA_TOPK_UNION_GQA=1) pending full-model quality verification; the probes passed, so it becomes the default with =0 as the opt-out. The overlap gate is unchanged and still declines wherever compaction would not pay.
LynxPDA
force-pushed
the
pr/sparse-fa-pp-fix
branch
from
September 12, 2026 06:13
1be0904 to
0b08831
Compare
github-actions
Bot
added
the
documentation
…roup g The union scan (one workgroup, ~4 ms per group at depth) ran serialized behind the previous group's FA behind full barriers. Per-group slots for the union index list and the kv-count word make scan(g+1) data-independent of FA(g), so it is issued right after it with no barrier and overlaps it on the GPU. Allocation: the prealloc_y sizing gains one union-list slot per group (gul_all = gul_sz * n_groups); the fa_union_stat slot stride is 16 bytes. The estimator still reads slot 0 (group 0's last count). test-backend-ops FLASH_ATTN_EXT: 13369/13369. pp4096 full model (Q4_K_S, RADV STRIX_HALO): d2048 494.5 (was 500, noise band) d32768 392.2 (was 343, +14%) d65536 352.1 (was 287, +23%) d131072 301.7 (was 248, +22%)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On
strix-halo-vulkan@dff600487, qwen4exp prefill throughput degrades linearly with context depth while decode stays flat. The QSA top-k indexer already selects ~2051 blocks per query and the attention mask is-INFINITYeverywhere else — but flash attention still computes over the whole cache, so its cost grows linearly withn_kv.Profiling one prefill graph (ub512, pp4096, d131072,
GGML_VK_PERF_LOGGER=1):The fix
Port of upstream PR ggml-org#28105's sparse flash-attention compaction, wired to the qwen4exp QSA prefill mask:
ggml_flash_attn_ext_set_sparsestores a per-row finite bound inop_params[5];flash_attn_sparse_compact.compshader builds a per-row index list of the finite positions with a deterministic subgroup-ballot scan (upstream's atomic slot assignment is a race, and the list order is the softmax accumulation order — the scan makes the sparse path bit-stable);USE_SPARSEspecialization and read K/V/mask through the per-row index list, so each row attends only ~n_kv_maxcells instead ofn_kv: prefill FA drops from O(n_kv) to O(top-k width).The dense (non-sparse) path is untouched:
USE_SPARSEis off without a sparse mask. CUDA fattn receives the hint but ignores it by default.Benchmarks
Full model
Qwen3.8-Flash-Next-UD-Q4_K_XL(176.9 B, Q4_K_M, 103.68 GiB), RADV STRIX_HALO,-ngl 999 -fa 1 --load-mode mmap -ub 512 -r 1 -p 4096:dff600487tg128 is unchanged within run-to-run noise; the decode path is not touched by this commit (the deltas above are the same-order CPU-side variation visible across back-to-back runs).
Correctness
-fa 1 -c 2048 -b 2048 -ub 2048), full model:dff600487: PPL = 4.0328 +/- 0.02283pf512/pd64=72b4c0209c...— identical to baserm16/4096=6ae0c8c9b8...— identical to basetest-backend-ops: 33153/33153 pass.AI usage disclosure: YES
An AI coding assistant was used to draft this description and to run the benchmarks, profiling, and correctness checks above. The design decisions, the debugging that produced them, and the end-to-end validation runs were directed and verified by me, and I am responsible for every line submitted.