Skip to content

kv-cells: resolve get_prev_tokens in O(log n) from the sequence position index #28040

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ServeurpersoCom merged 1 commit into from
Sep 1, 2026
+33 −95
Merged

kv-cells: resolve get_prev_tokens in O(log n) from the sequence position index #28040

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 5 additions & 54 deletions src/llama-kv-cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
#include "llama-context.h"

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstring>
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}
}
Expand Down
69 changes: 28 additions & 41 deletions src/llama-kv-cells.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
#include <bitset>
#include <cassert>
#include <cstring>
#include <map>
#include <limits>
#include <set>
#include <vector>

Expand Down Expand Up @@ -248,7 +248,7 @@ class llama_kv_cells {
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;
Expand All @@ -272,7 +272,7 @@ class llama_kv_cells {
seq[i].reset();

seq[i].set(seq_id);
seq_pos_inc(seq_id, pos[i]);
seq_pos_inc(seq_id, i);

return false;
}
Expand Down Expand Up @@ -318,28 +318,22 @@ class llama_kv_cells {
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;
}

const auto m = seq[i] & seqs;
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);

// a cell carries a handful of sequences at most, out of LLAMA_MAX_SEQ
size_t left = m.count();
const auto & sp = seq_pos[seq_id];

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
Expand All @@ -349,7 +343,7 @@ class llama_kv_cells {
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
Expand All @@ -376,8 +370,6 @@ class llama_kv_cells {
return -1;
}

assert(seq_pos[seq_id].begin()->second > 0);

return seq_pos[seq_id].begin()->first;
}

Expand All @@ -391,8 +383,6 @@ class llama_kv_cells {
return -1;
}

assert(seq_pos[seq_id].rbegin()->second > 0);

return seq_pos[seq_id].rbegin()->first;
}

Expand Down Expand Up @@ -523,36 +513,33 @@ class llama_kv_cells {
// 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);
}
}
}
Expand All @@ -561,7 +548,7 @@ class llama_kv_cells {
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);
}
}
}
Expand Down
Loading