wylab/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-06 17:57:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama: add token ID tracking to KV cell (#27762)

Browse Source 
* kv: track token id

* rm get_prev_tokens, move it to the main pr

* nits

* add get_prev_tokens
This commit is contained in:
Xuan-Son Nguyen
2026-08-26 23:34:28 +02:00
committed by GitHub
parent 539f24529b
commit 925e117994
4 changed files with 133 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/llama.h
+2 -2
View File
@@ -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" {
src/llama-kv-cache.cpp
+92 -11
View File
@@ -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;
@@ -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);
}
@@ -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;
}
}
}
}
size_t llama_kv_cache::total_size() const {
size_t size = 0;
@@ -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));
}
@@ -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
@@ -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);
@@ -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);
@@ -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);
}
src/llama-kv-cache.h
+11
View File
@@ -219,6 +219,14 @@ public:
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;
@@ -401,6 +409,9 @@ public:
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;
src/llama-kv-cells.h
+28 -1
View File
@@ -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);
@@ -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{};
}
};
@@ -305,6 +309,29 @@ 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
// 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());