Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

kv-cells: look up the n-gram history in the sequence position index (#28040)

Browse Source 
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.
This commit is contained in:
Pascal
2026-09-01 20:16:07 +02:00
committed by GitHub
parent dfc29b64eb
commit b356fa2624
2 changed files with 33 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-kv-cache.cpp
+5 -54
View File
@@ -6,7 +6,6 @@
#include "llama-context.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstring>
@@ -1836,58 +1835,10 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
return;
}
// note: apply_ubatch() has already stored the current ubatch
// the window below thus covers tokens of this very ubatch as well, which is what we want
llama_pos p_min = std::numeric_limits<llama_pos>::max();
llama_pos p_max = std::numeric_limits<llama_pos>::min();
std::bitset<LLAMA_MAX_SEQ> seqs;
for (uint32_t i = 0; i < n_tokens; ++i) {
p_min = std::min(p_min, ubatch.pos[i]);
p_max = std::max(p_max, ubatch.pos[i]);
}
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
seqs.set(ubatch.seq_id_unq[s]);
}
const llama_pos w0 = p_min - (llama_pos) n;
// (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token
std::unordered_map<uint64_t, llama_token> hist;
const auto key = [](llama_seq_id seq_id, llama_pos pos) {
return ((uint64_t) seq_id << 32) | (uint32_t) pos;
};
// handle M-RoPE gaps: multiple tokens share the same temporal pos
// TODO @ngxson : improve this in the future
std::array<std::pair<llama_pos, llama_token>, LLAMA_MAX_SEQ> below;
below.fill({ -1, LLAMA_TOKEN_NULL });
for (uint32_t s = 0; s < n_stream; ++s) {
// p_max inclusive: an embd token looks up cells at its own (shared) position
v_cells[s].for_each_token_in(seqs, 0, p_max + 1,
[&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
if (pos >= w0) {
hist[key(seq_id, pos)] = tok;
} else if (pos > below[seq_id].first) {
below[seq_id] = { pos, tok };
}
});
}
// the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap
const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token {
for (llama_pos q = p; q >= w0; --q) {
const auto it = hist.find(key(seq_id, q));
if (it != hist.end()) {
return it->second;
}
}
return below[seq_id].second;
};
// note: apply_ubatch() has already stored the current ubatch, so the cells cover the tokens
// of this very ubatch as well, which is what we want
// the nearest cell at or before a position also resolves M-RoPE gaps, where multiple tokens
// share the same temporal pos
// an embd (multimodal) ubatch can repeat one position for a whole image, so positions
// do not encode the token order; resolve its predecessors by ubatch order instead
@@ -1925,7 +1876,7 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
continue;
}
res[i*n + j] = lookup(seq_id, p);
res[i*n + j] = v_cells[seq_to_stream[seq_id]].seq_pos_tok_le(seq_id, p);
}
}
}