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);
}
}
}
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);
}
}
}