qwen4exp: gather-based sparse attention for QSA decode (+ cuda top-k batching fix) - #165
Open
abdel-darwish-27 wants to merge 4 commits into
Open
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.
…iceTopK With CCCL >= 3.2 available, ggml_cuda_op_top_k ran cub::DeviceTopK::MaxPairs in a serial per-row loop. Decode-shaped calls (nrows == 1) are fine, but prompt processing hands this op hundreds of rows: profiling one 141K-token prefill of a qwen4exp model showed 903,702 DeviceTopK invocations (512-row batches x 12 QSA layers x ubatches), each launching 3-4 kernels. Keep DeviceTopK for nrows <= 4 and fall back to the segmented-sort path for larger batches until DeviceSegmentedTopK exists (NVIDIA/cccl#6391).
llama_kv_cache::get_prev_tokens (the n-gram/PLE predecessor lookup, called once per ubatch) iterated the `used` std::set — an RB-tree walk over every used cell — and then tested all LLAMA_MAX_SEQ (256) bits of each matching cell's sequence bitset. At 141K context this made a single call cost ~40 ms and consume 58% of total decode CPU time (perf: 50.9% llama_kv_cache::get_prev_tokens + 7.1% std::_Rb_tree_increment), leaving both GPUs ~15% utilized. Replace the tree walk with a contiguous scan of the pos/seq/ext arrays (same cells, same index order; empty cells have pos == -1) and hoist the queried seq ids out of the per-cell loop (queries carry a handful of ids, usually one). Decode at 141K context (Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX A6000): masked attention path: 9.0 -> 13.1 tok/s gather path: 11.6 -> 21.7 tok/s Outputs byte-identical before/after (greedy, fixed seed). This helps every model that uses get_prev_tokens (all PLE/n-gram architectures) at long context.
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.
Author
|
Heads-up: qwen4exp landed on ggml-org master via ggml-org#27742 while this was in flight, so I have ported this work to mainline and re-measured everything there (master
On master the combined fixes take 129.6K-context decode from 16.9 to 29.8 tok/s (+75%) and long prefill +25%, with retrieval/factual outputs byte-identical. Happy to keep this PR open for the fork branch or close it in favor of the mainline ones — your call. |
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.
Summary
At decode time the 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. On 2x RTX A6000 (sm_86), UD-IQ4_XS, decode falls from 53 tok/s shallow to ~7 tok/s at 141K context; nsys showsflash_attn_ext_f16at ~15 ms per QSA layer per token.This PR adds a decode-only gather path (541f63a, details in the commit message): block-level top-k -> gather the selected 2048 cells + their mask values -> dense attention over the gathered set.
QWEN4EXP_QSA_GATHER=0opts out at runtime.Three supporting fixes surfaced while profiling the same path:
ggml_cuda_op_top_kloopedDeviceTopK::MaxPairsper row; a 512-token prompt ubatch makes that 512 serial invocations per QSA layer. Now dispatches to the segmented-sort path for nrows > 4 (untilDeviceSegmentedTopKlands, [RFE] AddDeviceSegmentedTopKNVIDIA/cccl#6391).for_each_token_initerated theusedstd::set (RB-tree) and tested allLLAMA_MAX_SEQbits per cell; via the PLE predecessor lookup this was ~58% of decode CPU at 141K. Now a bounded contiguous array scan with the queried seq ids hoisted. Not qwen4exp-specific — happy to split this into its own PR.n_kvkq_mask (~18 MB of host->device mask traffic per token at 141K).Correctness: greedy outputs are byte-identical to the masked path at 75K and 141K depth; buried-needle retrieval passes in both modes at all depths.
Decode throughput (single stream, q8_0 KV cache), all four commits vs baseline: 34K 17.4 -> 27.8, 68K 12.0 -> 28.5, 141K 7.1 -> 18-22 tok/s (2.5-3.1x at 141K). Shallow decode unchanged (53 tok/s). As a control, dense models on the same box (a 35B and a 27B) decay only 24-29% over the same depth range, vs 59% for this arch before the patch.
Known limitation / follow-up: both paths still recompute pooled block keys from the raw indexer cache every 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.
nsys profiles and the full benchmark methodology (contention-tagged interleaved A/B, byte-exact output comparison) available on request.