Skip to content

llama: add token ID tracking to KV cell #27762

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
ngxson merged 4 commits into from
Aug 26, 2026
Merged

llama: add token ID tracking to KV cell #27762

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
980256b
kv: track token id
ngxson Aug 26, 2026
c0ea2b3
rm get_prev_tokens, move it to the main pr
ngxson Aug 26, 2026
d07a0ae
nits
ngxson Aug 26, 2026
2bd7078
add get_prev_tokens
ngxson Aug 26, 2026
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
4 changes: 2 additions & 2 deletions include/llama.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,10 @@
#define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'

#define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN
#define LLAMA_SESSION_VERSION 9
#define LLAMA_SESSION_VERSION 10

#define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ
#define LLAMA_STATE_SEQ_VERSION 2
#define LLAMA_STATE_SEQ_VERSION 3

#ifdef __cplusplus
extern "C" {
Expand Down
103 changes: 92 additions & 11 deletions src/llama-kv-cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include <limits>
#include <map>
#include <stdexcept>
#include <unordered_map>

static bool ggml_is_power_of_2(int n) {
return (n & (n - 1)) == 0;
Expand Down Expand Up @@ -1128,11 +1129,18 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &

cells.pos_set(idx, ubatch.pos[i]);

if (ubatch.is_pos_2d()) {
llama_kv_cell_ext ext {
/*.x =*/ ubatch.pos[i + ubatch.n_tokens*2],
/*.y =*/ ubatch.pos[i + ubatch.n_tokens],
};
if (ubatch.is_pos_2d() || ubatch.token) {
llama_kv_cell_ext ext;

if (ubatch.is_pos_2d()) {
ext.x = ubatch.pos[i + ubatch.n_tokens*2];
ext.y = ubatch.pos[i + ubatch.n_tokens];
}

if (ubatch.token) {
ext.tok = ubatch.token[i];
}

cells.ext_set(idx, ext);
}

Expand Down Expand Up @@ -1805,6 +1813,69 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
}

bool llama_kv_cache::has_cell_ext() const {
return hparams.n_pos_per_embd() > 1;
}

void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
const uint32_t n_tokens = ubatch.n_tokens;

res.clear();
res.resize(n_tokens*n, LLAMA_TOKEN_NULL);

if (n == 0) {
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]);
}

// (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;
};

for (uint32_t s = 0; s < n_stream; ++s) {
v_cells[s].for_each_token_in(seqs, p_min - (llama_pos) n, p_max,
[&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
hist[key(seq_id, pos)] = tok;
});
}

for (uint32_t i = 0; i < n_tokens; ++i) {
// TODO: a token that belongs to more than one sequence has an ambiguous history.
// the n-gram architectures have to reject such batches
const llama_seq_id seq_id = ubatch.seq_id[i][0];

for (uint32_t j = 0; j < n; ++j) {
const llama_pos p = ubatch.pos[i] - (llama_pos) (n - j);
if (p < 0) {
continue;
}

const auto it = hist.find(key(seq_id, p));
if (it != hist.end()) {
res[i*n + j] = it->second;
}
}
}
}
Comment on lines +1820 to +1877

@ggerganov ggerganov Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A bit difficult to digest, but I guess should be fine. Maybe at some point we can move the logic to llama_kv_cells and write unit+perf tests.

@ngxson ngxson Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah the logic here is not very idea: because kv cells can be out-of-order, we need to firstly write them to std::unordered_map, then for each token in a batch we find the list of N predecessor of that token

I think probably at some point it's better to track an ordered list of cells (in case of unified kv), can be a refactoring in the future


size_t llama_kv_cache::total_size() const {
size_t size = 0;

Expand Down Expand Up @@ -2106,7 +2177,7 @@ void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t
io.write(&pos, sizeof(pos));
io.write(&n_seq_id, sizeof(n_seq_id));

if (hparams.n_pos_per_embd() > 1) {
if (has_cell_ext()) {
const llama_kv_cell_ext ext = cells.ext_get(i);
io.write(&ext, sizeof(ext));
}
Expand Down Expand Up @@ -2243,12 +2314,17 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
return false;
}

if (hparams.n_pos_per_embd() > 1) {
if (has_cell_ext()) {
llama_kv_cell_ext ext;
io.read(&ext, sizeof(ext));

ubatch.pos[i + ubatch.n_tokens] = ext.y;
ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
if (hparams.n_pos_per_embd() > 1) {
ubatch.pos[i + ubatch.n_tokens] = ext.y;
ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
}

// apply_ubatch() below restores ext.tok from the ubatch tokens
ubatch.token[i] = ext.tok;
}

// read the sequence id, but directly discard it - we will use dest_seq_id instead
Expand All @@ -2268,7 +2344,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
return false;
}

// TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet
// note: apply_ubatch() rebuilds llama_kv_cell_ext from the ubatch
// only ext.tok and the M-RoPE 2D position round-trip through it
// see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
apply_ubatch(sinfo, ubatch);

Expand Down Expand Up @@ -2301,7 +2378,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32

cells.pos_set(i, pos);

if (hparams.n_pos_per_embd() > 1) {
if (has_cell_ext()) {
llama_kv_cell_ext ext;
io.read(&ext, sizeof(ext));
cells.ext_set(i, ext);
Expand Down Expand Up @@ -2652,3 +2729,7 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const {
kv->set_input_v_rot(dst);
}

void llama_kv_cache_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
kv->get_prev_tokens(ubatch, n, res);
}
11 changes: 11 additions & 0 deletions src/llama-kv-cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,14 @@ class llama_kv_cache : public llama_memory_i {
void set_input_k_rot(ggml_tensor * dst) const;
void set_input_v_rot(ggml_tensor * dst) const;

// true if llama_kv_cell_ext holds information that has to survive a state save/restore
bool has_cell_ext() const;

// for every token of the ubatch, the ids of the n tokens that precede it in its sequence
// entries with no matching cell are set to LLAMA_TOKEN_NULL
// note: used by n-gram input embeddings
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;

private:
const llama_model & model;
const llama_hparams & hparams;
Expand Down Expand Up @@ -401,6 +409,9 @@ class llama_kv_cache_context : public llama_memory_context_i {
void set_input_k_rot(ggml_tensor * dst) const;
void set_input_v_rot(ggml_tensor * dst) const;

// see llama_kv_cache::get_prev_tokens()
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;

private:
llama_memory_status status;

Expand Down
29 changes: 28 additions & 1 deletion src/llama-kv-cells.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@ struct llama_kv_cell_ext {
llama_pos x = 0;
llama_pos y = 0;

// when tok = LLAMA_TOKEN_NULL when the cell is produced by embedding input (i.e. multimodal)
// use case: n-gram embeddings hash
llama_token tok = LLAMA_TOKEN_NULL;

// return true if the current 2D spatial position is greater than other
bool is_2d_gt(llama_pos ox, llama_pos oy) const {
return (y > oy) || (y == oy && x > ox);
Expand All @@ -23,7 +27,7 @@ struct llama_kv_cell_ext {
void reset() {
static_assert(std::is_trivially_copyable_v<llama_kv_cell_ext>);

memset(this, 0, sizeof(*this));
*this = llama_kv_cell_ext{};
}
};

Expand Down Expand Up @@ -305,6 +309,29 @@ 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
// 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;
if (m.none()) {
continue;
}

for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
if (m.test(s)) {
f(s, pos[i], ext[i].tok);
}
}
}
}

// note: call only if the cell is not empty and the seq_id is not in the cell
void seq_add(uint32_t i, llama_seq_id seq_id) {
assert(i < pos.size());
Expand Down
Loading