qwen4exp: reduce the generation slowdown as context grows by ServeurpersoCom · Pull Request #27977 · ggml-org/llama.cpp

ServeurpersoCom · 2026-08-29T17:39:25Z

Overview

qwen4exp: reduce the generation slowdown as context grows

Additional information

Five things I found while profiling the long context slowdown on this model.

The big one is the n-gram predecessor lookup. It was walking every cell of the cache and checking all 256 possible sequences, when a cell basically always belongs to one. At 88k that's 20 million iterations per token just to find three positions. Now it stops as soon as it has seen the sequences the cell actually holds. Generation went from 44 to 63 t/s and the GPU from 62% to 91% busy, which is what put me on the trail in the first place: the card was drawing 315W instead of 570W during generation.

Then the attention. QSA picks around 2000 cells per query, but that selection was only used to build a mask, and the whole cache was still handed to flash attention. On the CUDA side the only thing that gets skipped is a trailing cut, and the selection always keeps the tail of the context, so it never skipped anything. The sparsity was doing nothing. The selected cells now get gathered into their own window per query. Generation only, prefill still scans because gathering 2048 windows costs more than reading everything. Worth 6.7% at 132k, nothing at 55k, and it grows with depth.

Third one is in the indexer. Summing the four heads went through a transpose plus a sum_rows with ne0 = 4, which is about the worst shape for that kernel, and the transpose was copying a 671 MB surface twice for nothing. Three adds over strided views do the same thing.

Fourth is the same kind of thing somewhere else. The predecessor scan started at position 0 when we only care about the last few positions. The full range was only there for a fallback that fires on position gaps, which never happens on text, so it only runs when it's actually needed now.

And a small one: the set of used cells was a std::set, so one tree node per cell with no locality at all. It's a bitmap now, walked 64 bits at a time.

Tested on a real 149k agentic conversation, applying and removing the patches both ways to be sure: generation 30 to 56 t/s, prefill 1617 to 2148 t/s.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES

The indexer score reduction went through a permute, a cont and a sum_rows over the head axis. That left sum_rows with ne0 = 4, one block per row for a four element reduction, and the transpose copied the whole block by token surface twice on the way in. The heads are adjacent on ne[1], so each one is a strided view and the sum is a short chain of adds, the same shape as the pooling loop above it. Measured on RTX PRO 6000, UD-Q4_K_XL, fa on, per_layer_token_embd on CPU, llama-bench r=2: pp2048 @ d32768 2163 -> 2356 t/s pp2048 @ d65536 1497 -> 1666 t/s tg32 @ d32768 67 -> 69 t/s tg32 @ d65536 45 -> 48 t/s Greedy output is unchanged token for token on an 80k token retrieval probe.

for_each_token_in() walked the full LLAMA_MAX_SEQ width for every used cell. get_prev_tokens(), its only caller, passes p0 = 0, so no cell is filtered out by position: the scan visits every cell of the cache to keep the few positions of the n-gram window. A cell carries a handful of sequences at most, so the scan stops once their count is reached. Behaviour is unchanged. The cost this removes is proportional to the number of used cells, so it shows up at long context and is invisible on short prompts. It applies to any n-gram model that resolves predecessors from the KV cells, not just qwen4exp.

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 flash_attn_mask_to_KV_max, a trailing cut, and the selection always keeps the current tail, so the cut never fires: the selection 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 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 prefill on the existing path; only decode and short chunks take the gather. Flash attention is required, since the gather reads the value side as rows. The win grows with depth, as the scan cost follows n_kv while the gather cost does not: nothing at 55k context, and generation goes from 52.05 to 55.55 t/s at 132k. Prefill is unchanged. Retrieval from a 132k token context is unaffected.

The scan ran over the whole position range, from 0 up to p_max, so every cell of the cache passed the position filter and paid the sequence mask work. The full range was only needed for below[], which holds the nearest token before the window and is read solely when a lookup finds nothing in [w0, p], that is on an M-RoPE gap. The main scan now covers the window alone, and the below[] pass is deferred to the first lookup that falls through. Contiguous positions never trigger it, so text decoding pays one windowed pass instead of a full one. Measured on Qwen3.8-Flash-Next at 55k context on an RTX PRO 6000, warm runs with the first one discarded: generation 74.4 -> 76.3 t/s Prompt processing is unaffected, since the scan is amortised over the whole ubatch there. What is left of this call is the traversal of the used-cell set itself, which the range restriction cannot avoid.

