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

[tvanderka]

Overview

While testing qwen4exp I noticed decode is cpu limited. Found get_prev_tokens() scans all used cells to resolve the n-gram predecessor tokens. This implements a TODO by @ngxson from llama_kv_cache::get_prev_tokens, mostly just to show the impact.

On my test HW with 2xL40s this PR achieves 2.7x speedup at 240k ctx.

Additional information

llama_kv_cells now maintains a per-seq index of cell rows per position (seq_pos: pos -> set) instead of refcount. Updated by the existing seq_pos_inc/dec, so add/remove/defrag/copy (memory_seq_cp) paths update it. Multimodal falls back to original scan.

There is a bit of debug/verification plumbing, using 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
• off just maintains index, its not used

Verify plumbing would be removed if this is the way. Test suite passed.

This adds a small memory/cpu cost for every other model, while only qwen4exp is using it.

measured (qwen4exp, 2xL40S, unified KV, 256k ctx):

ctx	verify (≈stock)	fast	gain
0	61.13 ± 1.66	61.12 ± 1.35	1.00×
4096	57.11 ± 1.47	61.20 ± 1.95	1.07×
16384	43.86 ± 0.88	55.59 ± 1.37	1.27×
65536	24.00 ± 0.25	44.14 ± 0.93	1.84×
131072	14.93 ± 0.12	34.99 ± 0.54	2.34×
240000	8.92 ± 0.10	24.25 ± 0.38	2.72×

numbers from: llama-bench -hf "unsloth/Qwen3.8-Flash-Next-GGUF:UD-Q4_K_XL" -ngl 999 -fa 1 -p 0 -n 32 -fitt 128 -d 0,4096,16384,65536,131072,240000

Possibly related upstream work: #27941 (qsa correctness; likely fixes the 65535 gridDim.y abort at n_kv 262144), #27977 (shrinks the general scan constants + qsa gather windows). complementary; can be combined.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, debugging, tracing, coding possible solutions, verification, explanations using Qwen38-Next-Flash itself

[nibor1896]

Second data point, different hardware class and a real workload rather than
llama-bench. Independent of the author.

Setup

	this report	PR description
GPU	1× RTX 5090, 32,607 MiB	2× L40S
placement	-ncmoe 40, experts of 40/48 layers in system RAM	-ngl 999, GPU-resident
host RAM	2× 32 GB DDR5-5600, dual channel (89.6 GB/s)	—
quant	UD-Q2_K_XL	UD-Q4_K_XL
workload	llama-server, real agent turns with tool calls	llama-bench -p 0 -n 32
base	pin 6c84c7d5d (#27742), this PR applied alone	master
driver	616.56	—

A/B at 32k

	control (bare pin)	pin + this PR
prefill tok/s	964.92 (936.31–981.07)	970.44 (941.07–985.32)
decode tok/s	29.05 (28.49–29.84)	32.44 (31.06–33.20)

+11.7 % decode, no overlap between the ranges. Prefill +0.6 %.

Three rounds interleaved against a same session control, one boot per run, one
31,979-token cold-cache turn each. Same day drift on this machine reached 9.4 %,
which is why the control is same session rather than a stored number.

Per-token saving 3.59 ms against the 4.45 ms the verify path reports for the
scan it removes. Gain scales with depth as described: ~3 % at a few-hundred-token
prompt, 11.7 % at 32k.

Correctness

LLAMA_KV_PREV_TOKENS=verify in the live server under an agent workload with
tool calls: 0 mismatches over 250 calls at 31,979 tokens of context.
Ten-task gate 10/10 twice, token counts byte-identical to the control.

Sustained load

One agent run, 16 milestones, driven to completion by the model itself.

span	5 h 52 m, 407 requests, no restart
decode	424,465 tokens @ 25.01 tok/s
prefill	327,878 tokens @ 415.91 tok/s
total	752,343 tokens
model time	4 h 55 m = 83.9 % of the span, 95.6 % of that decode
deepest context	183,274 tokens
rollovers	3, at 181,501 / 180,874 / 180,364

25.01 tok/s is a cross-section over every depth from 10k to 183k and is not
comparable to the 32.44 above, which is one cold 31,979-token turn.

Fitted over all 407 requests, decode ms/token against the context depth reported
by each request's own release line:

ms/token = 30.740 ms + 0.0898 ms per 1,000 context tokens
r = 0.848   r² = 0.720   n = 407   depths 9,588 - 183,274

depth	n	decode
0–30k	63	31.57 tok/s
30–60k	71	28.86 tok/s
60–90k	55	26.40 tok/s
90–120k	59	24.58 tok/s
120–150k	68	23.26 tok/s
150–200k	91	22.29 tok/s

Raw rows, one per request:
https://github.com/nibor1896/Crow/blob/main/docs/measurements/qwen4exp-depth-407.csv

Conditions and the fit:
https://github.com/nibor1896/Crow/blob/main/docs/measurements/README.md

A depth term of 0.0898 ms per 1,000 survives this PR. At 180k that is
16.2 ms of a 46.9 ms token, 34.5 %. The index removes the O(n_kv) scan;
what is left still grows with depth.

Single arm, no control at these depths, this is not a second speedup
measurement. It supports the shape, the stability of the build across six hours
and 407 requests without a restart, and the size of what remains.

Found during testing

1. tests/test-kv-prev-tokens.cpp does not compile as shipped.
It includes the internal src/llama-kv-cells.h; every other test uses the
public include/llama-cpp.h. tests/CMakeLists.txt registers the test but adds
no search path, so it fails with C1083 under MSVC. A
target_include_directories for this target pointing at src/ was enough here.
After that: 9,480 lookups, 0 failures.

2. The verify heartbeat cannot fire on a normal turn.
It reports every 2,000 calls. A turn here makes ~300, so "no message" reads as
"no mismatch" a reviewer running verify against a normal workload sees silence
either way. Lowering the cadence is what made the 250-call figure above
observable.

Architectures

The PR notes it charges every other model a little. One measurement on this
line: prefill flat at +0.6 %, inside the run-to-run range. Not a multi-arch
survey.

Not measured

• decode at a full 200k window as a controlled A/B
• any A/B above 32k, the depth series above is a single arm
• CPU-side cost on architectures other than qwen4exp

Requirements

I have read and agree with the contributing guidelines

AI Disclosure : Claude helped me Coding, build tables (like the csv file) + I reviewed, wrote code, wrote this text (ai helped translating some parts since no native english speaker).

[tvanderka]

Better fix in #28040 and performance within margin of error.