qwen4exp : gather-based sparse attention for QSA decode

[abdel-darwish-27]

Overview

This changes QSA sparse attention for qwen4exp during single-token decode so attention only runs over the tokens selected by the indexer.

The indexer used to pick the top 2048 KV cache entries, but those selections were turned into a mask over the full KV cache, so attention still ran across the entire context with unselected positions masked out. This meant the attention cost continued to grow with context length even though only ~2k tokens were actually being attended to. I noticed this because Qwen3.8 Flash was really slowing down at longer contexts compared to other models.

This patch instead gathers the selected keys and values into a compact buffer and runs regular dense attention over the gathered set. The attention mask is derived from the existing per-cell bias values, so this does not require any new model inputs.

Prompt processing and batched inference are unchanged and continue to use the existing masked path.

The gather path can be disabled at runtime with:

QWEN4EXP_QSA_GATHER=0

This was also used for A/B benchmarking, so the numbers below compare the same binary with the gather path enabled and disabled.

Additional information

Tested on dual RTX A6000 GPUs with an IQ4_XS model, q8_0 KV cache, and temperature 0:

• 130k context: 15.7 up to 23.6 tok/s (+50%)
• 62k context: 26.5 up to 31.6 tok/s (+19%)
• 31k context: 36.5 up to 38.5 tok/s (+6%)

You can see the improvement gets larger as the context grows. At 141k context, the old sparse attention kernel was taking roughly 15 ms per sparse layer per decoded token, across 12 sparse layers.

The gather graphs also no longer reference the full attention mask, which avoids uploading roughly 17 MB of mask data per token at 130k context.

The change is about 110 lines across src/models/qwen4exp.cpp, src/models/models.h, and a small guard in src/llama-graph.cpp. The guard skips filling the large attention mask when the graph does not allocate a buffer for it, matching the existing handling for other optional attention inputs.

I made sure to test that retrieval and short factual answers came out byte identical between the gather and masked paths at 31k, 62k, and 130k contexts every time.

Long open-ended generations can diverge after roughly 150+ tokens, but I observed the same behaviour between repeated runs of the unpatched build. Seems to be just normal GPU run-to-run nondeterminism rather than a difference introduced by the gather path.

All 53 CPU tests pass, and I also validated the CUDA build end-to-end with the real model.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - this code was written with AI assistance. I personally profiled the existing implementation, benchmarked the change, and validated the results on my own hardware.

[Rhonstin]

Tested on a much slower/lower-end box than the datacenter targets, and the gather path is a clear win — but it exposes the next bottleneck, which I think is worth spelling out.

Setup: Qwen3.8-Flash-Next UD-Q3_K_XL (84 GB, 3 shards), 8 GPUs (2× RTX 3090 + 6× CMP 90HX, all PCIe Gen1 x4), host is 2× Xeon E5-2620 with no AVX2/BMI2, 15 GB RAM. Single 160K slot, q8_0 KV, no speculative decoding. Cherry-picked onto current master + our local qwen4exp patches.

Method: 65,715-token prompt, 80-token decode, same binary toggled via QWEN4EXP_QSA_GATHER.

config	decode @60k (t/s)
gather OFF (QWEN4EXP_QSA_GATHER=0)	15.5
gather ON (default)	19.6
decode @2k (both)	39.8

So the gather path gives +26% at 60K and no regression at short context. Confirmed working, thanks.

But — with attention now O(top_k) and constant, decode at 60K is still 2× slower than at 2K (19.6 vs 39.8). The remaining per-step cost that scales with context is the selection itself: build_qsa_top_k runs ggml_top_k over the full expanded score tensor every decode step, i.e. O(N) per token regardless of how cheap the subsequent attention is. On this box that's the dominant term at 60K.

