diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 65afbd8c377..5ede0ea054e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -8,7 +8,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -1826,6 +1828,88 @@ bool llama_kv_cache::has_cell_ext() const { return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0; } +// +// (seq, pos) -> token index for the n-gram predecessor lookups in get_prev_tokens() +// +// the llama_kv_cells::seq_pos index stores the rows of the cells per (seq, pos), which allows resolving the +// predecessor tokens in O(needs * log n) instead of scanning all used cells once per ubatch. controlled by the +// LLAMA_KV_PREV_TOKENS environment variable: +// +// fast (default): use the index for eligible ubatches, fall back to the general scan when not applicable +// verify: run the general scan (source of truth) and the index for every call, compare them and log +// mismatches (with a heartbeat of the scan / index cost). intended for testing and debugging +// off: general scan only (the index is still maintained, but never queried) +// + +enum class llama_prev_tokens_idx_mode { + VERIFY, + FAST, + OFF, +}; + +static llama_prev_tokens_idx_mode llama_prev_tokens_idx_mode_get() { + static const llama_prev_tokens_idx_mode mode = []() { + const char * env = getenv("LLAMA_KV_PREV_TOKENS"); + + if (env && strcmp(env, "verify") == 0) { return llama_prev_tokens_idx_mode::VERIFY; } + if (env && strcmp(env, "off") == 0) { return llama_prev_tokens_idx_mode::OFF; } + + return llama_prev_tokens_idx_mode::FAST; + }(); + + return mode; +} + +// diagnostic counters - decode is serialized per context, benign otherwise (reporting only): +static struct { + uint64_t n_calls = 0; + uint64_t n_needs = 0; + uint64_t n_mismatch = 0; + uint64_t n_not_applic = 0; // ubatches the index cannot serve (multimodal) + uint64_t n_fast_hits = 0; + double t_scan_us = 0.0; + double t_index_us = 0.0; +} g_prev_tokens_diag; + +bool llama_kv_cache::get_prev_tokens_indexed(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { + // an embd (multimodal) ubatch can repeat one position for a whole image, so positions do not encode the + // token order - the general scan resolves those predecessors by ubatch order instead + if (!ubatch.token) { + return false; + } + + res.clear(); + res.resize(ubatch.n_tokens*n, LLAMA_TOKEN_NULL); + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + // same choice as the general scan's consumer (the n-gram architectures reject shared tokens) + 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; + } + + llama_token tok = LLAMA_TOKEN_NULL; + + for (uint32_t s = 0; s < n_stream; ++s) { + llama_token t = LLAMA_TOKEN_NULL; + + // later stream wins, matching the general scan which merges all streams into one hist + if (v_cells[s].seq_pos_token_le(seq_id, p, &t)) { + tok = t; + } + } + + res[i*n + j] = tok; + } + } + + return true; +} + void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { const uint32_t n_tokens = ubatch.n_tokens; @@ -1836,6 +1920,45 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st return; } + const llama_prev_tokens_idx_mode idx_mode = llama_prev_tokens_idx_mode_get(); + const auto t_start = std::chrono::steady_clock::now(); + + if (idx_mode == llama_prev_tokens_idx_mode::FAST) { + const auto t_i0 = std::chrono::steady_clock::now(); + + if (get_prev_tokens_indexed(ubatch, n, res)) { + const auto t_i1 = std::chrono::steady_clock::now(); + + g_prev_tokens_diag.n_calls++; + g_prev_tokens_diag.n_needs += (uint64_t) n_tokens*n; + g_prev_tokens_diag.n_fast_hits++; + g_prev_tokens_diag.t_index_us += std::chrono::duration(t_i1 - t_i0).count(); + + if (g_prev_tokens_diag.n_calls % 2000 == 0) { + LLAMA_LOG_WARN("%s: PLE-IDX mode=fast calls=%llu needs=%llu hits=%llu not_applicable=%llu mismatches=%llu idx_avg=%.3f us used=%u\n", + __func__, + (unsigned long long) g_prev_tokens_diag.n_calls, + (unsigned long long) g_prev_tokens_diag.n_needs, + (unsigned long long) g_prev_tokens_diag.n_fast_hits, + (unsigned long long) g_prev_tokens_diag.n_not_applic, + (unsigned long long) g_prev_tokens_diag.n_mismatch, + g_prev_tokens_diag.t_index_us / g_prev_tokens_diag.n_calls, + v_cells[0].get_used()); + } + + return; + } + + // the index is not applicable to this ubatch - run the general scan below + g_prev_tokens_diag.n_not_applic++; + + res.clear(); + res.resize(n_tokens*n, LLAMA_TOKEN_NULL); + + LLAMA_LOG_WARN("%s: PLE-IDX mode=fast fell back to general scan (n_tokens=%u multimodal=%d)\n", + __func__, n_tokens, !ubatch.token); + } + // 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(); @@ -1928,6 +2051,64 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st res[i*n + j] = lookup(seq_id, p); } } + + if (idx_mode == llama_prev_tokens_idx_mode::OFF) { + return; + } + + // verify the index against the general scan, which remains the source of truth (no speedup intended here): + // any MISMATCH line means the index and the scan disagree and must not be trusted yet + { + const auto t_g1 = std::chrono::steady_clock::now(); + + std::vector res_idx; + + const auto t_i0 = std::chrono::steady_clock::now(); + const bool applicable = get_prev_tokens_indexed(ubatch, n, res_idx); + const auto t_i1 = std::chrono::steady_clock::now(); + + g_prev_tokens_diag.n_calls++; + g_prev_tokens_diag.t_scan_us += std::chrono::duration(t_g1 - t_start).count(); + g_prev_tokens_diag.t_index_us += std::chrono::duration(t_i1 - t_i0).count(); + + if (!applicable) { + g_prev_tokens_diag.n_not_applic++; + } else { + for (uint32_t i = 0; i < n_tokens; ++i) { + for (uint32_t j = 0; j < n; ++j) { + const llama_token a = res[i*n + j]; + const llama_token b = res_idx[i*n + j]; + + if (a == b) { + continue; + } + + if (g_prev_tokens_diag.n_mismatch < 32) { + LLAMA_LOG_ERROR("%s: PLE-IDX MISMATCH calls=%llu i=%u j=%u seq=%d pos[i]=%d scan=%d index=%d\n", + __func__, + (unsigned long long) g_prev_tokens_diag.n_calls, + i, j, ubatch.seq_id[i][0], ubatch.pos[i], a, b); + } + + g_prev_tokens_diag.n_mismatch++; + } + } + + g_prev_tokens_diag.n_needs += (uint64_t) n_tokens*n; + } + + if (g_prev_tokens_diag.n_calls % 2000 == 0) { + LLAMA_LOG_WARN("%s: PLE-IDX mode=verify calls=%llu needs=%llu mismatches=%llu not_applicable=%llu scan_avg=%.1f us index_avg=%.3f us used=%u\n", + __func__, + (unsigned long long) g_prev_tokens_diag.n_calls, + (unsigned long long) g_prev_tokens_diag.n_needs, + (unsigned long long) g_prev_tokens_diag.n_mismatch, + (unsigned long long) g_prev_tokens_diag.n_not_applic, + g_prev_tokens_diag.t_scan_us / g_prev_tokens_diag.n_calls, + g_prev_tokens_diag.t_index_us / g_prev_tokens_diag.n_calls, + v_cells[0].get_used()); + } + } } size_t llama_kv_cache::total_size() const { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c4d8699def1..eab552de383 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -243,6 +243,11 @@ class llama_kv_cache : public llama_memory_i { // note: used by n-gram input embeddings void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + // (seq, pos) -> token index variant of get_prev_tokens - O(needs * log n) instead of scanning all cells + // returns false if the ubatch is not eligible (multimodal predecessors are resolved by ubatch order) + // see llama_kv_cells::seq_pos_token_le() and the LLAMA_KV_PREV_TOKENS modes in llama-kv-cache.cpp + bool get_prev_tokens_indexed(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + private: const llama_model & model; const llama_hparams & hparams; diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index a4292c79e44..78da75953d5 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -246,7 +246,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, pos[i], i); if (seq[i].none()) { pos[i] = -1; @@ -270,7 +270,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, pos[i], i); return false; } @@ -340,7 +340,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, pos[i], i); } // return the sequence id of this cell @@ -367,7 +367,7 @@ class llama_kv_cells { return -1; } - assert(seq_pos[seq_id].begin()->second > 0); + assert(!seq_pos[seq_id].begin()->second.empty()); return seq_pos[seq_id].begin()->first; } @@ -382,11 +382,36 @@ class llama_kv_cells { return -1; } - assert(seq_pos[seq_id].rbegin()->second > 0); + assert(!seq_pos[seq_id].rbegin()->second.empty()); return seq_pos[seq_id].rbegin()->first; } + // the token of the newest cell (largest row) at the largest position <= p for the given sequence + // this reproduces the result of an ascending scan over all cells building a (seq, pos) -> token map, where + // the last cell written for a (seq, pos) pair wins - see llama_kv_cache::get_prev_tokens() + // return values: + // +1: found, *tok is set + // 0: the sequence has no cell at or before p + int seq_pos_token_le(llama_seq_id seq_id, llama_pos p, llama_token * tok) const { + assert(seq_id >= 0); + assert(seq_id < LLAMA_MAX_SEQ); + + const auto & m = seq_pos[seq_id]; + + auto it = m.upper_bound(p); + if (it == m.begin()) { + return 0; + } + + --it; + assert(!it->second.empty()); + + *tok = ext[*it->second.rbegin()].tok; + + return 1; + } + // note: call only if the cell is not empty llama_pos pos_get(uint32_t i) const { assert(i < pos.size()); @@ -516,7 +541,7 @@ 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 + // the set seq_pos[s][p] contains the rows (cell indices) of all cells currently occupying position p for sequence s // if the position p is not present, seq_pos[s][p] is not set // this way seq_pos[s].begin() and seq_pos[s].rbegin() give us the min/max positions currently in the cache // @@ -524,28 +549,34 @@ class llama_kv_cells { // - during performing a cache reuse via (rm + add) // - some vision models have input embeddings with repeating positions // - std::map seq_pos[LLAMA_MAX_SEQ]; + // keeping the rows (and not just a reference count) allows resolving the token of a (seq, pos) pair without + // scanning all cells. the rows are sorted, so *rbegin() reproduces the "last cell wins" result of an ascending + // scan over all cells - see llama_kv_cache::get_prev_tokens() and llama_kv_cells::seq_pos_token_le() + // + std::map> 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) { + void seq_pos_dec(llama_seq_id s, llama_pos p, uint32_t row) { auto it = seq_pos[s].find(p); assert(it != seq_pos[s].end()); - if (--it->second == 0) { + it->second.erase(row); + + if (it->second.empty()) { seq_pos[s].erase(it); } } - void seq_pos_inc(llama_seq_id s, llama_pos p) { - seq_pos[s][p]++; + void seq_pos_inc(llama_seq_id s, llama_pos p, uint32_t row) { + seq_pos[s][p].insert(row); } // 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, pos[i], i); } } } @@ -554,7 +585,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, pos[i], i); } } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fe3d14ffc55..8dd844e7383 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -289,6 +289,7 @@ llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p " set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) llama_build_and_test(test-arg-parser.cpp) +llama_build_and_test(test-kv-prev-tokens.cpp) llama_build_and_test(test-model-resolution.cpp) # the test serves its repos from an httplib server, and the library links it privately target_link_libraries(test-model-resolution PRIVATE cpp-httplib) diff --git a/tests/test-kv-prev-tokens.cpp b/tests/test-kv-prev-tokens.cpp new file mode 100644 index 00000000000..ad15e91d154 --- /dev/null +++ b/tests/test-kv-prev-tokens.cpp @@ -0,0 +1,210 @@ +// unit test: llama_kv_cells::seq_pos_token_le vs brute-force ascending scan +// +// the reference implements exactly what llama_kv_cache::get_prev_tokens() does: +// scan all cells ascending (for_each_token_in visits `used` in ascending row order), build +// hist[(seq, pos)] = token with last-write-wins (= largest row), then resolve each need by +// walking q = p..0 and taking the first hit. compares against the O(log n) index lookup. +// +// (llama_kv_cells is header-inline; built as a standalone unit test via tests/CMakeLists.txt) + +#include "llama-kv-cells.h" + +#include +#include +#include +#include +#include +#include + +static int n_fail = 0; +static int n_checked = 0; + +using hist_t = std::map, llama_token>; + +static hist_t build_hist(const llama_kv_cells & cells, uint32_t p_max) { + hist_t hist; + std::bitset all; + all.set(); + + cells.for_each_token_in(all, 0, p_max + 1, [&](llama_seq_id s, llama_pos p, llama_token t) { + hist[{s, p}] = t; // ascending rows: last write (largest row) wins + }); + + return hist; +} + +// reference answer for (seq, p): token at the largest position <= p, per the hist +static llama_token ref_lookup(const hist_t & hist, llama_seq_id s, llama_pos p) { + for (llama_pos q = p; q >= 0; --q) { + auto it = hist.find({ s, q }); + if (it != hist.end()) { + return it->second; + } + } + return LLAMA_TOKEN_NULL; +} + +static void check(const llama_kv_cells & cells, uint32_t p_max, const char * stage) { + const hist_t hist = build_hist(cells, p_max); + + for (llama_seq_id s = 0; s < 4; ++s) { + for (llama_pos p = -2; p <= (llama_pos) p_max + 2; ++p) { + if (p < 0) { + continue; + } + + llama_token tok = LLAMA_TOKEN_NULL; + const int rc = cells.seq_pos_token_le(s, p, &tok); + + const llama_token idx = rc > 0 ? tok : LLAMA_TOKEN_NULL; + const llama_token ref = ref_lookup(hist, s, p); + + n_checked++; + + if (idx != ref) { + printf("FAIL [%s] seq=%d p=%d: index=%d reference=%d\n", stage, s, p, idx, ref); + n_fail++; + } + } + } +} + +static void add_cell(llama_kv_cells & cells, uint32_t row, llama_pos p, llama_seq_id s, llama_token t) { + cells.pos_set(row, p); + llama_kv_cell_ext e; + e.tok = t; + cells.ext_set(row, e); + cells.seq_add(row, s); +} + +int main() { + // scenario 1: sequential append + { + llama_kv_cells cells; + cells.resize(64); + for (uint32_t i = 0; i < 20; ++i) { + add_cell(cells, i, (llama_pos) i, 0, (llama_token) (100 + i)); + } + check(cells, 19, "sequential"); + } + + // scenario 2: holes (removed middle cells) -> nearest <= fallback + { + llama_kv_cells cells; + cells.resize(64); + for (uint32_t i = 0; i < 20; ++i) { + add_cell(cells, i, (llama_pos) i, 0, (llama_token) (100 + i)); + } + // remove positions 7..12 via seq_rm + cleanup + for (uint32_t i = 7; i <= 12; ++i) { + if (cells.seq_rm(i, 0)) { /* becomes empty */ } + } + check(cells, 19, "holes"); + } + + // scenario 3: duplicate positions (checkpoint-style copies), incl. removal of copies + { + llama_kv_cells cells; + cells.resize(64); + for (uint32_t i = 0; i < 10; ++i) { + add_cell(cells, i, (llama_pos) i, 0, (llama_token) (100 + i)); + } + // "checkpoint" copy of rows 0..9 into rows 32..41, same positions, same seq, different tokens + for (uint32_t i = 0; i < 10; ++i) { + add_cell(cells, 32 + i, (llama_pos) i, 0, (llama_token) (900 + i)); + } + check(cells, 15, "dup-positions"); + + // erase the originals (rows 0..9) -> copies must still resolve + for (uint32_t i = 0; i < 10; ++i) { + cells.seq_rm(i, 0); + } + check(cells, 15, "dup-orig-removed"); + + // erase the copies too + for (uint32_t i = 32; i < 42; ++i) { + cells.seq_rm(i, 0); + } + check(cells, 15, "dup-copies-removed"); + } + + // scenario 4: restore via set() (prompt-cache load), rows out of position order + { + llama_kv_cells cells; + cells.resize(64); + add_cell(cells, 0, 0, 0, 100); + add_cell(cells, 1, 1, 0, 101); + + // simulate memory_seq_cp: write a state block at a LOW row range with HIGH positions + llama_kv_cells src; + src.resize(8); + for (uint32_t j = 0; j < 8; ++j) { + add_cell(src, j, (llama_pos) (100 + j), 1, (llama_token) (500 + j)); + } + cells.set(2, src); // rows 2..9 now hold positions 100..107 for seq 1 + + // decode continues at 108 on a high row + add_cell(cells, 10, 108, 1, 600); + check(cells, 110, "restore-low-rows"); + } + + // scenario 5: multi-seq shared cells + { + llama_kv_cells cells; + cells.resize(32); + for (uint32_t i = 0; i < 6; ++i) { + add_cell(cells, i, (llama_pos) i, 0, (llama_token) (100 + i)); + cells.seq_add(i, 1); + } + for (uint32_t i = 6; i < 10; ++i) { + add_cell(cells, i, (llama_pos) i, 1, (llama_token) (200 + i)); + } + check(cells, 12, "multi-seq"); + } + + // scenario 6: randomized fuzz - random adds/removes/multi-seq vs reference + // the mirror tracks (row -> seq set) so that seq_rm/seq_add are called only when valid + { + std::mt19937 rng(42); + llama_kv_cells cells; + cells.resize(256); + + std::map> mirror; + + for (int step = 0; step < 3000; ++step) { + const uint32_t row = rng() % 256; + + auto it = mirror.find(row); + if (it == mirror.end()) { + const llama_seq_id s = rng() % 3; + add_cell(cells, row, (llama_pos) (rng() % 60), s, (llama_token) (1000 + rng() % 1000)); + mirror[row].insert(s); + } else if (it->second.size() < 2 && rng() % 5 == 0) { + const llama_seq_id s = rng() % 3; + if (!it->second.count(s)) { + cells.seq_add(row, s); + it->second.insert(s); + } + } else { + const llama_seq_id s = *it->second.begin(); + const bool empty = cells.seq_rm(row, s); + it->second.erase(s); + if (empty != it->second.empty()) { + printf("FAIL [fuzz-consistency] row=%u seq_rm returned %d, mirror left %zu\n", row, (int) empty, it->second.size()); + n_fail++; + } + if (it->second.empty()) { + mirror.erase(it); + } + } + + if (step % 97 == 0) { + check(cells, 64, "fuzz"); + } + } + check(cells, 64, "fuzz-final"); + } + + printf("%s: %d lookups checked, %d failures\n", n_fail ? "FAILED" : "PASSED", n_checked, n_fail); + return n_fail ? 1 : 0; +}