kv-cells: resolve get_prev_tokens in O(log n) from the sequence position index - #28040
Conversation
|
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. |
|
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 : |
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.
e68fafc to
db997ff
Compare
|
Rebased for a quick test |
|
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
left a comment
ggerganov
left a comment
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.
4 participants
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