used was a std::set<uint32_t>, so a full pass over the used cells chased one tree node per index with no locality. The set is dense, a large fraction of the cache, which makes a bitmap both smaller and far cheaper to walk: one word per 64 cells and a trailing-zero count per set bit. llama_kv_idx_set exposes the same operations the class relied on, and visits indices in ascending order like the set did. insert, erase and size stay O(1); first and last scan the words and are called once per ubatch. Measured on Qwen3.8-Flash-Next at 55k context on an RTX PRO 6000, warm runs with the first one discarded: generation 76.2 -> 77.4 t/s The gain scales with the number of used cells, so it grows with context and is invisible on short ones. This touches the KV cache for every model, though only the n-gram path walks the set often enough to notice.

Green-Sky · 2026-08-29T20:11:38Z

Do we have 3 prs doing some of the same things now?
#27879
#27941

ServeurpersoCom · 2026-08-29T20:19:08Z

Do we have 3 prs doing some of the same things now? #27879 #27941

The three don't do the same thing. Mine is pure performance work, and four of my five commits don't touch any shared area: for_each_token_in and get_prev_tokens aren't modified by either #27879 or #27941, and neither is the used-cell bitmap.

The real overlap is on build_qsa_top_k, which all three PRs modify. In particular my "attend the selected cells instead of masking the window" commit builds on build_attn_qsa and the per-cell top_k, which #27879 replaces with a mask built differently. That one will conflict head-on, not cosmetically.

Given that #27879 also fixes an n_blocks calculation and adds tests, it should probably go in first, and I'll rebase that commit on top of it.

Green-Sky · 2026-08-29T20:24:20Z

The three don't do the same thing.

Yea, I said some things. :)

ServeurpersoCom · 2026-08-29T20:30:38Z

The three don't do the same thing.

Yea, I said some things. :)

Ha, I did influence my LLM on the translation :) I read it straight as "same thing" :)

Green-Sky · 2026-08-29T20:47:49Z

The three don't do the same thing.

Yea, I said some things. :)

Ha, I did influence my LLM on the translation :) I read it straight as "same thing" :)

Ah ye, technically is "some of the same". Anyway, point gotten across. :D

sammcj · 2026-08-30T03:06:08Z

Apple Silicon datapoint (M5 Max 128GB, Metal). Target: Unsloth UD-IQ4_XS with the #27836 MTP head merged into the file set, --spec-type draft-mtp --spec-draft-n-max 6 --spec-draft-p-min 0.7 --spec-draft-backend-sampling, q8_0 KV, -ub 2048, source-code prompts, both-orders A/B per config.

Baseline is #27836 + master. Also tested #27992, which attacks the same get_prev_tokens() cost via a kv-cell index.

prompt depth	PP base	PP #27977	TG base	TG #27992	TG #27977
32K	663	685 (+3%)	~33	~33	34.6
74K	488	512 (+4.8%)	24.4	28.8 (+18%)	28.8 (+18%)
115K	387	409 (+5.5%)	~25	24.8	25.9

Either PR recovers the same TG at depth (+18% at 74K over the plain PR build); head to head they're a wash on generation here. The CUDA-sized gains don't transfer - unified memory already makes the CPU-side scan comparatively cheap, so on Metal the win saturates earlier.
The PP improvement from this PR is consistent at every depth and is what kv-cache : index (seq,pos) cells to make ngram prev-token lookups O(log n) for qwen4exp decode speedup #27992 doesn't have.
Carrying both PRs together measurably regressed TG at 32K (27.6/29.0 t/s vs 34.8/34.4 for this PR alone, mirrored order) - worth noting in whichever lands second.
Output sanity-checked coherent at 32K with the MTP draft path active.

am17an · 2026-08-30T05:58:22Z

Please create one PR for each change for easier review. Preferably don't make changes to llama-kv-cache*

@200k

get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.

@16k

