vulkan: sparse prefill flash attention for qwen4exp top-k masks - #10
Draft
LynxPDA wants to merge 15 commits into
Draft
vulkan: sparse prefill flash attention for qwen4exp top-k masks#10LynxPDA wants to merge 15 commits into
LynxPDA wants to merge 15 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%)
The grouped union prices a GQA batch per group, so the estimate slot must be keyed by the group actually measured, min(64, n_batch): pricing a speculative 2-4 row batch as a 64-row group read past the end of the top-k tensor and poisoned the prefill slot, which made prefill decline on every depth for as long as decoding ran. The stat buffer now has one slot per group of the largest supported batch, and the estimator read is clamped.
ggml-alloc assigns a buffer only to tensors some node reads. A hybrid graph whose recurrent side is unused (an empty recurrent set, e.g. a qwen4exp MTP draft context with an all-false recurrent filter) leaves s_copy unallocated, and set_input then aborted on the null buffer. Skip it like any dead input.
The NextN/MTP block is a full-attention QSA layer: it ships its own trained indexer tensors and its norms sit far from their zero initialization, so the block's attention used the indexer during training. The draft was running that layer dense, which made server pp512 degrade with depth (a dense full-context FA over the whole ubatch, growing linearly): 6936/4494/2235/790 t/s at d0/8k/32k/131k on the micro model against 10631/6445/3943/1983 without the draft. The sidecar's compress_ratios[n_layer] == 0 is a padding artifact of llama-model-saver (the trunk's array padded to n_layer_all), not a statement about the block, so the loader restores the ratio from the trunk's QSA layers when the nextn layer really carries indexer tensors. The MTP context now gets llama_memory_hybrid_idx with the trunk's filters inverted: attention + indexer over il >= n_layer(), and a recurrent filter that is always false. An empty recurrent set allocates nothing; the old comment claiming its buffer allocation fails was wrong. The draft's sparse path only pays on prefill-sized batches, so graph_mtp enables QSA for n_tokens >= 16 and decode/verify batches run dense while still writing the raw indexer keys (pooled keys are recomputed above the pool watermark from the stored keys). The hybrid graph inputs also skip an s_copy no node reads. Micro model, llama-server, pp512/tg128, one session, dense draft vs QSA draft: pp 6258/4118/2120/734 vs 6157/4579/2350/1019, tg 124/114/94/53 vs 124/112/87/49. Recipe: docs/qwen4exp-mtp-sparse-draft.md
The radix selection used 8-bit digits, so a 32-bit key took four passes over the row. 11+11+10 bits does it in three, which is one full row of traffic less per token. Any digit width is exact for radix select, so the result is unchanged. The emit scan read the row once per class and once more per chunk. Both classes now share one pass: values strictly above the threshold fill [0, n_above) and the ties at the threshold fill [n_above, k). Each class keeps ascending element order, so the output equals the two-pass (above, then ties) result bit for bit while the row is read once. The two classes cannot be merged into a single ">= threshold" pass: the ties must land after all strictly larger values or the sparse attention summation order changes. Slots still come from a deterministic exclusive scan, so the output never depends on scheduling. The emit scan batches EMIT_W elements per invocation and prefix-sums them with subgroup shuffles; a per-chunk ballot path (EMIT_W=1) is kept for devices without subgroup shuffle. Spec constant 2 selects between them. Measured on the qwen4exp prefill shape: plain TOP_K 678 -> 532 ms per window.
The QSA indexer built its cell-major score by materializing the (1,0,2,3) transpose of the block score and then gathering cells from that copy. At 135k context the copy is 1.09 GB and runs at ~10 GB/s, ~20x below DRAM: each workgroup reads a strided column of the block-major source. The fused kernel wants values, not a layout, so it now reads the score in its native block-major layout, where the run of consecutive blocks that cell_blk yields walks consecutive addresses. The transpose is folded away and never executed; the pattern anchor moves from the gather to the transpose in front of it. A masked cell sums to exactly -inf whatever the block index is, so the kernel skips the block lookup and the score read for it. That is not just a shortcut: the block index of a masked cell is unconstrained, so the shortcut and the general path agree by construction rather than by luck. Every node of the pattern still exists for the unfused fallback (small k) and for the other backends. The gather is value-identical, so the selection and the fingerprints are unchanged. Measured on the qwen4exp prefill shape: fused TOPK_QSA 23.3 -> 14.1 ms per window, the standalone transpose 2.70 s -> 0.52 s, prefill wall 45.2 -> 39.5 s. The fused kernel reads the block score while it writes the cell indices, so it can only write them into the destination when nothing it reads shares storage with it. ggml-alloc does hand the output the memory of the block score, which made the fusion guard decline the whole site: at 131k, 5 of the 12 QSA layers fell back to the unfused transpose + top-k, ~125 ms per site per window. The guard is right, so the output moves instead: when an input overlaps the destination the indices go to the tail of the shared scratch and a buffer copy fills the output after a barrier. The test is on the kernel's read set, which is the whole hazard, rather than on the pattern's elided intermediates - so the guard can be skipped for this fusion without hiding a read-after-write. GGML_VK_QSA_PRIV_FORCE=1 admits the private route without an overlap. Which route a site takes is a property of the allocator's layout, so no test case reaches the private route on its own; the knob is its A/B arm.
The perf logger names a node's interval after the fusion it was chosen for, and the name is written after the fusion guard has already had its say. A declined fusion therefore kept the fused name while its nodes ran one by one, so the logger billed the fallback under the fused name. That is not cosmetic: it is how a profile showed "12 x TOPK_QSA" at 131k when 7 sites were fused and 5 had declined and were paying for an unfused transpose + top-k. The label now follows the decision.
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.