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-cells.h
+28 -41
View File
@@ -6,7 +6,7 @@
#include <bitset>
#include <cassert>
#include <cstring>
#include <map>
#include <limits>
#include <set>
#include <vector>
@@ -248,7 +248,7 @@ public:
assert(seq_id >= 0);
seq[i].reset(seq_id);
seq_pos_dec(seq_id, pos[i]);
seq_pos_dec(seq_id, i);
if (seq[i].none()) {
pos[i] = -1;
@@ -272,7 +272,7 @@ public:
seq[i].reset();
seq[i].set(seq_id);
seq_pos_inc(seq_id, pos[i]);
seq_pos_inc(seq_id, i);
return false;
}
@@ -318,28 +318,22 @@ public:
return seq[i].test(seq_id);
}
// gather the token ids of the cells in `seqs` with position in [p0, p1)
// the callback receives (seq_id, pos, token) for every such (cell, seq) pair
// the token of the cell of sequence seq_id at the largest position <= p
// when several cells share that position, the one with the highest index wins
// return LLAMA_TOKEN_NULL if the sequence has no cell at or before p
// note: used by n-gram input embeddings to recover the tokens preceding a ubatch
template<typename F>
void for_each_token_in(const std::bitset<LLAMA_MAX_SEQ> & seqs, llama_pos p0, llama_pos p1, F && f) const {
for (const auto & i : used) {
if (pos[i] < p0 || pos[i] >= p1) {
continue;
}
llama_token seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const {
assert(seq_id >= 0);
assert(seq_id < LLAMA_MAX_SEQ);
const auto m = seq[i] & seqs;
const auto & sp = seq_pos[seq_id];
// a cell carries a handful of sequences at most, out of LLAMA_MAX_SEQ
size_t left = m.count();
for (llama_seq_id s = 0; left > 0 && s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) {
if (m.test(s)) {
f(s, pos[i], ext[i].tok);
--left;
}
}
auto it = sp.upper_bound({ p, std::numeric_limits<uint32_t>::max() });
if (it == sp.begin()) {
return LLAMA_TOKEN_NULL;
}
return ext[(--it)->second].tok;
}
// note: call only if the cell is not empty and the seq_id is not in the cell
@@ -349,7 +343,7 @@ public:
assert(!seq[i].test(seq_id));
seq[i].set(seq_id);
seq_pos_inc(seq_id, pos[i]);
seq_pos_inc(seq_id, i);
}
// return the sequence id of this cell
@@ -376,8 +370,6 @@ public:
return -1;
}
assert(seq_pos[seq_id].begin()->second > 0);
return seq_pos[seq_id].begin()->first;
}
@@ -391,8 +383,6 @@ public:
return -1;
}
assert(seq_pos[seq_id].rbegin()->second > 0);
return seq_pos[seq_id].rbegin()->first;
}
@@ -523,36 +513,33 @@ private:
// the bitset seq[i] tells us which sequences are currently occupying the i-th cell
std::vector<seq_set_t> seq;
// the set seq_pos[s][p] tells us how many times the position p is currently present for sequence s
// if the position p is not present, seq_pos[s][p] is not set
// the set seq_pos[s] holds one (pos, cell) pair per cell that carries sequence s, ordered by position
// this way seq_pos[s].begin() and seq_pos[s].rbegin() give us the min/max positions currently in the cache
// and upper_bound() on a position finds the nearest cell of the sequence in logarithmic time
//
// note that we cannot a use an std::set because in some cases a position can occur more than once for the same seq:
// the cell index is part of the key because a position can occur more than once for the same seq:
// - during performing a cache reuse via (rm + add)
// - some vision models have input embeddings with repeating positions
//
std::map<llama_pos, int> seq_pos[LLAMA_MAX_SEQ];
std::set<std::pair<llama_pos, uint32_t>> seq_pos[LLAMA_MAX_SEQ];
// helper functions for updating `seq_pos`, once cell at a time:
void seq_pos_dec(llama_seq_id s, llama_pos p) {
auto it = seq_pos[s].find(p);
assert(it != seq_pos[s].end());
if (--it->second == 0) {
seq_pos[s].erase(it);
}
void seq_pos_dec(llama_seq_id s, uint32_t i) {
const auto n = seq_pos[s].erase({ pos[i], i });
assert(n == 1);
GGML_UNUSED(n);
}
void seq_pos_inc(llama_seq_id s, llama_pos p) {
seq_pos[s][p]++;
void seq_pos_inc(llama_seq_id s, uint32_t i) {
seq_pos[s].insert({ pos[i], i });
}
// remove cell i
void seq_pos_rm(uint32_t i) {
for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {
if (seq[i].test(s)) {
seq_pos_dec(s, pos[i]);
seq_pos_dec(s, i);
}
}
}
@@ -561,7 +548,7 @@ private:
void seq_pos_add(uint32_t i) {
for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {
if (seq[i].test(s)) {
seq_pos_inc(s, pos[i]);
seq_pos_inc(s, i);
}
}
}