qwen4exp: attend the selected cells instead of masking the window - #28244
Open
ServeurpersoCom wants to merge 1 commit into
Open
qwen4exp: attend the selected cells instead of masking the window#28244ServeurpersoCom wants to merge 1 commit into
ServeurpersoCom wants to merge 1 commit into
Conversation
QSA names about n_top_k cells per query, but build_attn_qsa turned that selection into a mask over the whole window and handed the full K and V to flash attention. On CUDA the only skip is a trailing cut, and the selection always keeps the current tail, so it never fires: the sparsity saved no work at all. The selected cells are now gathered into a window of their own, one per query, and the queries ride the stream axis of the attention so each one carries the window its own selection named. The mask comes from gathering the attention mask at those cells, which leaves the same values the mask path would put there. The scan reads the window once for all the queries of a stream while the gather moves the selected cells twice, so the two meet at 2*n_tps*width == n_kv. Below that margin the scan still wins, which keeps prompt processing on the existing path; only generation and short chunks take the gather. Flash attention is required, since the gather reads the value side as rows. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs: 132k context generation 56.7 -> 61.1 t/s Nothing at 55k, and the gain grows with depth, as the scan cost follows n_kv while the gather cost does not. Prompt processing is unchanged. Retrieval from a 132k token context is unaffected.
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.
Overview
Reduce the generation slowdown as context grows. Split out of #27977 as requested, one PR per change.
Optimize Qwen3.8 Flash Next decode at depth: attend only the top-k cells QSA selected instead of masking the whole KV cache (+8% at 50k).
Additional information
QSA picks the top-k cells (2048 + tail) per query, but attention still ran over the whole KV cache with the unselected cells masked out, so the cost kept growing with context depth even though only ~2k cells matter.
This gathers the selected K/V cells (and their mask entries) into a compact buffer with get_rows and runs flash attention over just those. The gather path is taken at decode when the cache is at least 4x deeper than the top-k width; prompt processing keeps the masked path.
RTX PRO 6000, UD-Q4_K_XL, q8_0 KV, tg 1024:
Greedy output is unchanged at short context. At depth the reduction order in flash attention differs (2k gathered cells vs the masked window), so outputs are not bit-identical to master there.
Requirements