qwen4exp : gather-based sparse attention for QSA decode - #28130
Closed
abdel-darwish-27 wants to merge 2 commits into
Closed
qwen4exp : gather-based sparse attention for QSA decode#28130abdel-darwish-27 wants to merge 2 commits into
abdel-darwish-27 wants to merge 2 commits into
Conversation
At decode time build_attn_qsa built a full-n_kv mask (-INF everywhere with the top-k positions unmasked) and ran attention over the entire KV cache, so the indexer's top-k selection saved no attention compute. A kernel profile at 141K context shows flash_attn_ext_f16 at ~15 ms per QSA layer per token (~180 ms per token over 12 QSA layers); decode collapses from 53 tok/s shallow to ~7 tok/s at 141K on 2x RTX A6000. Add a decode-only gather path, taken for single-token-per-stream ubatches once the cache is at least twice the top-k width: - select winning *blocks* directly on the per-block indexer scores (the block bias already carries visibility), avoiding the O(n_kv) expansion of block scores to token scores and sorting n_blocks entries instead of n_kv - map winning blocks to cell indices via blk_cells and gather the selected cells' K/V (whole-cell rows; a cell's heads are contiguous in the cache) plus their kq_mask values, then attend densely over r*K_blk cells (2048 for Qwen3.8-Flash-Next), a multiple of 256 so flash attention padding holds - skip the O(n_kv) host-side cell_blk fill when the graph never references it QWEN4EXP_QSA_GATHER=0 restores the masked path (same binary A/B lever). Correctness: greedy outputs byte-identical to the masked path at 75K and 141K depth; mid-context needle retrieval passes in both modes at all tested depths. Decode throughput, UD-IQ4_XS, q8_0 KV, single stream (repeats within 0.1 t/s): depth masked gather 34K 17.4 19.4 (+11%) 68K 12.0 13.9 (+16%) 141K 7.1 8.8 (+23%) The remaining depth scaling in both modes is the indexer recomputing pooled block keys from the raw cache every layer per token; caching those incrementally is a follow-up.
The gather path only reads one row of the attention kq_mask (to carry each selected cell's visibility into the gathered attention), but referencing it kept the whole FA-padded tensor alive: an O(n_kv x GGML_KQ_MASK_PAD) host fill plus an n_kv x 64 x 2-byte upload every decode step (~18 MB/token at 141K ctx, measured as the largest H2D stream during decode, with the staging copy attributed to the driver at ~14% of decode CPU). Add a compact F32 [n_kv, n_tps, n_stream] visibility row to the QSA input set, filled in set_input_qsa alongside the existing per-token pass and gathered in place of the kq_mask row. The attention kq_mask then goes unreferenced in gather graphs and is neither filled nor uploaded; llm_graph_input_mem_hybrid now skips it when unallocated, matching llm_graph_input_attn_kv. Upload drops 18 MB -> 1 MB per token. Outputs remain byte-identical to the masked path; mid-context needle retrieval passes at 68K and 141K.
|
Hi @abdel-darwish-27, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
github-actions
Bot
removed
the
draft
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.
Problem
At decode time the qwen4exp QSA path builds a full-
n_kvmask and attends over the entire KV cache, so the lightning indexer's top-k saves memory only, not compute: nsys at 141K context showsflash_attn_ext_f16at ~15 ms per QSA layer per token, and decode throughput collapses with depth.Change
Decode-shaped ubatches (
n_tokens == n_stream) take a gather path instead: block-level top-k on the per-block indexer scores, winners mapped to cell indices viablk_cells, then K/V and the mask values for the selectedtop_kcells are gathered withggml_get_rowsand attended densely — the attention works over ~2K cells instead of the whole cache. Prefill and mixed batches keep the existing masked path.QWEN4EXP_QSA_GATHER=0opts out at runtime.The second commit feeds the gather path from a compact per-cell mask row instead of the FA-padded full-
n_kvkq_mask, removing ~18 MB of host->device mask traffic per token at 141K, and skips the O(n_kv) host-sidecell_blkfill when the graph does not reference it.Correctness
Retrieval and factual prompts are byte-identical between masked and gather at all tested depths (needle retrieval buried at 15% of a 129.6K-token document: correct and md5-identical in both modes; a 12-primes answer: md5-identical in both modes at 30.8K/61.8K/129.6K). Long open-ended generations (200+ tokens) can diverge after ~150 tokens — but the same divergence occurs between two runs of unpatched master with identical settings (measured), i.e. run-to-run fp nondeterminism, not a property of this change.
Performance
Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX A6000 (sm_86), q8_0 KV, single stream, temp 0, master
e4b9af00. Two reps per cell, contention-free samples only (foreign GPU activity sampled at 1 Hz during every measurement):Combined with #28128 and #28129, 129.6K-deep decode reaches 29.8 tok/s (+75% vs master).
Known limitation / follow-up
Both paths still recompute pooled block keys from the raw indexer cache every QSA layer per token, which keeps some O(n_kv) depth scaling. An incremental pooled-key cache would make QSA decode ~O(top_k); happy to discuss or attempt it.
Notes
This is the master port of unslothai#165 (same change against their staging branch, measured there before #27742 merged). nsys profiles and full benchmark methodology available on request.