llama: add token ID tracking to KV cell

@@ -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" {

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

@@ -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;@@ -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; 

@@ -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 @@ 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());