…Com) Early-exit ngram predecessor scan, gathered QSA decode window, indexer head-sum without transpose, bounded predecessor range, bitmap cell set. On gfx1151: depth curve flattens (tg64 @16k 17.8 -> 19.5, @32k -> 17.7) and file-rewrite @24K ctx 29.5 -> 34.8 tok/s (+18%). 8k unchanged.

drluoto · 2026-08-30T08:36:29Z

gfx1151 / Strix Halo datapoint (Ryzen AI Max+ 395, ROCm 7.1, UD-IQ4_XS 93.7 GB, applied on top of #27836 + #27941 + the hipCUB TOP_K path):

llama-bench tg64, no speculation:

depth	before	this PR
1024	22.70	22.83
4096	21.38	22.04
16384	17.80	19.48 (+9%)
32768	~15.2	17.68 (+17%)

End-to-end with draft-mtp,ngram-mod (n-max 6, p-min 0.7), real coding workloads, greedy:

file rewrite @ 24k-token prompt: 29.5 → 34.8 tok/s (+18%)
new code @ 24k: 25.1 → 25.9
@8k: unchanged within noise (43–45 / 30 / 32 / 26)

Greedy output verified clean at the 2.7k-token threshold. The pattern matches your profile: nothing at shallow depth, growing with context — on unified memory the ngram predecessor scan was stealing CPU the drafters need. Nice find; this is in our production stack as of today.

ServeurpersoCom · 2026-08-30T08:49:22Z

#28011 Split out this #27977 as requested, one PR per change. This is the quick win of the series: a handful of lines, and the largest gain of the whole set.

Squashed so the set reverts as one on rebase. Takes the line descending from the merged ggml-org#27742: ggml-org#27836 (NextN/MTP draft head) and ggml-org#27941 (danielhanchen's follow-up fixes), plus ggml-org#27977 and fifteen others. Drops ggml-org#27879 (third-party correctness fixes) -- it conflicts with ggml-org#27836 in qwen4exp.cpp, and MTP is not optional here. ggml-org#27836 supplies everything the MTP head needs: it reads nextn_predict_layers, loads the nextn block behind ml.load_mtp, and adds the MTP graph. Earlier local patches for those are dropped as redundant. Still ours: the two --ngram-on-disk bugs from halo-box#8, which live in code these PRs do not touch. ple_head_offsets was applied twice, and can_reuse dereferenced rows -- null whenever the table is on disk.

The head reduction went through a transpose and a sum_rows over ne[1], which left sum_rows with ne0 = 4, one block per row for a four element reduction, and the transpose copied the whole block by token surface twice on the way in. The heads are adjacent on ne[1], so each one is a strided view and the sum is a short chain of adds. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, 55k context, warm runs on top of ggml-org#27977: prompt processing 2170 -> 2366 t/s Generation is unaffected. The removed work scales with n_blocks by n_tokens, so the gain grows with context and with ubatch size.

@200k

get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.

@200k

get_prev_tokens() scans all used cells once per ubatch to resolve the n-gram predecessor tokens (qwen4exp PLE). the scan is O(used) per call and dominates decode at long context: measured ~46 ms per token at ~170k ctx on 2xL40S (out of ~89 ms/token total, ~65% of wall time). llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set<rows>), updated by the existing seq_pos_inc/dec funnels, so add/remove/defrag/copy (memory_seq_cp) paths update it by construction. prev_token(p) resolves to the token of the largest existing position <= p, with the 'last cell wins' tie-break of the general scan preserved. ubatches with shared temporal positions (multimodal) are detected as not-applicable and keep using the general scan: within a shared position, cells resolve by ubatch order, not by row. LLAMA_KV_PREV_TOKENS env: fast (default) | verify | off. verify runs both paths for every call and logs mismatches + a cost heartbeat; used to validate the index against production traffic (455k lookups across prompt-cache loads and checkpoint restores, 0 mismatches; scan 45.9 ms vs index 6.9 us avg). measured (qwen4exp, 2xL40S, unified KV, 256k ctx): tg @155k ctx: ~11.7 -> ~31 t/s tg @200k ctx: ~11.5 -> ~27 t/s related upstream work: ggml-org#27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), ggml-org#27977 (shrinks the general scan constants + qsa gather windows). complementary layers; can be combined.