For anyone hitting the same wall: the indexer selection (not the attention) is where the next quadratic→linear win lives. A CUDA radix/partial top-k path would mirror what already landed for Vulkan (#28032) and ROCm (#27466) — CUDA still uses the generic top-k here. Happy to A/B a candidate if one shows up.

[ServeurpersoCom]

tg 1024, all on master:

                     4k      50k
master            103.4    73.2
+ #28244          103.4    78.9
+ #28213          99.7     79.9

Same gain at depth, 28213 loses 4% at 4k (full per-cell bias + width rounded to 2304).
#28244

EDIT:
Same approach here, implemented in #27977 on Aug 29 (commit d7ca8ea) and now split out as #28244. Performance came out the same when I measured the two, though that was on your base, which sits on an older master. The difference is that deriving the mask from the per-cell bias instead of gathering the existing kq_mask leaves self_kq_mask unreferenced, which is why this needs a guard in llm_graph_input_mem_hybrid::set_input; reusing the mask keeps the change inside qwen4exp.cpp.

[Rhonstin]

Follow-up after running a few more days with the gather path on this box (Qwen3.8-Flash-Next UD-Q3_K_XL, 8 GPUs — 2×3090 + 6×90HX, all Gen1 x4; host 15 GB RAM, Xeon E5-2620 no AVX2). A/B on the same binary, 60K context, q8_0 KV, single slot:

change	decode @60k
no gather (pre-#28213)	15.5 t/s
gather (this PR)	19.6 t/s
+ #28040 (kv-cells get_prev_tokens O(N)→O(log n))	21.6 t/s
@2k baseline	~40 t/s

Two observations that might help:

1. The remaining slowdown at depth is NOT the QSA selection. I A/B'd the two candidates directly: running top-k over only the width-sized trim instead of the full n_kv, and shrinking the gathered attention width from 2304 to 256 — neither moved the needle (@60k stayed ~19-20 t/s). CPU profiling shows the process sleeping in CUDA event handling (no host-side hotspot), so the residual cost is in some GPU kernel that still scales with the cache even when attention is gathered.

2. On this hardware the CUDA top_k fallback was not the bottleneck either — with CUDA 12.4 / CCCL 2.x we're on the argsort fallback (CUB_TOP_K_AVAILABLE needs CCCL ≥ 3.2), but replacing the top-k input with a tiny trim changed nothing. Worth knowing before anyone assumes the argsort path is the depth-limiting factor.

#28040 was the only thing that moved decode at depth here (+10%); thanks for splitting #27977 into reviewable pieces, the incremental commits are much easier to A/B than the original mega-PR.

[dpmm99]

Just for your awareness: build_attn_mha was changed in #27970 , so the build would break if you merged this in right now, but it won't be a merge conflict according to git. Qwen3.8-27B says the fix is to add n_topk before kq_scale, il where you called it in the if (gather) block.

[abdel-darwish-27]

Thanks for that, it's already fixed by the rebase on the PR.

[abdel-darwish-27]

I think the remaining depth scaling is probably coming from rebuilding the indexer's pooled/roped block keys every decode step over the full cached context. Since only the newest block changes, caching those block keys and updating them as you go seems like the fix, but that needs its own cache lifecycle and is a bit beyond this PR.

[TacoTakumi]

Data point from a PCIe multi-GPU desktop, plus one thing I found trying it with speculative decoding.

Rig: 2x RTX 3090, 1x 5060 Ti, 2x 4060 Ti (96 GB VRAM total), Qwen3.8-Flash-Next UD-Q4_K_XL, q8_0 KV, part of the experts on CPU via --fit. Same binary, toggled with QWEN4EXP_QSA_GATHER, greedy, 256 tokens generated, no speculation:

prompt	gather off	gather on
34,692 tok	20.14 t/s	21.64 t/s (+7%)
101,109 tok	13.90 t/s	16.60 t/s (+19%)

Prefill unchanged, no asserts, run-to-run spread under 1.5%. Matches the shape of your A6000 numbers.

With the unsloth MTP head (--spec-type draft-mtp, n-max 2) the patch does nothing, though: the gate is n_tokens == n_stream and verification passes draft+1 tokens per ubatch, so the gather path never runs. Measured 22.35 vs 22.86 t/s at 101k, inside noise. I think every number posted here so far is spec-off.

I tried opening the gate for small ubatches. Each token in the ubatch already has its own top-k list from build_qsa_top_k, so in build_attn_qsa I loop over the n_tps tokens, gather K/V by that token's list and its own bias row, run build_attn_mha per token over the 2304 gathered rows, and concat. Gate becomes n_tps <= 16 && n_kv >= 2widthn_tps, which is the same condition at n_tps 1. It is exact, each query sees the same set the masked path unmasks for it. At 101k with draft-mtp n-max 2 the verification step goes from 114.3 ms to 106.5 ms (-6.8%); the gain is smaller than spec-off because the step also carries two draft forwards. I also tried a batched version, one get_rows over all the lists and one attention over width*n_tps rows with a block-diagonal mask, and it was slower (109.0 ms). Happy to put the per-token diff up if you want it in this PR or as a follow-up.

One small thing for anyone cherry-picking this onto a tree without #27970: the build_attn_mha call in the gather block needs the extra n_topk argument dropped, the inverse of the master-side note above.

[abdel-darwish-27]

Gather now only kicks in once the cache is over 4x the top-k width (~9k cells), below that it's the unchanged masked path. My testing shows this fixes it:

decode tok/s      4k         31k    130k
masked (master)   52.5-52.9  36.5   15.7-17.0
gather (this PR)  53.0-53.3  38.5   23.8
sparse-fa TODO    51.3       37.5   18.2

[TacoTakumi]

Profiled where the rest of the depth cost goes on the gather path, on the 140 MB micro model from Lynxpda/micro-qwen4exp (same indexer shapes as the real model: 4 heads, key 128, top_k 2048, ratio 4, 3 QSA layers), one 3090, this PR's tree, nsys with CUDA graphs off, last 31 decode tokens of a tg32 run. Per token, 3 QSA layers:

n_kv	0	32768	91024	131072
kernel time, us	2155	2959	4433	5393
of which get_rows	62	640	1509	2103
add + rms_norm + rope + scale	514	698	1211	1530
flash_attn_ext (gathered, 2304)	33	43	43	43

Flash attention is flat. All of the growth is in build_qsa_top_k, per QSA layer at n_kv 131072:

1. Re-pooling the block keys every token: get_rows over all n_kv indexer keys (219 us), the four cont slices of the members (4 x 42 us, they come out as D2D memcpys), the three adds, scale, rms_norm and rope over all n_blocks (188 + 41 + 74 + 54 us). About 0.76 ms per layer. These all run at 400 to 1900 GB/s, so the cost is only that they cover every block every token. Caching the pooled+normed+roped keys for completed blocks and pooling only the tail block would remove it, which is what you described above.

2. Expanding the block scores back to cells: get_rows of 131072 rows of one f32 each runs at 2.6 GB/s (425 us per layer) because the CUDA kernel launches one block per row. That one is a ggml-cuda kernel shape issue, independent of the cache: either a small-row path in get_rows or taking top-k over blocks and expanding only the winners.

For the 48-layer model that is 12 QSA layers, so roughly 14 ms per token at 131k on a 3090 before any of the attention work, which matches the gap we still see between shallow and 100k+ decode with the gather on.

[sjordan0228]

Metal data point, Apple M2 Ultra (192 GB). Same build for every row: this PR's head (beed2f78a, build 10815) with #28349's one-liner behind an env switch. Qwen3.8-Flash-Next Q5_K_M (header patched to compress_ratios = [0,0,0,4]x12), F16 KV, -c 131072 -np 1 --kv-unified, no speculative decoding. Fixed prompts, cache_prompt: false, 512 tokens greedy, warm-up plus 3 runs, median. Gather confirmed active via GGML_SCHED_DEBUG. Noise floor 0.1-0.2%.

Decode tok/s:

path	2k	32k	128k
dense (zeroed ratios, reference)	35.04	31.21	21.66
masked (master)	32.10	24.98	15.02
gather (this PR)	32.02	23.62	14.84
sparse-FA (#28349)	31.32	24.05	14.38

Prefill is unchanged by the gather (522 / 451 / 332 vs masked 522 / 454 / 335).

On Metal the gather is neutral against masked at both depths, so no regression, but also none of the CUDA gain: the FA-vec kernel here already skips fully masked 32-cell chunks, so the masked scan was never the expensive part. All three QSA paths sit 20% (32k) to 31% (128k) below dense, and TOP_K in isolation is only 198 us at 131072 cells (about 2.4 ms per token over 12 layers), which leaves the block-key re-pooling as the cost, in line with the nsys breakdown above.