Conversation
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.
…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.
The fix/qwen4exp recipes live in the project's local docs tree, not in the upstream-facing PR diff.
Author
|
On upstream ggml-org#28105 ("vulkan: support sparse Flash Attention") Linking for reference: the same sparse-FA compaction approach landed upstream (merged 2026-09-15, fc82583). This PR already contains a port of it — 20edd90 ("vulkan: sparse flash attention for qwen4exp top-k masks"), followed by two fork-specific fixes. Summary of the differences and why we are not cherry-picking further from upstream: History note: upstream PR ggml-org#28105 was actually the source of this port — our flash_attn_sparse_compact.comp is derived from the PR branch of that change, so the overlap is intentional, not coincidental. What we carry beyond the upstream final state:
What upstream has that we do not (and why we skip it for now):
|
The CUDA compact-mask kernel and the use_sparse plumbing in fattn-common/mma/tile/vec were ported together with the Vulkan sparse path in 20edd90, but this fork targets Strix Halo RADV/Vulkan only, nothing on the box or in CI builds or tests CUDA, and without a maintained consumer the code would silently rot. Backends without the patch read a sparse-tagged FA op as a dense masked one, which stays correct.
Owner
|
Impressive work! A lot of this can be used to add or replace some of the work in the most recent v7.6 release, im running some testing now |
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
2 participants
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.
Authorship & AI Disclosure
Summary
On
strix-halo-vulkan@dff600487, qwen4exp (Qwen3.8-Flash-Next, QSA / query-aware sparse attention) prefill throughput degrades linearly with context depth while decode stays flat.The QSA top-k indexer already selects ~2051 compressed cells per query and the attention mask is
-INFINITYeverywhere else — but flash attention still walks the whole uncompressed cache row (n_kv wide) for every query row: prefill is O(depth²) by construction.Profiling one prefill graph (ub 512, pp4096, d131072,
GGML_VK_PERF_LOGGER=1), dense:Worse, the sparse shared-tile path — used in decode — was silently wrong for prefill: one workgroup per (kv tile, head, batch) reads the mask at
iq3 == 0, but GQA prefill tiles share that workgroup across all query rows of the batch, so every row after the first attended row 0's selection. Output was garbage past the first token — not just slow.What this PR does
n_head_kv > 1— correctness fix first (20edd90d6).flash_attn_union.comp— scan the group's 64 × 2051 compressed-cell indices into a shared bitmap, prefix-sum + compact into a union index list in scratch; dynamic union count feeds the FA dispatch size (kv_c).flash_attn_gather_union.comp— gather K/V from the full cache into compact scratch ([kv_c, n_head_kv], no head duplication) and build the compact mask from the union list.expand_qsa_block_indices, per-row widthtoken_topk + compress_ratio − 1 = 2051).n_kv ≥ 2 · PAD(4·union_cand, 256): dense stays cheaper until the union actually compresses (measured crossover ≈ n_kv 43k at n_top_k = 2051, union/cand 0.1–0.3).GGML_VK_FA_UNION_FORCE=1bypasses for A/B.0b08831cf); opt-outGGML_VK_FA_TOPK_UNION_GQA=0.test-backend-ops: permanent union tests + 13369/13369 in both default and FORCE modes.Numbers
Real-request A/B (the authoritative comparison)
Real text corpus (mixed Python code + prose,
scripts/prof-real.sh), one server start at ctx 131328,-ub 512, MTP speculative decoding on, temperature 0, 2 reps per depth, temperature 0, real tokenizer. Full modelQwen3.8-Flash-Next-Q4_K_S(105 GiB, 12 attn layers, 24 Q / 2 KV heads, n_top_k 2051, compress_ratio 4, f16 cache), RADV STRIX_HALO.Base arm = the clean branch this PR was cut from (
strix-halo-vulkan@5d8c07b44), same harness, same corpus.A causal prefill averages over all depths up to the prompt end, so the "effective depth" of a prompt of N tokens is ~N/2 - that is the column to read against flat-depth bench rows below.
5d8c07b44tg at deep context was inconclusive in this harness run: at temperature 0 on the repetitive corpus the PR arm hit EOS after 9 generated tokens (short window, 43-45 t/s not comparable), while the base arm generated the full 128 (20.8-22.2 t/s). The flat-depth reference below (tg128 @ d131072, ~0%) remains the comparison for decode.
The real-request gains line up with the bench numbers once the causal averaging is accounted for: +14% at effective depth ~39k vs bench +16.7% at flat d32768; +23% at ~56k vs bench +28.3% at flat d65536. Real text is the harder workload (top-k selections vary more than on synthetic prompts), so these are the numbers to trust.
Flat-depth llama-bench reference
llama-bench -p 4096(synthetic prompts, every token processed at the SAME depth; fine for attention shapes, but it does not exercise realistic top-k overlap and MoE routing, so treat it as secondary): basedff600487(dense) vs this branch, union default (no env):dff600487(dense)Degradation ratio pp4096 d2048→d131072: 2.85x → 1.64x.
(Intermediate revision
0b08831cfmeasured 499.94 / 343.41 / 287.36 / 247.71 at d2048/32768/65536/131072; the scan/FA overlap0209047a3added the rest.) The gate keeps shallow depths dense, so there is no shallow-depth regression beyond noise.Robustness beyond the single-slot case
-np 4 --cont-batching, 8 concurrent mixed requests across 2 waves - all 8 answers correct, all 4 slots exercised, 0 server errors.-ngl 24(pp ~96 t/s),-ngl 8, and-ngl 0(CPU-only) - correct answers in every mode.Quality (union engaged): needle probes with the passphrase planted at 25%/50% (32.4k prompt) and 10%/60% (99k prompt), grammar-pinned answer, temperature 0 — all exact; the 99k probe passes through the union regime (engages from n_kv ≈ 43k, verified via
GGML_VK_FA_UNION_STATS=512).Testing
test-backend-ops -o FLASH_ATTN_EXTon Vulkan0: 13369/13369, in both build trees (build-release, secondarybuilds/pr-929aa8d5b), and in both arms — default (grouped union for GQA) andGGML_VK_FA_UNION_FORCE=1(gate bypassed, so every sparse-capable shape takes the union path).GGML_VK_FA_TOPK_UNION_GQA=0(dense fallback) is a third arm of the same sweep. Coverage is asserted as an exactN/Nagainst the tree's test count, never as "0 failures".kv=32768(R fills the bitmap exactly, soscan_wordsis a multiple of the emit chunk) and then_top_k=2051, hs=256, nh=24, nh_kv=2GQA set; these are the first cases re-run after any change to the union path.GGML_VK_FA_UNION_STATS).llama-benchprocess at a time.GGML_VK_QSA_PRIV_FORCE=1takes the private-output route on every site, so the route the full model picks by accident gets a deterministic run.test-backend-ops -o TOPK_QSAis 5/5 in both arms, and thellama-ab-runfingerprints are byte-identical in both.Updates after the first revision
188988f2d+380b00a2d+929aa8d5b— the QSA indexer pipeline (the pp tail at depth).Three separate costs, profiled on the full model:
TOP_K k=2051cost 678 ms per bench window in a 4-pass radix select. The passes are now 11 bits each (three passes, one full row of traffic less per token) and the emit scan processesEMIT_W=8elements per invocation with a subgroup ballot fallback: 678 -> 532 ms.(1,0,2,3)transpose of the block score and gathering from that copy. At 135k the copy is 1.09 GB and runs at ~10 GB/s,~20x below DRAM, because 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 thatcell_blkyields walks consecutive addresses.TOPK_QSA23.3 -> 14.1 ms per window, the standalone transpose 2.70 s -> 0.52 s, prefill wall 45.2 -> 39.5 s (micro, pf131072).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=1admits 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, and the knob is its A/B arm.929aa8d5bis a diagnostic that had been lying: the perf logger writes a node's fusion name after the guard has already run, so a declined fusion kept the fused name while its nodes ran one by one. That is how a profile read "12 x TOPK_QSA" at 131k when 7 sites were fused and 5had declined and were paying for the unfused transpose.
None of this changes arithmetic: the gather is value-identical, radix select is exact for any digit width, and the
llama-ab-runfingerprints (micro) are byte-identical in both the default and the FORCE arm.0209047a3— scan/FA overlap (the change behind the d32k/64k/131k numbers above). The union scan of group g+1 is issued immediately after FA(g) with no barrier between them, using per-group scratch slots; the scan no longer owns a serialized ~32 ms per layer. This is what took d131072 from ~248 to ~302 t/s. An earlier attempt to issue all scans up front behind a single barrier produced wrong results on nb=128 shapes and was abandoned; the dependency could not be characterized, so do not retry that shape.Parallel multi-workgroup scan — implemented, verified, reverted (negative result). A mark(64 workgroups into private bitmaps) → count(merge + chunk bases) → emit pipeline was written and made bit-identical to the single-workgroup scan (13369/13369 in both trees), but measured slower on the full model: 350/300/251 t/s at d32k/64k/131k against 390/355/305 for the single-workgroup version. Four dispatches plus two barriers per group cost more than the internal parallelism saves, and the scan is already hidden behind FA(g) — there is nothing left to parallelize away. Recipe and the debugging traps are recorded in
docs/opt-fa-union-parallel-scan-negative.md; the code is preserved on branchwip/parallel-union-scan(local, archived inbuilds/branches-backup-2026-09-15.bundle) rather than shipped.The one bug this exercise did surface is worth stating because it is silent: calling
subgroupAdd()insideif (lane == 0)makes only lane 0 participate, so the "total" collapses to lane 0's own value. Counts stayed correct (the same pattern with the collective outside the branch) while the emitted index list came out unsorted with plausible keys — a failure mode that looks like a data race and is not one.Profiles and findings
Remaining projected win: a fused per-row-index FA prefill (SGLang qsa shape), ~310–360 t/s.
GGML_VK_FA_UNION_FORCE=1is the second arm of every FA test sweep.Env knobs
GGML_VK_FA_TOPK_UNION_GQA=0GGML_VK_FA_UNION_FORCE=1GGML_VK_FA_UNION_STATS=NGGML_VK_FA_TOPK=0GGML_VK_FA_SPARSE_DISABLE=1GGML_VK_QSA_PRIV_FORCE=1Commits
20edd90d6vulkan: sparse flash attention for qwen4exp top-k masks (decode path + dense prefill gate)7e56d92e8vulkan: size the sparse-FA compaction per-subgroup tally to the workgroupc46ecbcd3docs: recipe — shared per-subgroup tally sizing bug in the sparse FA compaction949cc2cd4docs: recipe — the sparse-FA shared-tile tiling contractec2265962vulkan: fix grouped union prefill host addressing for non-square head/batch shapes161ab5976docs: recipe — push constants carry the destination shape, not the dispatch tile's0b08831cffeat: enable the grouped union prefill for GQA caches by default0209047a3vulkan: overlap the union scan of group g+1 with flash attention of group g188988f2dvulkan: widen the top-k radix pass and batch the emit scan380b00a2dvulkan: fold the QSA score transpose into the top-k gather929aa8d5bvulkan: clear the fusion label when a fusion is declined