The scan ran from position 0 over every used cell (and walked the full LLAMA_MAX_SEQ width per cell) although only cells inside the n-gram window feed hist[], and the before-window below[] fallback only fires on M-RoPE gaps, which contiguous text never produces. The main scan now covers the window alone and the below[] pass is deferred to the first lookup that falls through. Reference: ggml-org/llama.cpp#27977 (28e3792). Semantics preserved: PPL 3.7892 +/- 0.033 on wiki.test (60 chunks), identical to baseline. Measured with the pooled QSA cache active (qwen4exp, APU 8060s, 48k ctx): decode 9.1 -> 16.5 t/s (+81%), context decay -54% -> -19%. Assisted-by: Sisyphus

The ggml-org#27977 pick assumed upstream's w0/below refactor, which our get_prev_tokens predates; replace the whole function with the PR's final form (db40b22 state) so the window base and the deferred below[] pass are defined. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

LynxPDA · 2026-09-01T09:00:19Z

@ServeurpersoCom, Saw you were profiling the long-context slowdown and dealing with the pain of loading the full 103GB model for testing. I put together a 324MB micro GGUF that keeps the exact QSA/GDN architecture dimensions (same indexer top_k=2048, compress_ratio=4, etc.) but with random weights.

It reproduces the same performance degradation curve as the full model, so you can quickly profile decode/prompt-eval at various context depths without needing multi-GB downloads or massive VRAM. Might be useful for validating these optimizations.

https://huggingface.co/Lynxpda/micro-qwen4exp

ServeurpersoCom · 2026-09-01T09:29:56Z

@ServeurpersoCom, Saw you were profiling the long-context slowdown and dealing with the pain of loading the full 103GB model for testing. I put together a 324MB micro GGUF that keeps the exact QSA/GDN architecture dimensions (same indexer top_k=2048, compress_ratio=4, etc.) but with random weights.

It reproduces the same performance degradation curve as the full model, so you can quickly profile decode/prompt-eval at various context depths without needing multi-GB downloads or massive VRAM. Might be useful for validating these optimizations.

https://huggingface.co/Lynxpda/micro-qwen4exp

That could be useful, and it should let me profile things I could not reach before and spot other bottlenecks. At that size the matmuls disappear from the profile and only what scales with context depth is left, which is exactly what I am chasing. Random weights are not a problem for that either: for a refactor the check is base against patched on the same prompt in greedy, and the garbage just has to be the same garbage on both sides.

The one thing it cannot cover is MTP acceptance: random weights the acceptance rate would mean nothing anyway, so that part still has to be measured on the real model.

Small thing: the table in the card does not match the file, which has 4 blocks, n_embd 768, 4 indexer heads and n_ctx 65536.

ServeurpersoCom · 2026-09-01T10:09:37Z

Closed because it was split into mini-PRs.

LynxPDA · 2026-09-01T10:18:05Z

@ServeurpersoCom Thanks, that's exactly the use case. Agreed on MTP acceptance — random weights make the rate meaningless, that part still has to be measured on the real model.

The card/file mismatch you caught was real: the first upload was a 4-layer prototype (n_embd 768, n_ctx 65536) while the table described a different draft. That's fixed. The repo now matches the files, and the topology is closer to Qwen3.8-Flash-Next:

trunk is 12 layers = 3 repeats of the real GDN, GDN+PLE, GDN, QSA cycle (the 56B has 12 of those)
QSA geometry is 1:1 with the full model: 24/2 heads, d_h=256, 4 indexer heads, key_length=128, top_k=2048, compress_ratio=4, n_ctx=262144, M-RoPE sections [11, 11, 10, 0]
same split layout as Unsloth: trunk, detached MTP sidecar (blk.12 NextN head), and a qwen3vl_merger mmproj
Q4_K_M of trunk + MTP via llama-quantize (indexer stays F16, same as the real file keeping those BF16)

So you can profile decode/pp vs context, greedy A/B of a refactor (same garbage both sides), and the MTP graph/rollback path with --mtp -md …. Acceptance quality still needs the real checkpoint.

