From db997ff129bf5b982f2cb3eec204edcb5ca5de00 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 30 Aug 2026 18:51:13 +0200 Subject: [PATCH] kv-cells: look up the n-gram history in the sequence position index get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every ubatch by walking all used cells, while llama_kv_cells already keeps an ordered index of the positions of each sequence in seq_pos, updated on every cell mutation to serve seq_pos_min() and seq_pos_max(). The index now stores (pos, cell) pairs in a std::set instead of a position -> count map, so a repeated position (cache reuse via rm + add, vision inputs with shared positions) yields distinct entries and the removal of a cell erases its own pair. The new seq_pos_tok_le() returns the token of the cell at the largest position <= p in logarithmic time, which is exactly what the old window lookup and its M-RoPE gap fallback computed together. get_prev_tokens() shrinks to a direct lookup per (token, offset) and for_each_token_in() goes away with its only caller. The kv-cache keeps no n-gram logic of its own. Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%), pp unchanged at ~2720 t/s, greedy output identical, needle retrieved. --- src/llama-kv-cache.cpp | 59 +++--------------------------------- src/llama-kv-cells.h | 69 +++++++++++++++++------------------------- 2 files changed, 33 insertions(+), 95 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index fd7ce0bb6e4..3e4a4d56f92 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -6,7 +6,6 @@ #include "llama-context.h" #include -#include #include #include #include @@ -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::max(); - llama_pos p_max = std::numeric_limits::min(); - - std::bitset 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 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, 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 @@ -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); } } } diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index e9adffc0989..5d567a6ed0b 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -6,7 +6,7 @@ #include #include #include -#include +#include #include #include @@ -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; @@ -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; } @@ -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 - void for_each_token_in(const std::bitset & 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::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 @@ -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 @@ -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; } @@ -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; } @@ -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; - // 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 seq_pos[LLAMA_MAX_SEQ]; + std::set> 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); } } } @@ -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); } } }