qwen4exp: compute QSA cache-layout inputs once per ubatch, not per layer - #150
Open
talexeyft wants to merge 2 commits into
Open
qwen4exp: compute QSA cache-layout inputs once per ubatch, not per layer#150talexeyft wants to merge 2 commits into
talexeyft wants to merge 2 commits into
Conversation
set_input_qsa results depend only on the indexer cells and the ubatch, yet the graph asked for them once per attention layer (48x), each time running O(n_kv) serial host loops and filling a separate set of input tensors. At long context this dominated decode: GPU util sat at ~31% while host code burned tens of ms per token. Two changes: 1. Share one set of QSA input tensors (k_idxs, cell_blk, blk_cells, blk_pos, bias) across all layers with the same compress ratio (graph::qsa_shared). One fill + one H2D copy per ubatch instead of 48. 2. Memoize set_input_qsa on llama_memory_hybrid_idx, keyed by a fingerprint of per-stream cell generations (new llama_kv_cells::get_generation(), bumped by every mutator) and the ubatch tokens. Decode-step repeats hit the memo; prefill batches above a 256MB bias cap bypass it. Measured on Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX 3090 (NVLink), f16 KV: decode tps 32k 131k 176k 229k before 29.1 16.2 OOM - after 32.1 21.1 18.2 15.5 prefill tps 32k 131k 176k 229k before 336 120 - - after 492 315 260 220
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
2 participants
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.
qwen4exp: compute QSA cache-layout inputs once per ubatch, not per layer
Problem
set_input_qsaresults depend only on the indexer cells and the ubatch, yet the graphasked for them once per attention layer (48x), each time running O(n_kv) serial host
loops and filling a separate set of input tensors. At long context this dominated
decode: GPU util sat at ~31% while host code burned tens of ms per token.
Changes
across all layers with the same compress ratio (
graph::qsa_shared). One fill +one H2D copy per ubatch instead of 48.
set_input_qsaonllama_memory_hybrid_idx, keyed by a fingerprint ofper-stream cell generations (new
llama_kv_cells::get_generation(), bumped by everymutator) and the ubatch tokens. Decode-step repeats hit the memo; prefill batches
above a 256MB bias cap bypass it.
Measured
Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX 3090 (NVLink), f16 KV cache:
Full write-up (sweeps, nsys profiles, remaining bottlenecks):
https://github.com/talexeyft/llama.cpp/tree/qwen4exp-qsa-perf (branch commit message),
local: admin_gpu/ai_experiments/qwen38_flash_next_speed/REPORT.md