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

* 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
+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;