vulkan: support sparse Flash Attention by 0cc4m · Pull Request #28105 · ggml-org/llama.cpp

0cc4m · 2026-08-31T17:38:40Z

Overview

Vulkan support for #27970

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES, Claude was used to assist.

On gfx1151 RADV the gather loses at every depth (11k -4.4, 115k -2.7 t/s vs masked, draft-mtp n-max 2): masked FA already skips fully-masked tiles, the per-block bias avoids the mask upload the gather exists to dodge, and gather mode forces the per-cell bias whose upload costs more than the gather saves. QWEN4EXP_QSA_GATHER=1 re-enables it for A/B. Revisit when Vulkan sparse FA (upstream ggml-org#28105) lands.

LynxPDA · 2026-09-11T10:42:49Z

@0cc4m Hi, I tried this PR on a Strix Halo machine with the Vulkan backend. Short contexts seem to work fine, but with long contexts global attention appears to break. The generation stays locally coherent, but the model intermittently loses parts of the previous history/messages.
A simple repro is to ask it to analyze a known-working code file of about 20-30k tokens. Without this PR, it gives a detailed analysis and confirms the code works. With this PR, it claims the code is broken, stitched together from drafts/hallucinations, and contains many truncated functions and unrelated fragments. It looks like some parts of the long context are not being attended to properly. Happy to provide more details if needed.

0cc4m · 2026-09-11T11:28:01Z

Please check with latest version. I'll try to reproduce it.

0cc4m · 2026-09-11T12:36:43Z

I tested it and didn't see an issue, it gives correct responses with long context input.

LynxPDA · 2026-09-11T13:49:39Z

@0cc4m it seems that draft had prefill-sparse. There was an issue with qwen4exp that I described earlier.
I can confirm that the issue is not present in the current PR.

0cc4m · 2026-09-11T16:19:04Z

qwen4exp is not supported by sparse FA yet, as far as I know.

jeffbolznv

I haven't gone through all the code yet, but wanted to get some early feedback out..

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.

jeffbolznv · 2026-09-14T13:14:19Z

-                    const uint v_row = j * Bc + row;
+                    uint32_t vcol;
+                    bool kv_active = fa_kv_index(j * Bc + row, vcol);
+                    const uint v_row = USE_SPARSE ? vcol : (j * Bc + row);


Can this just use vcol unconditionally? It doesn't follow the same pattern as other replacements.

Yes, that wasn't necessary. Fixed.

jeffbolznv · 2026-09-14T13:15:35Z

+    const int r = data_sparse[sparse_base + blockCoords[0]];
+    if (r < 0) { return f16vec4(0); }
+    const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
+    return f16vec4(data_kf16[o], data_kf16[o + 1], data_kf16[o + 2], data_kf16[o + 3]);


I think it should be possible to declare an f16vec4 binding and just do one load.

jeffbolznv · 2026-09-14T13:18:50Z

        const bool k_use_decode = (bs_k > 1u);
-        if (k_use_decode) {
+        if (USE_SPARSE) {
+            coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK);


Should FAGATHERK just be called FADECODEK now and remove the branch?

They use different addressing, so I don't think that's possible here.

I don't understand. I'm just talking about renaming the #define, and then combining line 392 with line 394. I think this should work?

They are distinguished by a spec constant flag in the same shader compile, I can only separate them if I compile them separately.

github-actions Bot added model Model specific testing Everything test related Vulkan Issues specific to the Vulkan backend ggml changes relating to the ggml tensor library for machine learning CUDA Related to the CUDA backend labels Aug 31, 2026

LynxPDA mentioned this pull request Sep 10, 2026

vulkan: sparse prefill flash attention for qwen4exp top-k masks Nathanw1014/llama.cpp#10

Draft

0cc4m added 3 commits September 11, 2026 11:52

vulkan: add sparse Flash Attention support for DSV4/GLM

f4d1e06

tune implementation

9d624c7

add tests

5cfc32c

0cc4m force-pushed the 0cc4m/vulkan-fa-sparse branch from bad0e3e to 5cfc32c Compare September 11, 2026 11:24

0cc4m marked this pull request as ready for review September 11, 2026 11:34

0cc4m requested review from a team and ggerganov as code owners September 11, 2026 11:34

jeffbolznv reviewed Sep 12, 2026

View reviewed changes

0cc4m added 2 commits September 14, 2026 08:34

avoid nondeterministic atomicAdd

4191ad1

add cm2 decode vector support

8d9b181

jeffbolznv reviewed Sep 14, 2026

View reviewed changes

0cc4m added 2 commits September 14, 2026 16:27

simplify logic and make variable names more consistent

de18aa9

add cm2 f16vec4 binding for decode vector

201031a

jeffbolznv approved these changes Sep 14, 2026

View reviewed changes

Uh oh!

Conversation

0cc4m commented Aug 31, 2026

Overview

Requirements

Uh oh!

LynxPDA commented Sep 11, 2026

Uh oh!

0cc4m commented Sep 11, 2026

Uh oh!

0cc4m commented Sep 11, 2026

Uh oh!

LynxPDA commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

0cc4m commented Sep 11, 2026

Uh oh!

jeffbolznv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

LynxPDA commented Sep 11, 2026 •

edited

Loading