kv-cells: resolve get_prev_tokens in O(log n) from the sequence position index - #28040
Merged
ServeurpersoCom merged 1 commit intoSep 1, 2026
Merged
Conversation
Contributor
Author
|
Net effect on the tree is -62 lines: the hash map, the M-RoPE fallback array and the descending walk in llama-kv-cache.cpp all become unnecessary once the index answers the question directly. I need to check the mtmd to be sure. -> Multimodal is the case the removed fallback existed for, so it was checked explicitly: single image and two images of different sizes, greedy, byte identical output to master, with the position gaps showing up in the logs as expected. |
Contributor
Author
|
I check prefill performance, On a long-running llama-server, prefill throughput on a resumed conversation shows much lower than a full prefill on the same content after a restart. Possibly a reporting artefact of the cache reuse path rather than an actual slowdown, needs a proper look. Nothing to do with this PR. Edit: OK I found my unrelated problem. I must use this on my router for this model : |
ngxson
approved these changes
Sep 1, 2026
get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
ServeurpersoCom
force-pushed
the
kvcells-prev-tokens-index
branch
from
September 1, 2026 12:03
e68fafc to
db997ff
Compare
Contributor
Author
|
Rebased for a quick test |
CISC
approved these changes
Sep 1, 2026
Contributor
Author
|
Rebased on the latest master and rechecked out of caution, same binaries on both sides except libllama, three alternated pairs per context with the first run discarded. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, needle prompt, greedy output byte identical to master with the needle retrieved.
The runs sit within 0.1 t/s of each other, so the gap is well clear of the noise. It grows with context, which is what replacing the scan with a lookup predicts, and prefill stays where it was. Worth noting this comes on top of #28011, which already took generation from 56.3 to 74.3 t/s at 55k and from 33.6 to 50.9 t/s at 132k: the early exit removed most of the scan, and the index removes what was left of it. |
ggerganov
approved these changes
Sep 1, 2026
ggerganov
left a comment
ggerganov
left a comment
Member
There was a problem hiding this comment.
Looks OK to me. Though this reminds me we should add exhaustive unit tests for llama_kv_cells soon.
TrevorS
added a commit
to TrevorS/llama.cpp
that referenced
this pull request
Sep 4, 2026
…ationale Three fork commits are now upstream's and were dropped at the rebase onto 3466812: the rollback allowlist entry and the per-token conv snapshots (ggml-org#28123) and the windowed n-gram scan (ggml-org#28040). The env toggles that went with them are gone too. The entry for the TENSOR_READ_LAZY revert claimed the GPU read the mmap'd PLE pages directly. It does not: integrated is hardcoded false on CUDA, GET_ROWS is never offloaded, and the gather runs on one CPU thread demand-faulting the mapping. The measurement stands; the explanation did not.
thecodacus
pushed a commit
to thecodacus/llama.cpp
that referenced
this pull request
Sep 7, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
SteelPh0enix
pushed a commit
to SteelPh0enix/llama.cpp-qwen4exp
that referenced
this pull request
Sep 8, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved. (cherry picked from commit b356fa2)
SteelPh0enix
pushed a commit
to SteelPh0enix/llama.cpp-qwen4exp
that referenced
this pull request
Sep 8, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved. (cherry picked from commit b356fa2)
x1250
pushed a commit
to x1250/llama.cpp
that referenced
this pull request
Sep 9, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
zbrad
pushed a commit
to zbrad/llama.cpp
that referenced
this pull request
Sep 10, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
pl752
pushed a commit
to pl752/llama.cpp
that referenced
this pull request
Sep 15, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
TrevorS
added a commit
to TrevorS/llama.cpp
that referenced
this pull request
Sep 16, 2026
…ationale Three fork commits are now upstream's and were dropped at the rebase onto 3466812: the rollback allowlist entry and the per-token conv snapshots (ggml-org#28123) and the windowed n-gram scan (ggml-org#28040). The env toggles that went with them are gone too. The entry for the TENSOR_READ_LAZY revert claimed the GPU read the mmap'd PLE pages directly. It does not: integrated is hardcoded false on CUDA, GET_ROWS is never offloaded, and the gather runs on one CPU thread demand-faulting the mapping. The measurement stands; the explanation did not.
zsogitbe
pushed a commit
to zsogitbe/llama.cpp
that referenced
this pull request
Sep 17, 2026
…gml-org#28040) get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
adderek
pushed a commit
to adderek/llama.cpp
that referenced
this pull request
Sep 20, 2026
Cherry-picked from unslothai/llama.cpp topic branches onto d0e7ec2, then merged here. The case for taking these is CORRECTNESS, not speed — the two performance commits in the set touch the GPU path and nothing in the verification below measured throughput. 82b4198 fix llm_graph_input_ple reuse e339f5d model: qwen4exp: reduce number of graph splits (ggml-org#27880) 6f06845 qwen4exp: fix QSA correctness defects and harden metadata loading f46d73a qwen4exp: trim comments 4594877 qwen4exp: enable rms_norm + mul fusion (ggml-org#28896) Three conflicts were resolved by hand and are worth knowing about: llama-memory-hybrid-idx.cpp BOTH sides kept. Upstream adds hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE (cached indexer keys are raw; rotation happens after pooling at read time, so a K-shift must not rotate them). Our no_turbo() guard stays: the indexer keys are read back with GET_ROWS, TurboQuant stores them WHT-rotated and has no get_rows, so that one cache falls back to f16. Independent fixes at the same spot, not alternatives. qwen4exp.cpp metadata guards took upstream. GGML_ASSERT(dsv4_hc_mult > 0) becomes a throw with a readable message plus a > 1 check — transformers, vLLM and SGLang all reject hc_count <= 1. build_ple -> build_inp_ple took the upstream refactor. Our GGML_UNUSED(inp) was dropped: inp is genuinely used in the new build_ple, via build_conv_state_at. Verified on this box, 2026-09-20: - compiles and links, zero warnings under -Wall -Wextra -Wpedantic - loads Qwen3.8-Flash-Next UD-IQ4_XS (88 GB, the only qwen4exp weights here) CPU-only and generates correctly - byte-identical output against the pre-merge build (d0e7ec2) at temp 0 seed 42, on a short prompt AND on a ~2515-token prompt split into ~5 ubatches, which is where the PLE predecessor history can actually diverge - needle recall intact in the long run (a broken n-gram hash fails retrieval first, so this is the check that matters) NOT verified: anything on a GPU. No HIP build was made, so the two perf commits are unmeasured and the preserved no_turbo() guard ran with f16 KV, where it had nothing to convert. Its real test is --cache-type-k turbo4 on a card. Deliberately not taken: dbac7c4 (superseded one day later by b356fa2 / ggml-org#28040, which replaces the whole hash-map rebuild with a logarithmic lookup and carries a measured +4.9% tg at 71k ctx) and ca14269 (the tip of a 7-commit NextN/MTP chain that also needs draft weights we do not have). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
4 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.
Overview
get_prev_tokens was scanning every used cell of the cache once per ubatch to recover the few tokens preceding it. llama_kv_cells already keeps a per sequence map of positions, so it only needs the cell index next to the position to answer the same question with an upper_bound.
Putting the cell index in the key of a set rather than next to a refcount also removes the reason the counter existed in the first place, since a position occurring twice for the same sequence is now two distinct entries.
This stays inside llama_kv_cells as suggested in #28011, and llama-kv-cache.cpp loses the hash map, the M-RoPE fallback array and the descending walk, so the nearest cell at or before a position now resolves gaps through the same path as everything else.
Model side, needle retrieval at 55k context is unchanged, and 12 sequences with -np 4 greedy on a fixed seed produce byte identical output to master.
RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs, alternated A/B in both orders:
55k context generation 74.5 -> 77.6 t/s
132k context generation 52.0 -> 56.7 t/s
Additional information
Follow up to #28011
Alternative to #27992
Requirements