https://huggingface.co/Lynxpda/micro-qwen4exp

Port upstream PR ggml-org#27977 (serveurperso) onto the bmoe/expert-ready-hook line: - qwen4exp: build_qsa_gather gathers only the selected cells per query via ggml_get_rows(k_cells, idx_stream) instead of masking the full cache window; decode dispatch picks the gather path when flash_attn and 4*n_tps*width < n_kv (windows that fit a decode graph's compute buffer), else falls back to the scan path - qwen4exp: sum the indexer heads by slices instead of a transpose that carried the whole block-by-token surface twice over - kv-cells: keep the used-cell set as a bitmap (llama_kv_idx_set) and stop the sequence scan once all sequences are seen - kv-cache: scan only the n-gram window in get_prev_tokens Verified: test-backend-ops 900/900 MUL_MAT_VEC_FUSION and 416/416 TOPK_MOE on CUDA0 (incl. the (32,8)/(32,9) boundary cases); test-recurrent-state-rollback regression identical to the pre-port baseline; bmoe byte-identity gates 13/16 (the three pre-existing streaming-init failures only).

ServeurpersoCom added 5 commits August 29, 2026 19:02

ServeurpersoCom requested review from CISC and ggerganov as code owners August 29, 2026 17:39

github-actions Bot added the model Model specific label Aug 29, 2026

nits

db40b22

tvanderka mentioned this pull request Aug 29, 2026

kv-cache : index (seq,pos) cells to make ngram prev-token lookups O(log n) for qwen4exp decode speedup #27992

Closed

ServeurpersoCom mentioned this pull request Aug 30, 2026

kv-cells: stop the sequence scan once all sequences are seen #28011

Merged

ServeurpersoCom marked this pull request as draft August 30, 2026 08:48

ServeurpersoCom mentioned this pull request Aug 30, 2026

qwen4exp: sum the indexer heads by slices #28023

Merged

Nathanw1014 mentioned this pull request Aug 30, 2026

WIP: DSv4 + Qwen3.8 Flash runtime fixes (MTP spec decode, converters, checkpoint-on-device, GDN graph order) halo-box/strix-llama.cpp#14

Closed

ServeurpersoCom closed this Sep 1, 2026

This was referenced Sep 2, 2026

qwen4exp: attend the selected cells instead of masking the window #28244

Open

kv-cells: keep the used-cell set as a bitmap #28245

Closed

Navigation Menu

Uh oh!

qwen4exp: reduce the generation slowdown as context grows - #27977

qwen4exp: reduce the generation slowdown as context grows#27977
ServeurpersoCom wants to merge 6 commits into
ggml-org:masterfrom
ServeurpersoCom:qwen4exp-optimize

ServeurpersoCom commented Aug 29, 2026

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

ServeurpersoCom commented Aug 29, 2026

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

ServeurpersoCom commented Aug 29, 2026

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

sammcj commented Aug 30, 2026

Uh oh!

am17an commented Aug 30, 2026

Uh oh!

drluoto commented Aug 30, 2026

Uh oh!

ServeurpersoCom commented Aug 30, 2026

Uh oh!

LynxPDA commented Sep 1, 2026

Uh oh!

ServeurpersoCom commented Sep 1, 2026 •

edited

Loading

Uh oh!

ServeurpersoCom commented Sep 1, 2026

Uh oh!

LynxPDA commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

ServeurpersoCom commented Aug 29, 2026

Overview

Additional information

Requirements

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

ServeurpersoCom commented Aug 29, 2026

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

ServeurpersoCom commented Aug 29, 2026

Uh oh!

Green-Sky commented Aug 29, 2026

Uh oh!

sammcj commented Aug 30, 2026

Uh oh!

am17an commented Aug 30, 2026

Uh oh!

drluoto commented Aug 30, 2026

Uh oh!

ServeurpersoCom commented Aug 30, 2026

Uh oh!

LynxPDA commented Sep 1, 2026

Uh oh!

ServeurpersoCom commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ServeurpersoCom commented Sep 1, 2026

Uh oh!

LynxPDA commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ServeurpersoCom commented Sep 1, 2026 •

edited

Loading