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.