diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index ef76fff1..95cf877e 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -102,6 +102,8 @@ add_library(${TARGET} reasoning-budget.h sampling.cpp sampling.h + scr.cpp + scr.h speculative.cpp speculative.h subproc.cpp diff --git a/common/arg.cpp b/common/arg.cpp index 2487497d..e096bfad 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1760,6 +1760,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_classifier_timeout = std::max(value, 1); } ).set_env("LLAMA_ARG_CACHE_CLASSIFIER_TIMEOUT").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--scr"}, + {"--no-scr"}, + "suffix cache reuse for requests that do not set \"scr\" themselves (default: off)", + [](common_params & params, bool value) { + params.scr = value; + } + ).set_env("LLAMA_ARG_SCR").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index 76028932..550bef61 100644 --- a/common/common.h +++ b/common/common.h @@ -626,6 +626,7 @@ struct common_params { int32_t cache_ttl_keep = 30*86400; std::string cache_classifier; // URL that classifies new prompt cache entries, empty = off int32_t cache_classifier_timeout = 20; // seconds + bool scr = false; // suffix cache reuse for requests that do not set "scr" std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/common/scr.cpp b/common/scr.cpp new file mode 100644 index 00000000..23fa285e --- /dev/null +++ b/common/scr.cpp @@ -0,0 +1,201 @@ +#include "scr.h" + +#include "../src/llama-ext.h" + +#include +#include + +int32_t common_scr_plan::resume() const { + return runs.empty() ? restore : runs.back().new_start + runs.back().n; +} + +int32_t common_scr_plan::n_reused() const { + int32_t n = 0; + for (const auto & r : runs) { + n += r.n; + } + return n; +} + +std::string common_scr_plan::str() const { + std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + std::to_string(adopt); + for (const auto & r : runs) { + s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n); + } + return s; +} + +// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once +static std::vector scr_find_spans( + const std::vector & old, + const std::vector & cur, + int32_t from, + const common_scr_params & params) { + std::vector spans; + + const int32_t n_old = (int32_t) old.size(); + const int32_t n_cur = (int32_t) cur.size(); + const int32_t k = params.seed; + + if (n_old - from < params.min_span || n_cur - from < params.min_span) { + return spans; + } + + auto gram = [k](const std::vector & v, int32_t i) { + uint64_t h = 1469598103934665603ULL; + for (int32_t j = 0; j < k; ++j) { + h = (h ^ (uint32_t) v[i + j]) * 1099511628211ULL; + } + return h; + }; + + // the first few old positions of every seed n-gram + const size_t max_bucket = 16; + std::unordered_map> index; + index.reserve(n_old - from); + for (int32_t o = from; o + k <= n_old; ++o) { + auto & b = index[gram(old, o)]; + if (b.size() < max_bucket) { + b.push_back(o); + } + } + + std::vector used(n_old, false); + + for (int32_t p = from; p + k <= n_cur; ) { + const auto it = index.find(gram(cur, p)); + + int32_t best_o = -1; + int32_t best_n = 0; + if (it != index.end()) { + for (const int32_t o : it->second) { + int32_t n = 0; + while (p + n < n_cur && o + n < n_old && !used[o + n] && + cur[p + n] == old[o + n] && cur[p + n] != LLAMA_TOKEN_NULL) { + ++n; + } + if (n > best_n) { + best_n = n; + best_o = o; + } + } + } + + if (best_n >= params.min_span) { + spans.push_back({ p, best_o, best_n }); + std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true); + p += best_n; + } else { + ++p; + } + } + + return spans; +} + +common_scr_plan common_scr_plan_build( + const std::vector & old, + const std::vector & cur, + const std::vector & snapshots, + bool recurrent, + const common_scr_params & params) { + common_scr_plan plan; + + const int32_t n_cur = (int32_t) cur.size(); + + int32_t a = 0; + while (a < (int32_t) old.size() && a < n_cur && old[a] == cur[a] && cur[a] != LLAMA_TOKEN_NULL) { + ++a; + } + plan.n_prefix = a; + plan.restore = a; + + if (recurrent) { + plan.restore = 0; + for (const int32_t q : snapshots) { + if (q <= a) { + plan.restore = std::max(plan.restore, q); + } + } + } + + auto spans = scr_find_spans(old, cur, a, params); + + // the largest K spans, back in prompt order + std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.n > y.n; }); + if ((int32_t) spans.size() > params.max_runs) { + spans.resize(params.max_runs); + } + std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.new_start < y.new_start; }); + + // the last span must end at a snapshot, which also ends the reuse: everything after it is prefilled from that state + while (!spans.empty()) { + auto & last = spans.back(); + + int32_t end = -1; + if (recurrent) { + for (const int32_t q : snapshots) { + const int32_t n = q - last.old_start; + if (n >= params.min_span && n <= last.n && last.new_start + n < n_cur) { + end = std::max(end, q); + } + } + } else { + end = last.old_start + std::min(last.n, n_cur - 1 - last.new_start); + } + + if (end - last.old_start >= params.min_span) { + last.n = end - last.old_start; + plan.adopt = recurrent ? end : -1; + break; + } + + spans.pop_back(); + } + + plan.runs = std::move(spans); + + return plan; +} + +int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan) { + // a run may move once its target holds no source of a run still to move + std::vector order; + { + std::vector pending = plan.runs; + while (!pending.empty()) { + auto next = std::find_if(pending.begin(), pending.end(), [&](const common_scr_run & r) { + return std::none_of(pending.begin(), pending.end(), [&](const common_scr_run & o) { + return &o != &r && o.old_start < r.new_start + r.n && r.new_start < o.old_start + o.n; + }); + }); + if (next == pending.end()) { + return 0; + } + order.push_back(*next); + pending.erase(next); + } + } + + std::vector sources = plan.runs; + std::sort(sources.begin(), sources.end(), [](const auto & x, const auto & y) { return x.old_start < y.old_start; }); + + llama_pos pos = plan.restore; + for (const auto & s : sources) { + if (s.old_start > pos && !llama_scr_kv_rm(ctx, seq_id, pos, s.old_start)) { + return -1; + } + pos = std::max(pos, (llama_pos) (s.old_start + s.n)); + } + if (!llama_scr_kv_rm(ctx, seq_id, pos, -1)) { + return -1; + } + + for (const auto & r : order) { + if (!llama_scr_kv_relocate(ctx, seq_id, r.old_start, r.old_start + r.n, r.new_start - r.old_start)) { + return -1; + } + } + + return 1; +} diff --git a/common/scr.h b/common/scr.h new file mode 100644 index 00000000..03a74367 --- /dev/null +++ b/common/scr.h @@ -0,0 +1,59 @@ +#pragma once + +#include "llama.h" + +#include +#include +#include + +// +// suffix cache reuse (SCR) +// +// after an edit in the middle of a prompt, the cached attention KV of the text that survived is moved to its new +// positions (K re-rotated), the recurrent state of a hybrid model continues from a snapshot taken inside the last +// moved span, and only the rest is prefilled. the moved rows were computed over the pre-edit text, so the result +// approximates a fresh prefill of the edited prompt. +// + +struct common_scr_params { + int32_t min_span = 256; // shortest span worth moving + int32_t max_runs = 6; // most spans moved per request + int32_t seed = 32; // n-gram length that seeds a match +}; + +// a span of the cached tokens [old_start, old_start + n) that reappears at [new_start, new_start + n) +struct common_scr_run { + int32_t new_start; + int32_t old_start; + int32_t n; +}; + +struct common_scr_plan { + int32_t n_prefix = 0; // leading tokens equal in both prompts + int32_t restore = 0; // position the prefill restarts from: a snapshot at or before n_prefix + int32_t adopt = -1; // old position of the snapshot the recurrent state continues from, -1 without one + + std::vector runs; // in new order, all at or after n_prefix + + bool empty() const { return runs.empty(); } + int32_t resume() const; // new position the final prefill starts at + int32_t n_reused() const; + + std::string str() const; +}; + +// plans the reuse of a state that holds old[0, n_old) for the prompt cur. +// snapshots: sorted old positions q at which the recurrent state after old[0, q) can be restored (the live state +// counts as old.size()); with recurrent = false every position can be resumed from. +// LLAMA_TOKEN_NULL never matches. the plan always leaves at least one token of cur to prefill. +common_scr_plan common_scr_plan_build( + const std::vector & old, + const std::vector & cur, + const std::vector & snapshots, + bool recurrent, + const common_scr_params & params); + +// rearranges the attention KV of seq_id for a plan: drops every cell at or after plan.restore that no run moves, +// then moves the runs. returns 1 when done, 0 when the runs cannot be moved in any order (nothing changed), +// -1 when a move failed after the cache was changed (the sequence must then be cleared). +int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan); diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552..73c8d9e3 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -262,7 +262,7 @@ bool llama_batch_allocr::init( const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; if (batch.token) { - if (p0 >= 0 && p0 >= seq_pos_min(s)) { + if (p0 >= 0 && p0 >= seq_pos_min(s) && !seq_spliced[s]) { LLAMA_LOG_ERROR( "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n" " - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n" @@ -274,7 +274,7 @@ bool llama_batch_allocr::init( } } else { // embedding inputs can have overlapping positions - if (p0 >= 0 && p0 > seq_pos_min(s)) { + if (p0 >= 0 && p0 > seq_pos_min(s) && !seq_spliced[s]) { LLAMA_LOG_ERROR( "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n" " - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n" diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d..add9d90c 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -99,6 +99,10 @@ public: // call once before splitting the batch to reset the internal state void split_reset(); + // sequences being spliced (suffix cache reuse): their text batches may start below positions the memory + // already holds, because the splice prefills text in front of spans it moved there + std::bitset seq_spliced; + // simple split, unknown number of sequence sets of unequal lengths llama_ubatch split_simple(uint32_t n_ubatch); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 1d186472..453d2ee3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4082,6 +4082,68 @@ float * llama_get_embeddings_seq(llama_context * ctx, llama_seq_id seq_id) { return ctx->get_embeddings_seq(seq_id); } +// the attention part of a context's memory, or nullptr +static llama_kv_cache * scr_kv(llama_memory_i * mem) { + if (auto * hybrid = dynamic_cast(mem)) { + return hybrid->get_mem_attn(); + } + return dynamic_cast(mem); +} + +void llama_context::set_seq_spliced(llama_seq_id seq_id, bool value) { + GGML_ASSERT(seq_id >= 0 && seq_id < LLAMA_MAX_SEQ); + balloc->seq_spliced[seq_id] = value; +} + +bool llama_context::memory_relocate(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { + auto * kv = scr_kv(memory.get()); + if (kv == nullptr) { + return false; + } + + synchronize(); + + // the relocation graph resets the scheduler, so the previous graph cannot be reused + gf_res_prev->reset(); + + return kv->relocate(this, seq_id, p0, p1, delta); +} + +bool llama_scr_kv_relocate(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { + return ctx->memory_relocate(seq_id, p0, p1, delta); +} + +void llama_scr_set_spliced(llama_context * ctx, llama_seq_id seq_id, bool value) { + ctx->set_seq_spliced(seq_id, value); +} + +bool llama_scr_kv_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + auto * kv = scr_kv(ctx->get_memory()); + return kv != nullptr && kv->seq_rm(seq_id, p0, p1); +} + +bool llama_scr_recr_set_pos(llama_context * ctx, llama_seq_id seq_id, llama_pos pos) { + auto * hybrid = dynamic_cast(ctx->get_memory()); + if (hybrid == nullptr) { + return false; + } + + auto * recr = hybrid->get_mem_recr(); + + if (pos < 0) { + return recr->seq_rm(seq_id, -1, -1); + } + + const llama_pos cur = recr->seq_pos_max(seq_id); + if (cur < 0) { + return false; + } + + recr->seq_add(seq_id, cur, cur + 1, pos - cur); + + return recr->seq_pos_max(seq_id) == pos; +} + void llama_set_embeddings_nextn(llama_context * ctx, bool value, bool masked) { ctx->set_embeddings_nextn(value, masked); } diff --git a/src/llama-context.h b/src/llama-context.h index bced1881..c32a1b6f 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -76,6 +76,10 @@ struct llama_context { // return true if the memory was updated bool memory_update(bool optimize); + // suffix cache reuse: see llama_scr_kv_relocate and llama_scr_set_spliced + bool memory_relocate(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); + void set_seq_spliced(llama_seq_id seq_id, bool value); + enum llama_pooling_type pooling_type() const; float * get_logits(); diff --git a/src/llama-ext.h b/src/llama-ext.h index c6cfc35b..263961c8 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -91,6 +91,25 @@ LLAMA_API ggml_backend_dev_t llama_model_get_device(const struct llama_model * m LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx); +// +// suffix cache reuse (SCR): splice primitives for one sequence +// + +// moves the attention KV of seq_id at positions [p0, p1) by delta and rotates K to the new positions. +// the cells must belong to seq_id alone and the target range must hold no other cell of seq_id. +// returns false when it cannot (nothing changed, or the moved cells dropped if the rotation failed midway) +LLAMA_API bool llama_scr_kv_relocate(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); + +// while set, text batches of seq_id may start below positions the memory already holds (M-RoPE models check +// otherwise that positions only grow): a splice prefills the edited text in front of the spans it moved +LLAMA_API void llama_scr_set_spliced(struct llama_context * ctx, llama_seq_id seq_id, bool value); + +// removes the attention KV of seq_id at positions [p0, p1) and leaves a recurrent state as it is +LLAMA_API bool llama_scr_kv_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1); + +// sets the last position absorbed by seq_id's recurrent state, or clears the state for pos < 0 (hybrid memory only) +LLAMA_API bool llama_scr_recr_set_pos(struct llama_context * ctx, llama_seq_id seq_id, llama_pos pos); + // Set whether the context outputs nextn embeddings or not // If masked == true, output the embeddings only for the tokens with batch.logits != 0 // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3df5cedd..d4c3c794 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2163,6 +2163,170 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co return gf; } +bool llama_kv_cache::relocate(llama_context * lctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { + // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] + if (other) { + return false; + } + + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); + + if (delta == 0 || p0 >= p1) { + return true; + } + + const auto & cparams = lctx->get_cparams(); + + // a second rotation also scales by the YaRN magnitude factor, and a bias is subtracted before the rotation + if (cparams.yarn_ext_factor > 0.0f || !k_bar.empty() || p0 + delta < 0 || model.arch == LLM_ARCH_STEP35) { + return false; + } + + const uint32_t strm = seq_to_stream[seq_id]; + auto & cells = v_cells[strm]; + + std::vector idxs; + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i) || !cells.seq_has(i, seq_id)) { + continue; + } + if (cells.pos_in(i, p0, p1)) { + if (cells.seq_count(i) != 1) { + return false; + } + idxs.push_back(i); + } else if (cells.pos_in(i, p0 + delta, p1 + delta)) { + return false; + } + } + + if (idxs.empty()) { + return true; + } + + const bool mrope = hparams.n_pos_per_embd() > 1; + + if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) { + auto * sched = lctx->get_sched(); + + // 2048 cells of a 1024-wide K row are 8 MiB per F32 intermediate + const uint32_t n_chunk = 2048; + const size_t n_nodes = 16 + 16*layers.size(); + const uint32_t n_pos = mrope ? 4 : 1; + + std::vector idx32; + std::vector idx64; + std::vector pos; + + for (size_t off = 0; off < idxs.size(); off += n_chunk) { + const uint32_t n = std::min(n_chunk, idxs.size() - off); + + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead()*n_nodes + ggml_graph_overhead_custom(n_nodes, false), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx_ptr { ggml_init(params) }; + ggml_context * ctx = ctx_ptr.get(); + + ggml_cgraph * gf = ggml_new_graph_custom(ctx, n_nodes, false); + + ggml_tensor * t_idx32 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n); + ggml_tensor * t_idx64 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n); + ggml_tensor * t_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, (int64_t) n*n_pos); + ggml_set_input(t_idx32); + ggml_set_input(t_idx64); + ggml_set_input(t_pos); + + ggml_tensor * t_rot = build_input_k_rot(ctx); + + for (const auto & layer : layers) { + const uint32_t il = layer.il; + + if (!hparams.has_rope(il)) { + continue; + } + + const int64_t n_head_kv = hparams.n_head_kv(il); + const int64_t n_embd_head_k = hparams.n_embd_head_k(il); + const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); + + ggml_tensor * k = layer.k_stream[strm]; + + ggml_tensor * cur = ggml_get_rows(ctx, k, t_idx32); + cur = ggml_reshape_3d(ctx, cur, n_embd_head_k, n_head_kv, n); + + if (t_rot) { + cur = llama_mul_mat_hadamard(ctx, cur, t_rot); + } + + const float freq_base = model.get_rope_freq_base (cparams, il); + const float freq_scale = model.get_rope_freq_scale(cparams, il); + + if (mrope) { + int sections[4]; + std::copy(hparams.rope_sections.begin(), hparams.rope_sections.end(), sections); + + cur = ggml_rope_multi(ctx, cur, t_pos, nullptr, + hparams.n_rot(il), sections, hparams.rope_type, cparams.n_ctx_orig_yarn, freq_base, freq_scale, + 0.0f, 1.0f, cparams.yarn_beta_fast, cparams.yarn_beta_slow); + } else { + cur = ggml_rope_ext(ctx, cur, t_pos, model.get_rope_factors(cparams, il), + hparams.n_rot(il), hparams.rope_type, cparams.n_ctx_orig_yarn, freq_base, freq_scale, + 0.0f, 1.0f, cparams.yarn_beta_fast, cparams.yarn_beta_slow); + } + + if (t_rot) { + cur = llama_mul_mat_hadamard(ctx, cur, t_rot); + } + + cur = ggml_reshape_2d(ctx, cur, n_embd_k_gqa, n); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx, k, cur, t_idx64)); + } + + ggml_backend_sched_reset(sched); + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + LLAMA_LOG_ERROR("%s: failed to allocate the relocation graph\n", __func__); + return false; + } + + idx32.resize(n); + idx64.resize(n); + pos.assign((size_t) n*n_pos, delta); + for (uint32_t j = 0; j < n; ++j) { + idx32[j] = idxs[off + j]; + idx64[j] = idxs[off + j]; + } + + ggml_backend_tensor_set(t_idx32, idx32.data(), 0, ggml_nbytes(t_idx32)); + ggml_backend_tensor_set(t_idx64, idx64.data(), 0, ggml_nbytes(t_idx64)); + ggml_backend_tensor_set(t_pos, pos.data(), 0, ggml_nbytes(t_pos)); + if (t_rot) { + ggml_backend_tensor_set(t_rot, attn_rot_hadamard.at(t_rot->ne[0]).data(), 0, ggml_nbytes(t_rot)); + } + + if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) { + // rows already written by earlier chunks are rotated, so the cells are lost + LLAMA_LOG_ERROR("%s: failed to compute the relocation graph\n", __func__); + for (const uint32_t i : idxs) { + cells.rm(i); + } + return false; + } + } + + ggml_backend_sched_synchronize(sched); + ggml_backend_sched_reset(sched); + } + + for (const uint32_t i : idxs) { + cells.pos_move(i, delta, mrope); + } + + return true; +} + void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 4d2229b9..5ff230ec 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -207,6 +207,12 @@ public: bool update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info); + // moves the cells of seq_id at positions [p0, p1) by delta and rotates their K to the new positions + // (suffix cache reuse). the cells must hold seq_id only, and no other cell of seq_id may sit in the target + // range. returns false without changing anything when it cannot; if the rotation fails midway the moved + // cells are dropped. + bool relocate(llama_context * lctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); + // find a slot of kv cells that can hold the ubatch // if cont == true, then the slot must be continuous // return empty slot_info on failure diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index fddd31a0..5326e181 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -436,6 +436,25 @@ public: return false; } + // pos[i] = pos[i] + d, and the 2D position with it when move_ext is set (M-RoPE) + // does not modify "has_shift": the caller has already rotated the cell's K + // note: call only if the cell is not empty and pos[i] + d >= 0 + void pos_move(uint32_t i, llama_pos d, bool move_ext) { + assert(i < pos.size()); + assert(pos[i] != -1); + assert(pos[i] + d >= 0); + + seq_pos_rm(i); + + if (move_ext) { + ext[i].x += d; + ext[i].y += d; + } + pos[i] += d; + + seq_pos_add(i); + } + // pos[i] = pos[i] / d // sets "has_shift" to true // note: call only if the cell is not empty diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index da1b2bc7..47455f49 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -18,6 +18,7 @@ else() add_subdirectory(gguf-split) add_subdirectory(imatrix) add_subdirectory(kv-mean-center) + add_subdirectory(scr-eval) add_subdirectory(llama-bench) add_subdirectory(completion) add_subdirectory(perplexity) diff --git a/tools/scr-eval/CMakeLists.txt b/tools/scr-eval/CMakeLists.txt new file mode 100644 index 00000000..5fddd24f --- /dev/null +++ b/tools/scr-eval/CMakeLists.txt @@ -0,0 +1,4 @@ +set(TARGET llama-scr-eval) +add_executable(${TARGET} scr-eval.cpp) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/tools/scr-eval/scr-eval.cpp b/tools/scr-eval/scr-eval.cpp new file mode 100644 index 00000000..f2e0c30b --- /dev/null +++ b/tools/scr-eval/scr-eval.cpp @@ -0,0 +1,498 @@ +// scr-eval: quality of suffix cache reuse (SCR) against a fresh prefill of the edited prompt +// +// reads cases (one JSON object per line, file given with -f): +// {"id": "...", "old": [tokens], "new": [tokens], "snapshots": [old positions], "answer": "text" (optional)} +// "old" is what the cached state holds; "snapshots" are the old positions where the server would hold a +// recurrent-state checkpoint (the end of "old" always counts). for every case it prepares the state for +// new[0, N-1) in several ways (arms), then scores the same 256 teacher tokens on each and compares the +// distributions with the fresh arm F: +// F fresh prefill of the edited prompt, its greedy continuation is the teacher +// floor fresh prefill in 512-token chunks (numerics floor) +// S SCR: moved attention KV, recurrent state from the pre-edit snapshot at the end of the last moved span +// S0 SCR without the snapshot: the recurrent state continues from the prefilled text before the span +// RT fresh state whose moved spans are relocated away and back (two K re-rotations, no staleness) +// prints one JSON line per case and arm. SCR_ARMS (comma list) selects arms. + +#include "arg.h" +#include "common.h" +#include "log.h" +#include "scr.h" +#include "llama.h" + +#include "../../src/llama-ext.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using json = nlohmann::ordered_json; + +static double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +struct evaluator { + llama_context * ctx; + const llama_vocab * vocab; + int32_t n_vocab; + + std::map> snaps; + + void clear() { + llama_memory_clear(llama_get_memory(ctx), true); + } + + std::vector save_seq() const { + const size_t n = llama_state_seq_get_size(ctx, 0); + std::vector buf(n); + GGML_ASSERT(llama_state_seq_get_data(ctx, buf.data(), n, 0) == n); + return buf; + } + + void load_seq(const std::vector & buf) { + clear(); + GGML_ASSERT(llama_state_seq_set_data(ctx, buf.data(), buf.size(), 0) == buf.size()); + } + + std::vector snap() const { + const size_t n = llama_state_seq_get_size_ext(ctx, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + std::vector buf(n); + GGML_ASSERT(llama_state_seq_get_data_ext(ctx, buf.data(), n, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == n); + return buf; + } + + void load(const std::vector & buf) const { + GGML_ASSERT(llama_state_seq_set_data_ext(ctx, buf.data(), buf.size(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == buf.size()); + } + + // decodes t[from, to) at its own positions in chunks, snapshotting the recurrent state at the stop positions + void prefill(const std::vector & t, int32_t from, int32_t to, int32_t n_chunk, const std::set & stops = {}) { + llama_batch batch = llama_batch_init(n_chunk, 0, 1); + for (int32_t p = from; p < to; ) { + int32_t end = std::min(to, p + n_chunk); + for (const int32_t s : stops) { + if (s > p && s < end) { + end = s; + } + } + common_batch_clear(batch); + for (int32_t i = p; i < end; ++i) { + common_batch_add(batch, t[i], i, { 0 }, false); + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + if (stops.count(end)) { + snaps[end] = snap(); + } + p = end; + } + llama_batch_free(batch); + } + + // the cache holds the old text and snaps its snapshots: rearranges it for new_t per the plan and prefills new_t up + // to n_end, snapshotting the recurrent state at stops met in prefilled text. afterwards snaps describes the new + // state. adopt = false is the S0 ablation. returns the number of tokens prefilled + int32_t splice(const common_scr_plan & plan, const std::vector & new_t, int32_t n_end, int32_t n_chunk, bool adopt, + const std::set & stops = {}) { + const std::vector state_adopt = snaps.at(plan.adopt); + + GGML_ASSERT(common_scr_apply_kv(ctx, 0, plan) == 1); + llama_scr_set_spliced(ctx, 0, true); + if (plan.restore > 0) { + load(snaps.at(plan.restore)); + } else { + GGML_ASSERT(llama_scr_recr_set_pos(ctx, 0, -1)); + } + for (auto it = snaps.begin(); it != snaps.end(); ) { + it = it->first > plan.restore ? snaps.erase(it) : std::next(it); + } + + int32_t n_prefilled = 0; + int32_t seg = plan.restore; + for (const auto & r : plan.runs) { + if (r.new_start > seg) { + if (seg > 0) { + GGML_ASSERT(llama_scr_recr_set_pos(ctx, 0, seg - 1)); + } + prefill(new_t, seg, r.new_start, n_chunk, stops); + n_prefilled += r.new_start - seg; + } + seg = r.new_start + r.n; + } + if (adopt) { + load(state_adopt); + } + GGML_ASSERT(llama_scr_recr_set_pos(ctx, 0, plan.resume() - 1)); + llama_scr_set_spliced(ctx, 0, false); + prefill(new_t, plan.resume(), n_end, n_chunk, stops); + + return n_prefilled + n_end - plan.resume(); + } + + int32_t argmax(const float * row) const { + return (int32_t) (std::max_element(row, row + n_vocab) - row); + } + + // greedy continuation of n tokens, starting by decoding `first` at pos + std::vector generate(llama_token first, int32_t pos, int32_t n) { + std::vector out; + llama_batch batch = llama_batch_init(1, 0, 1); + llama_token tok = first; + for (int32_t i = 0; i < n; ++i) { + common_batch_clear(batch); + common_batch_add(batch, tok, pos + i, { 0 }, true); + GGML_ASSERT(llama_decode(ctx, batch) == 0); + tok = argmax(llama_get_logits_ith(ctx, 0)); + out.push_back(tok); + } + llama_batch_free(batch); + return out; + } + + // log-softmax rows for [first, teacher[0, n-1)] decoded at pos in one batch + std::vector score(llama_token first, const std::vector & teacher, int32_t pos) { + const int32_t n = (int32_t) teacher.size(); + llama_batch batch = llama_batch_init(n, 0, 1); + common_batch_add(batch, first, pos, { 0 }, true); + for (int32_t i = 0; i + 1 < n; ++i) { + common_batch_add(batch, teacher[i], pos + 1 + i, { 0 }, true); + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + std::vector rows((size_t) n * n_vocab); + for (int32_t i = 0; i < n; ++i) { + const float * l = llama_get_logits_ith(ctx, i); + float * r = rows.data() + (size_t) i * n_vocab; + const float mx = *std::max_element(l, l + n_vocab); + double z = 0.0; + for (int32_t v = 0; v < n_vocab; ++v) { + z += std::exp((double) l[v] - mx); + } + const float lz = (float) (mx + std::log(z)); + for (int32_t v = 0; v < n_vocab; ++v) { + r[v] = l[v] - lz; + } + } + llama_batch_free(batch); + return rows; + } + + // the state for prompt[0, N-1) is ready: optional greedy text, then the scored rows; leaves the cache as it found it + std::vector finish(const std::vector & prompt, std::vector & teacher, int32_t n_teacher, int32_t n_gen, std::string * text) { + const int32_t N = (int32_t) prompt.size(); + const int32_t n = teacher.empty() ? n_teacher : n_gen; + if (n > 0) { + const auto s = snap(); + auto gen = generate(prompt[N - 1], N - 1, n); + load(s); + GGML_ASSERT(llama_scr_kv_rm(ctx, 0, N - 1, -1)); + if (text) { + *text = common_detokenize(ctx, gen, false); + } + if (teacher.empty()) { + teacher = gen; + } + } + return score(prompt[N - 1], teacher, N - 1); + } +}; + +struct arm_stats { + double kl_mean = 0, kl_p50 = 0, kl_p99 = 0, kl_max = 0; + double top1 = 0, top1_teacher = 0, nll = 0; + int32_t first_div = -1; +}; + +// the server's rule: at least 512 reused tokens and a tenth of what a plain request would prefill +static bool scr_pays_eval(const common_scr_plan & plan, int32_t n_tokens) { + return !plan.empty() && plan.n_reused() >= std::max(512, (n_tokens - plan.restore) / 10); +} + +static arm_stats compare(const std::vector & ref, const std::vector & arm, const std::vector & teacher, int32_t n_vocab) { + arm_stats st; + const int32_t n = (int32_t) teacher.size(); + std::vector kls; + for (int32_t i = 0; i < n; ++i) { + const float * p = ref.data() + (size_t) i * n_vocab; + const float * q = arm.data() + (size_t) i * n_vocab; + double kl = 0.0; + for (int32_t v = 0; v < n_vocab; ++v) { + kl += std::exp((double) p[v]) * ((double) p[v] - (double) q[v]); + } + kls.push_back(std::max(0.0, kl)); + const int32_t ap = (int32_t) (std::max_element(p, p + n_vocab) - p); + const int32_t aq = (int32_t) (std::max_element(q, q + n_vocab) - q); + st.top1 += ap == aq; + st.top1_teacher += aq == teacher[i]; + st.nll -= q[teacher[i]]; + if (st.first_div < 0 && aq != teacher[i]) { + st.first_div = i; + } + } + std::vector s = kls; + std::sort(s.begin(), s.end()); + for (const double k : kls) { + st.kl_mean += k; + } + st.kl_mean /= n; + st.kl_p50 = s[n / 2]; + st.kl_p99 = s[std::min(n - 1, (int32_t) (0.99 * n))]; + st.kl_max = s.back(); + st.top1 /= n; + st.top1_teacher /= n; + st.nll /= n; + return st; +} + +int main(int argc, char ** argv) { + common_params params; + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + common_init(); + llama_backend_init(); + llama_numa_init(params.numa); + + auto llama_init = common_init_from_params(params); + llama_context * ctx = llama_init->context(); + if (ctx == nullptr) { + LOG_ERR("failed to create the context\n"); + return 1; + } + + evaluator ev { ctx, llama_model_get_vocab(llama_get_model(ctx)), 0, {} }; + ev.n_vocab = llama_vocab_n_tokens(ev.vocab); + + std::set arms = { "F", "floor", "S", "S0", "RT" }; + if (const char * e = getenv("SCR_ARMS")) { + arms.clear(); + std::string s = e; + for (size_t p = 0; p <= s.size(); ) { + const size_t q = std::min(s.find(',', p), s.size()); + arms.insert(s.substr(p, q - p)); + p = q + 1; + } + } + const int32_t n_teacher = getenv("SCR_TEACHER") ? atoi(getenv("SCR_TEACHER")) : 256; + const int32_t n_chunk = (int32_t) llama_n_batch(ctx); + + common_scr_params sp; + + // the state after prefilling the last distinct "old", with its snapshots, shared by the cases that start from it + std::vector cached_old; + std::vector cached_state; + std::map> cached_snaps; + + std::ifstream in(params.prompt_file); + std::string line; + while (std::getline(in, line)) { + if (line.empty()) { + continue; + } + const json c = json::parse(line); + const std::string id = c.at("id"); + + // a lineage: each step edits the previous step's text (prompt + reply). L splices every step from the previous + // spliced state, O splices from a fresh state of the previous text, F prefills fresh; L and O are compared with + // F on the step's reply + if (c.contains("chain")) { + const auto & steps = c.at("chain"); + + std::vector prev_t; + std::vector fresh_state, lin_state; + std::map> fresh_snaps, lin_snaps; + int32_t depth = 0; + + for (size_t i = 0; i < steps.size(); ++i) { + const std::vector s = steps[i].at("tokens"); + const int32_t n_prompt = steps[i].at("n_prompt"); + const std::vector sn = steps[i].at("snapshots"); + const std::set stops(sn.begin(), sn.end()); + const std::vector prompt(s.begin(), s.begin() + n_prompt); + const std::vector reply(s.begin() + n_prompt, s.end()); + + // the state holds prompt[0, n_prompt - 1): score the reply, then take in its last token too + const auto finish_step = [&]() { + auto rows = ev.score(prompt[n_prompt - 1], reply, n_prompt - 1); + ev.generate(s.back(), (int32_t) s.size() - 1, 1); + ev.snaps[(int32_t) s.size()] = ev.snap(); + return rows; + }; + + ev.clear(); + ev.snaps.clear(); + double t0 = now_ms(); + ev.prefill(s, 0, n_prompt - 1, n_chunk, stops); + const double t_fresh = now_ms() - t0; + const auto ref = finish_step(); + std::vector fresh_now = ev.save_seq(); + auto fresh_now_snaps = ev.snaps; + + if (i > 0) { + for (const std::string arm : { "O", "L" }) { + ev.load_seq(arm == "O" ? fresh_state : lin_state); + ev.snaps = arm == "O" ? fresh_snaps : lin_snaps; + + std::vector keys; + for (const auto & kv : ev.snaps) { + keys.push_back(kv.first); + } + const auto plan = common_scr_plan_build(prev_t, prompt, keys, true, sp); + + json r = { { "id", id }, { "step", i }, { "arm", arm }, { "n_prompt", n_prompt }, { "plan", plan.str() }, + { "n_reused", plan.n_reused() }, { "t_fresh_ms", t_fresh } }; + if (plan.empty() || !scr_pays_eval(plan, n_prompt)) { + r["skipped"] = true; + if (arm == "L") { + // the lineage restarts from an exact state, as the server would prefill plainly + lin_state = fresh_now; + lin_snaps = fresh_now_snaps; + depth = 0; + } + } else { + t0 = now_ms(); + ev.splice(plan, s, n_prompt - 1, n_chunk, true, stops); + r["t_scr_ms"] = now_ms() - t0; + const auto rows = finish_step(); + const auto st = compare(ref, rows, reply, ev.n_vocab); + r["kl_mean"] = st.kl_mean; r["kl_p99"] = st.kl_p99; r["top1"] = st.top1; r["first_div"] = st.first_div; + r["nll"] = st.nll; + r["nll_fresh"] = compare(ref, ref, reply, ev.n_vocab).nll; + if (arm == "L") { + depth++; + lin_state = ev.save_seq(); + lin_snaps = ev.snaps; + } + } + r["depth"] = arm == "L" ? depth : 1; + printf("%s\n", r.dump().c_str()); + fflush(stdout); + } + } else { + lin_state = fresh_now; + lin_snaps = fresh_now_snaps; + } + + fresh_state = std::move(fresh_now); + fresh_snaps = std::move(fresh_now_snaps); + prev_t = s; + } + continue; + } + + const std::vector old_t = c.at("old"); + const std::vector new_t = c.at("new"); + std::vector snapshots = c.value("snapshots", std::vector{}); + snapshots.push_back((int32_t) old_t.size()); + std::sort(snapshots.begin(), snapshots.end()); + + const std::string answer = c.value("answer", std::string()); + const int32_t n_gen = answer.empty() ? 0 : 64; + + std::vector teacher = c.value("teacher", std::vector{}); + + const int32_t N = (int32_t) new_t.size(); + + const auto plan = common_scr_plan_build(old_t, new_t, snapshots, true, sp); + + json head = { + { "id", id }, { "n_old", old_t.size() }, { "n_new", N }, { "plan", plan.str() }, + { "n_prefix", plan.n_prefix }, { "restore", plan.restore }, { "n_reused", plan.n_reused() }, + { "resume", plan.resume() }, { "runs", plan.runs.size() }, + }; + + // F first: it defines the teacher and the reference rows + std::string text; + ev.clear(); + double t0 = now_ms(); + ev.prefill(new_t, 0, N - 1, n_chunk); + const double t_fresh = now_ms() - t0; + const auto ref = ev.finish(new_t, teacher, n_teacher, n_gen, &text); + { + json r = head; + r["arm"] = "F"; + r["t_prefill_ms"] = t_fresh; + r["prefill_tokens"] = N - 1; + if (!answer.empty()) { + r["text"] = text; + r["hit"] = text.find(answer) != std::string::npos; + } + const auto st = compare(ref, ref, teacher, ev.n_vocab); + r["nll"] = st.nll; + r["top1_teacher"] = st.top1_teacher; + printf("%s\n", r.dump().c_str()); + fflush(stdout); + } + + auto report = [&](const std::string & arm, const std::vector & rows, double t_ms, int32_t n_prefilled, const std::string & txt) { + const auto st = compare(ref, rows, teacher, ev.n_vocab); + json r = head; + r["arm"] = arm; + r["t_prefill_ms"] = t_ms; + r["prefill_tokens"] = n_prefilled; + r["kl_mean"] = st.kl_mean; r["kl_p50"] = st.kl_p50; r["kl_p99"] = st.kl_p99; r["kl_max"] = st.kl_max; + r["top1"] = st.top1; r["top1_teacher"] = st.top1_teacher; r["first_div"] = st.first_div; r["nll"] = st.nll; + if (!answer.empty()) { + r["text"] = txt; + r["hit"] = txt.find(answer) != std::string::npos; + } + printf("%s\n", r.dump().c_str()); + fflush(stdout); + }; + + if (arms.count("floor")) { + ev.clear(); + t0 = now_ms(); + ev.prefill(new_t, 0, N - 1, 512); + const double t = now_ms() - t0; + std::string txt; + auto rows = ev.finish(new_t, teacher, n_teacher, n_gen, &txt); + report("floor", rows, t, N - 1, txt); + } + + if (arms.count("RT") && !plan.empty()) { + ev.clear(); + ev.prefill(new_t, 0, N - 1, n_chunk); + for (const auto & r : plan.runs) { + GGML_ASSERT(llama_scr_kv_relocate(ctx, 0, r.new_start, r.new_start + r.n, N)); + GGML_ASSERT(llama_scr_kv_relocate(ctx, 0, r.new_start + N, r.new_start + r.n + N, -N)); + } + std::string txt; + auto rows = ev.finish(new_t, teacher, n_teacher, n_gen, &txt); + report("RT", rows, 0, N - 1, txt); + } + + for (const std::string arm : { "S", "S0" }) { + if (!arms.count(arm) || plan.empty()) { + continue; + } + if (cached_old != old_t) { + ev.clear(); + ev.snaps.clear(); + ev.prefill(old_t, 0, (int32_t) old_t.size(), n_chunk, std::set(snapshots.begin(), snapshots.end())); + cached_old = old_t; + cached_state = ev.save_seq(); + cached_snaps = ev.snaps; + } + ev.load_seq(cached_state); + ev.snaps = cached_snaps; + + t0 = now_ms(); + const int32_t n_prefilled = ev.splice(plan, new_t, N - 1, n_chunk, arm == "S"); + const double t = now_ms() - t0; + + std::string txt; + auto rows = ev.finish(new_t, teacher, n_teacher, n_gen, &txt); + report(arm, rows, t, n_prefilled, txt); + } + } + + llama_backend_free(); + return 0; +} diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index b5e052a9..03155998 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -79,6 +79,10 @@ json server_slot_stats::to_json() const { {"predicted_per_second", n_gen_tps()}, }; + if (n_prompt_scr >= 0) { + base["scr_n"] = n_prompt_scr; + } + // report draft counters whenever speculation is enabled, so "nothing drafted" differs from "off" if (speculative || n_draft_tokens > 0) { base["draft_n"] = n_draft_tokens; @@ -629,6 +633,31 @@ llama_tokens server_tokens::get_text_tokens() const { return res; } +llama_tokens server_tokens::get_scr_ids() const { + llama_tokens res = tokens; + for (const auto & [start, chunk] : map_idx_to_media) { + const uint64_t h = std::hash{}(mtmd_input_chunk_get_id(chunk.get())); + const size_t n = mtmd_input_chunk_get_n_tokens(chunk.get()); + for (size_t i = 0; i < n; ++i) { + // negative and below LLAMA_TOKEN_NULL, so never a vocab id or the placeholder + res[start + i] = -2 - (llama_token) ((h + 0x9e3779b97f4a7c15ULL*i) & 0x3fffffff); + } + } + return res; +} + +std::pair server_tokens::media_span(size_t idx) const { + auto it = map_idx_to_media.upper_bound(idx); + if (it != map_idx_to_media.begin()) { + --it; + const size_t end = it->first + mtmd_input_chunk_get_n_tokens(it->second.get()); + if (idx < end) { + return { it->first, end }; + } + } + return { idx, idx }; +} + void server_tokens::set_token(llama_pos pos, llama_token id) { GGML_ASSERT(!has_mtmd); // only allow this if mtmd is disabled tokens[pos] = id; diff --git a/tools/server/server-common.h b/tools/server/server-common.h index a624ae6a..3e00ff5b 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -209,6 +209,13 @@ public: llama_tokens get_text_tokens() const; + // the ids the suffix-cache-reuse planner diffs: text tokens as they are, and for every media token an id + // that only the same media at the same offset shares + llama_tokens get_scr_ids() const; + + // [start, end) of the media chunk covering idx, or [idx, idx) when idx is a text token + std::pair media_span(size_t idx) const; + std::vector serialize() const; static server_tokens deserialize(const llama_tokens & packed, bool has_mtmd); @@ -345,6 +352,7 @@ json format_response_rerank( struct server_slot_stats { uint64_t n_prompt_cached = 0; uint64_t n_prompt_processed = 0; + int64_t n_prompt_scr = -1; // prompt tokens reused by a splice (suffix cache reuse), -1 without one uint64_t n_gen = 0; // speculative decoding stats diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bed091d9..f657fdf0 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,6 +14,8 @@ #include "log.h" #include "sampling.h" #include "speculative.h" +#include "scr.h" +#include "../../src/llama-ext.h" // suffix cache reuse primitives #include "mtmd.h" #include "mtmd-helper.h" @@ -254,8 +256,33 @@ struct server_slot { // the request that produced the prompt asked for "cache": "skip" bool cache_skip = false; + // a splice in progress (suffix cache reuse): the plan, the next run to take over, and the recurrent state the + // last run continues from. while set, the cache holds the moved runs beyond the prefill position + struct scr_splice { + common_scr_plan plan; + size_t i_run = 0; + std::vector adopt; + }; + + std::unique_ptr scr; + + // drops a splice that did not finish: the moved runs beyond the prefill position go, the prefix stays + void scr_abort() { + if (!scr) { + return; + } + SLT_WRN(*this, "dropping an unfinished splice at %d tokens\n", prompt.n_tokens()); + scr.reset(); + llama_scr_set_spliced(ctx_tgt, id, false); + if (ctx_dft) { + llama_scr_set_spliced(ctx_dft, id, false); + } + mem.seq_rm(id, prompt.tokens.pos_next(), -1); + } + bool prompt_save(server_prompt_cache & prompt_cache) const { - if (prompt.tokens.size() == 0 || cache_skip) { + // a spliced state is never stored: the cache would hand it to requests that did not ask for SCR + if (prompt.tokens.size() == 0 || cache_skip || prompt.scr_depth > 0) { return false; } @@ -506,6 +533,8 @@ struct server_slot { SLT_INF(*this, "stop processing: n_tokens = %d, truncated = %d\n", prompt.n_tokens(), truncated); + scr_abort(); + t_last_used = ggml_time_us(); state = SLOT_STATE_IDLE; @@ -1608,6 +1637,28 @@ private: } } + // a request that allows SCR takes the slot whose cached state it can reuse the most of + bool scr_keep = false; + if (task.params.scr && task.type == SERVER_TASK_TYPE_COMPLETION && task.id_slot == -1) { + int32_t n_best = ret ? (int32_t) ret->prompt.tokens.get_common_prefix(task.tokens) : 0; + for (server_slot & slot : slots) { + if (slot.is_processing()) { + continue; + } + const auto plan = scr_plan(slot, task.tokens); + if (scr_pays(plan, task.tokens.size()) && plan.n_prefix + plan.n_reused() > n_best) { + n_best = plan.n_prefix + plan.n_reused(); + ret = &slot; + scr_keep = true; + } + } + if (scr_keep) { + const size_t lcp = ret->prompt.tokens.get_common_prefix(task.tokens); + SLT_INF(*ret, "selected slot by suffix cache reuse, %d tokens reusable\n", n_best); + update_cache = float(lcp) / ret->prompt.tokens.size() < 0.5f; + } + } + // find the slot that has been least recently used if (ret == nullptr) { int64_t t_last = -1; @@ -1656,7 +1707,8 @@ private: ret->prompt_save(*prompt_cache); - if (!ret->prompt_load(*prompt_cache, task.tokens, skip)) { + // a slot chosen for SCR keeps its state: the splice reuses it + if (!scr_keep && !ret->prompt_load(*prompt_cache, task.tokens, skip)) { ret->prompt_clear(); } @@ -1664,6 +1716,10 @@ private: SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); } + + if (prompt_cache && task.params.scr && task.type == SERVER_TASK_TYPE_COMPLETION) { + scr_load_entry(*ret, task, update_cache); + } } return ret; @@ -2322,6 +2378,273 @@ private: return true; } + // suffix cache reuse: plans the reuse of a cached prompt whose recurrent state has absorbed its first n_state + // tokens, in token indices. the recurrent state can continue from that state and from the prompt's checkpoints. + // a media chunk is reused whole or not at all. empty when SCR does not apply + common_scr_plan scr_plan(const server_prompt & prompt, int32_t n_state, const server_tokens & tokens) const { + if (!llama_model_is_hybrid(model_tgt) || n_state <= 0 || n_state > prompt.n_tokens()) { + return {}; + } + + llama_tokens old = prompt.tokens.get_scr_ids(); + old.resize(n_state); + + std::vector snapshots = { n_state }; + for (const auto & ckpt : prompt.checkpoints) { + if (ckpt.n_tokens <= n_state && !ckpt.data_tgt.empty()) { + snapshots.push_back((int32_t) ckpt.n_tokens); + } + } + std::sort(snapshots.begin(), snapshots.end()); + + common_scr_params params; + common_scr_plan plan = common_scr_plan_build(old, tokens.get_scr_ids(), snapshots, true, params); + + // a run starting inside an image starts after it; one ending inside an image ends before it + // (the last run ends at a snapshot, which never falls inside an image) + std::vector runs; + for (size_t i = 0; i < plan.runs.size(); ++i) { + auto r = plan.runs[i]; + const bool last = i + 1 == plan.runs.size(); + + const auto [s0, s1] = tokens.media_span(r.new_start); + if (r.new_start > (int32_t) s0) { + const int32_t d = (int32_t) s1 - r.new_start; + r.new_start += d; + r.old_start += d; + r.n -= d; + } + if (!last) { + const auto [e0, e1] = tokens.media_span(r.new_start + r.n); + if (r.new_start + r.n > (int32_t) e0) { + r.n = (int32_t) e0 - r.new_start; + } + } + + if (r.n < params.min_span) { + if (last) { + return {}; + } + continue; + } + runs.push_back(r); + } + plan.runs = std::move(runs); + + return plan; + } + + // a plan in token indices as KV positions (they differ once an M-RoPE image takes fewer positions than tokens) + static common_scr_plan scr_plan_pos(const common_scr_plan & plan, const server_tokens & old, const server_tokens & cur) { + common_scr_plan res = plan; + res.restore = cur.pos_next(plan.restore); + for (auto & r : res.runs) { + const llama_pos old_start = old.pos_next(r.old_start); + r.n = old.pos_next(r.old_start + r.n) - old_start; + r.old_start = old_start; + r.new_start = cur.pos_next(r.new_start); + } + return res; + } + + // the tokens the slot's live recurrent state has absorbed + int32_t scr_n_state(const server_slot & slot) const { + return slot.prompt.tokens.size_up_to_pos(llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id) + 1); + } + + common_scr_plan scr_plan(const server_slot & slot, const server_tokens & tokens) const { + return scr_plan(slot.prompt, scr_n_state(slot), tokens); + } + + // the plan pays when it reuses at least 512 tokens and a tenth of what a plain request would prefill + static bool scr_pays(const common_scr_plan & plan, int32_t n_tokens) { + return !plan.empty() && plan.n_reused() >= std::max(512, (n_tokens - plan.restore) / 10); + } + + // loads the prompt-cache entry (RAM or disk) a splice can reuse more of than the slot's own state into the slot. + // the slot's state goes to the cache first unless saved says it is there already + void scr_load_entry(server_slot & slot, const server_task & task, bool saved) { + const auto value = [&](const common_scr_plan & plan, const server_prompt & prompt) { + return scr_pays(plan, task.tokens.size()) ? plan.n_prefix + plan.n_reused() : (int32_t) prompt.tokens.get_common_prefix(task.tokens); + }; + + // planning against every cached prompt is O(cache size); only prompts that start like this one are planned + const auto value_of = [&](const server_prompt & prompt, int32_t n_state) { + const int32_t lcp = prompt.tokens.get_common_prefix(task.tokens); + return lcp < 16 ? lcp : value(scr_plan(prompt, n_state, task.tokens), prompt); + }; + + const int32_t n_own = value(scr_plan(slot, task.tokens), slot.prompt); + + // an entry's own end stands in for its live state until it is loaded + int32_t n_best = n_own; + uint64_t id_best = 0; + for (const auto & s : prompt_cache->states) { + if (s.data.main.empty()) { + continue; + } + const int32_t n = value_of(s.prompt, s.prompt.n_tokens()); + if (n > n_best) { + n_best = n; + id_best = s.prompt.meta.id; + } + } + + // disk entries are text prompts without their checkpoints in the index + const std::string path = prompt_cache->disk_best([&](const llama_tokens & tokens) { + server_prompt prompt; + prompt.tokens = server_tokens(tokens, false); + return value_of(prompt, prompt.n_tokens()); + }, n_best); + + if (n_best == n_own) { + return; + } + + if (!saved && slot.prompt_save(*prompt_cache)) { + prompt_cache->update(); + } + + // saving can evict RAM entries, so a RAM entry is looked up again by its id + auto best = !path.empty() ? prompt_cache->disk_promote(path, task.tokens.has_mtmd) : + std::find_if(prompt_cache->states.begin(), prompt_cache->states.end(), [&](const server_prompt_cache_state & s) { return s.prompt.meta.id == id_best; }); + if (best == prompt_cache->states.end()) { + return; + } + + const bool skip = task.params.cache_explicit && task.params.cache_cls == CACHE_CLASS_SKIP; + + SLT_INF(slot, "scr: loading a cached prompt of %d tokens from %s, %d reusable (the slot's own state: %d)\n", + best->prompt.n_tokens(), path.empty() ? "RAM" : "disk", n_best, n_own); + if (!prompt_cache->load_entry(best, slot.prompt, ctx_tgt, ctx_dft, slot.id, skip && path.empty())) { + slot.prompt_clear(); + } + prompt_cache->update(); + } + + // starts a splice of the slot's cached state for its task (a request with "scr": true): rearranges the + // attention KV, restarts the recurrent state at the plan's restore point and leaves the rest to the prompt loop + // (see scr_take_run). returns the number of cached tokens to keep, or -1 when the request is served without SCR + int32_t scr_begin(server_slot & slot) { + const auto & tokens = slot.task->tokens; + + const common_scr_plan plan = scr_plan(slot, tokens); + if (!scr_pays(plan, tokens.size())) { + if (!plan.empty()) { + SLT_INF(slot, "scr: not worth it (%s, %d of %d tokens)\n", plan.str().c_str(), plan.n_reused(), (int) tokens.size()); + } + return -1; + } + + const int32_t n_state = scr_n_state(slot); + + const common_prompt_checkpoint * ckpt_restore = nullptr; + const common_prompt_checkpoint * ckpt_adopt = nullptr; + for (const auto & ckpt : slot.prompt.checkpoints) { + if (ckpt.n_tokens == plan.restore) { + ckpt_restore = &ckpt; + } + if (ckpt.n_tokens == plan.adopt) { + ckpt_adopt = &ckpt; + } + } + if (plan.restore > 0 && ckpt_restore == nullptr) { + return -1; + } + + // the state the last run continues from, taken before anything changes + std::vector adopt; + if (plan.adopt == n_state) { + adopt.resize(llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)); + if (llama_state_seq_get_data_ext(ctx_tgt, adopt.data(), adopt.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) != adopt.size()) { + return -1; + } + } else if (ckpt_adopt != nullptr) { + adopt = ckpt_adopt->data_tgt; + } else { + return -1; + } + + const int64_t t_start = ggml_time_us(); + + // the cache works in positions + const common_scr_plan plan_pos = scr_plan_pos(plan, slot.prompt.tokens, tokens); + + const int32_t rc = common_scr_apply_kv(ctx_tgt, slot.id, plan_pos); + if (rc == 0) { + return -1; + } + if (rc < 0) { + SLT_ERR(slot, "%s", "scr: moving the cache failed, prefilling from scratch\n"); + slot.prompt_clear(); + return 0; + } + + if (ctx_dft) { + // without its runs the draft only predicts worse + if (common_scr_apply_kv(ctx_dft, slot.id, plan_pos) != 1) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, plan_pos.restore, -1); + } + llama_scr_set_spliced(ctx_dft, slot.id, true); + } + llama_scr_set_spliced(ctx_tgt, slot.id, true); + + if (ckpt_restore != nullptr) { + ckpt_restore->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + common_speculative_set_state(spec.get(), slot.id, ckpt_restore->data_spec); + } else { + llama_scr_recr_set_pos(ctx_tgt, slot.id, -1); + } + + // checkpoints past the restore point belong to the old text + slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & ckpt) { return ckpt.n_tokens > plan.restore; }); + + const int32_t exact = plan.runs.front().new_start; + slot.prompt.scr_exact = slot.prompt.scr_depth > 0 ? std::min(slot.prompt.scr_exact, exact) : exact; + slot.prompt.scr_depth++; + + slot.stats.n_prompt_scr = plan.n_reused(); + + SLT_INF(slot, "scr: %s, reusing %d of %d tokens (depth %d, exact %d), %.2f ms\n", plan.str().c_str(), + plan.n_reused(), (int) tokens.size(), slot.prompt.scr_depth, slot.prompt.scr_exact, (ggml_time_us() - t_start) / 1000.0); + + slot.scr = std::make_unique(server_slot::scr_splice { plan, 0, std::move(adopt) }); + + return plan.restore; + } + + // the prefill of a splicing slot has reached its next run: the run's tokens are already in the cache. after the + // last run the recurrent state continues from the pre-edit snapshot, and the splice is done + void scr_take_run(server_slot & slot) { + auto & s = *slot.scr; + + const auto & tokens = slot.task->tokens; + const auto & run = s.plan.runs[s.i_run]; + for (int32_t i = run.new_start; i < run.new_start + run.n; ) { + if (tokens[i] == LLAMA_TOKEN_NULL) { + // the image is in the cache already: only its placeholder joins the prompt + const auto & chunk = tokens.find_chunk(i); + slot.prompt.tokens.push_back_placeholder(chunk.get()); + i += mtmd_input_chunk_get_n_tokens(chunk.get()); + } else { + slot.prompt.tokens.push_back(tokens[i++]); + } + } + + if (++s.i_run == s.plan.runs.size()) { + if (llama_state_seq_set_data_ext(ctx_tgt, s.adopt.data(), s.adopt.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) != s.adopt.size()) { + GGML_ABORT("scr: failed to restore the recurrent state\n"); + } + llama_scr_set_spliced(ctx_tgt, slot.id, false); + if (ctx_dft) { + llama_scr_set_spliced(ctx_dft, slot.id, false); + } + slot.scr.reset(); + } + + llama_scr_recr_set_pos(ctx_tgt, slot.id, slot.prompt.tokens.pos_next() - 1); + } + // n_tokens_cur: the number of tokens added to the batch for the current slot void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) { const int id_task = slot.task->id; @@ -3287,6 +3610,9 @@ private: // keep track how many tokens we can reuse from the previous state int n_past = 0; + // a splice (suffix cache reuse) has rearranged the cache for this prompt + bool scr_begun = false; + // empty prompt passed -> release the slot and send empty response if (input_tokens.empty()) { SLT_WRN(slot, "%s", "empty prompt - releasing slot\n"); @@ -3342,6 +3668,23 @@ private: // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); + // a request without "scr" reuses nothing a splice derived + if (slot.prompt.scr_depth > 0 && !slot.task->params.scr) { + const int32_t exact = slot.prompt.scr_exact; + SLT_INF(slot, "scr: the cached state is spliced, reusing at most its exact %d tokens\n", exact); + n_past = std::min(n_past, exact); + slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & ckpt) { return ckpt.n_tokens > exact; }); + slot.prompt.scr_depth = 0; + } + + if (slot.task->params.scr && n_past + 1 < slot.task->n_tokens() && slot.alora_invocation_start <= 0) { + const int32_t n_keep = scr_begin(slot); + if (n_keep >= 0) { + n_past = n_keep; + scr_begun = true; + } + } + // if there is an alora invoked, don't cache after the invocation start if (slot.alora_invocation_start > 0) { SLT_DBG(slot, "only caching to alora invocation start (n_past = %d, alora_invocation_start = %d)\n", n_past, slot.alora_invocation_start); @@ -3359,7 +3702,7 @@ private: } // reuse chunks from the cached prompt by shifting their KV cache in the new position - if (can_cache_reuse && n_cache_reuse > 0) { + if (!scr_begun && can_cache_reuse && n_cache_reuse > 0) { GGML_ASSERT(!slot.prompt.tokens.has_mtmd); size_t head_c = n_past; // cache @@ -3420,7 +3763,7 @@ private: // the largest pos_min required for a checkpoint to be useful const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); - if (n_past > 0 && n_past <= slot.prompt.n_tokens()) { + if (!scr_begun && n_past > 0 && n_past <= slot.prompt.n_tokens()) { const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); if (pos_min == -1) { SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.id, pos_min); @@ -3565,7 +3908,10 @@ private: SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - slot.mem.seq_rm(slot.id, p0, -1); + // a splice keeps its moved runs beyond the prefill position + if (!slot.scr) { + slot.mem.seq_rm(slot.id, p0, -1); + } // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3608,6 +3954,11 @@ private: break; } + // a splice's run starting with this image has it in the cache already + if (slot.scr && (int32_t) cur_token_idx == slot.scr->plan.runs[slot.scr->i_run].new_start) { + break; + } + // process the mtmd chunk // note: it submits its own decode, potentially be async // so the timing is queued and flushed on the next sync @@ -3646,6 +3997,15 @@ private: // add prompt tokens for processing in the current batch while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) { + // a splice takes its next run over once the prefill in front of it is decoded + if (slot.scr && slot.prompt.n_tokens() == slot.scr->plan.runs[slot.scr->i_run].new_start) { + if (batch.size() > n_tokens_prev) { + break; + } + scr_take_run(slot); + continue; + } + // get next token to process llama_token cur_tok = input_tokens[slot.prompt.n_tokens()]; if (cur_tok == LLAMA_TOKEN_NULL) { @@ -3745,8 +4105,8 @@ private: do_checkpoint = false; } - // do not checkpoint after mtmd chunks - do_checkpoint = do_checkpoint && !has_mtmd; + // do not checkpoint after mtmd chunks, or while a splice has the recurrent state between snapshots + do_checkpoint = do_checkpoint && !has_mtmd && !slot.scr; // no need to create checkpoints that are too close together, unless it's the last user message do_checkpoint = do_checkpoint && ( diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 8254e4a8..f011dfa7 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -41,6 +41,9 @@ std::vector> make_llama_cmpl_schema(const common_params & ctx.params.cache_explicit = true; })); + add((new field_bool("scr", params.scr)) + ->set_desc("Suffix cache reuse: after an edit in the middle of the prompt, reuse the cached state of the text after the edit instead of prefilling it again. The reused state was computed over the old text, so the output can differ from a fresh prefill")); + add((new field_num("cache_ttl", params.cache_ttl)) ->set_hard_limits(0, std::numeric_limits::max()) ->set_desc("Seconds the prompt stays in the prompt cache after its last use, instead of its class default")); @@ -540,6 +543,7 @@ task_params eval_llama_cmpl_schema( params.n_predict = params_base.n_predict; params.n_cache_reuse = params_base.n_cache_reuse; params.cache_prompt = params_base.cache_prompt; + params.scr = params_base.scr; params.antiprompt = params_base.antiprompt; params.sse_ping_interval = params_base.sse_ping_interval; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index e67050a9..0a83f4d0 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -2309,40 +2309,11 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok bool from_disk = false; if (!disk_best.empty()) { - namespace fs = std::filesystem; - - const int64_t t_start = ggml_time_us(); - server_prompt_cache_state state; - const bool ok = disk_read(disk_best, tokens_new.has_mtmd, state); - - std::lock_guard lock(disk_mutex); - std::error_code ec; - auto it = std::find_if(disk.begin(), disk.end(), [&](const disk_entry & e) { return e.path == disk_best; }); - if (ok) { - SRV_INF("prompt cache: read %zu tokens from disk in %.2f s, f_keep = %.3f, f_sim = %.3f\n", - state.prompt.tokens.size(), (ggml_time_us() - t_start) / 1e6, f_keep_best, f_sim_best); - if (it != disk.end()) { - // the file stays: it records the use, and is most recently used, here and after a restart - it->meta.last_used = unix_now(); - it->meta.hits++; - set_expiry(it->meta); - state.prompt.meta = it->meta; - state.prompt.meta.hits--; // the slot counts this use when the request continues the prompt - if (it->version == DISK_VERSION) { - write_meta(it->path, it->meta); - } - disk.splice(disk.end(), disk, it); - } - states.push_back(std::move(state)); - it_best = std::prev(states.end()); + const auto it = disk_promote(disk_best, tokens_new.has_mtmd); + if (it != states.end()) { + SRV_INF("prompt cache: f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best); + it_best = it; from_disk = true; - fs::last_write_time(disk_best, fs::file_time_type::clock::now(), ec); - } else { - SRV_WRN("prompt cache: cannot read %s, removing it\n", disk_best.c_str()); - if (it != disk.end()) { - disk.erase(it); - } - fs::remove(disk_best, ec); } } @@ -2350,56 +2321,119 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best); // a skip request restores a copy and leaves the RAM entry where it is - const bool consume = !keep_entry || from_disk; + if (!load_entry(it_best, prompt, ctx_tgt, ctx_dft, id_slot, keep_entry && !from_disk)) { + return false; + } + } - { - auto & data = it_best->data.main; + update_ram_next_expiry(); - const size_t size = data.size(); - const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0); - if (n != size) { - SRV_ERR("failed to restore state with size %zu\n", size); + return true; +} - return false; - } +std::list::iterator server_prompt_cache::disk_promote(const std::string & path, bool has_mtmd) { + namespace fs = std::filesystem; - if (consume) { - data.clear(); - data.shrink_to_fit(); - } + const int64_t t_start = ggml_time_us(); + server_prompt_cache_state state; + const bool ok = disk_read(path, has_mtmd, state); + + std::lock_guard lock(disk_mutex); + std::error_code ec; + auto it = std::find_if(disk.begin(), disk.end(), [&](const disk_entry & e) { return e.path == path; }); + if (!ok) { + SRV_WRN("prompt cache: cannot read %s, removing it\n", path.c_str()); + if (it != disk.end()) { + disk.erase(it); } + fs::remove(path, ec); + return states.end(); + } - { - auto & data = it_best->data.drft; + SRV_INF("prompt cache: read %zu tokens from disk in %.2f s\n", state.prompt.tokens.size(), (ggml_time_us() - t_start) / 1e6); + if (it != disk.end()) { + // the file stays: it records the use, and is most recently used, here and after a restart + it->meta.last_used = unix_now(); + it->meta.hits++; + set_expiry(it->meta); + state.prompt.meta = it->meta; + state.prompt.meta.hits--; // the slot counts this use when the request continues the prompt + if (it->version == DISK_VERSION) { + write_meta(it->path, it->meta); + } + disk.splice(disk.end(), disk, it); + } + states.push_back(std::move(state)); + fs::last_write_time(path, fs::file_time_type::clock::now(), ec); - if (!data.empty()) { - GGML_ASSERT(ctx_dft); + return std::prev(states.end()); +} - const size_t size = data.size(); - const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0); - if (n != size) { - SRV_WRN("failed to restore state with size %zu\n", size); +std::string server_prompt_cache::disk_best(const std::function & score, int32_t & n_best) { + std::string best; + std::lock_guard lock(disk_mutex); + if (!disk_thread.joinable()) { + return best; + } + expire_disk_locked(); + for (const auto & e : disk) { + const int32_t n = score(e.tokens); + if (n > n_best) { + n_best = n; + best = e.path; + } + } + return best; +} - return false; - } +bool server_prompt_cache::load_entry(std::list::iterator it, server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, bool keep_entry) { + const bool consume = !keep_entry; - if (consume) { - data.clear(); - data.shrink_to_fit(); - } - } + { + auto & data = it->data.main; + + const size_t size = data.size(); + const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0); + if (n != size) { + SRV_ERR("failed to restore state with size %zu\n", size); + + return false; } if (consume) { - prompt = std::move(it_best->prompt); + data.clear(); + data.shrink_to_fit(); + } + } - states.erase(it_best); - } else { - prompt = it_best->prompt.clone(); + { + auto & data = it->data.drft; + + if (!data.empty()) { + GGML_ASSERT(ctx_dft); + + const size_t size = data.size(); + const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0); + if (n != size) { + SRV_WRN("failed to restore state with size %zu\n", size); + + return false; + } + + if (consume) { + data.clear(); + data.shrink_to_fit(); + } } } - update_ram_next_expiry(); + if (consume) { + prompt = std::move(it->prompt); + + states.erase(it); + } else { + prompt = it->prompt.clone(); + } return true; } diff --git a/tools/server/server-task.h b/tools/server/server-task.h index c4c37df9..21554dfe 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -87,6 +87,9 @@ struct task_params { cache_class cache_cls = CACHE_CLASS_DEFAULT; bool cache_explicit = false; // "cache" was given int64_t cache_ttl = -1; // seconds, -1 = the class default + + // "scr": allow suffix cache reuse, an approximate reuse of the cached text that survived an edit + bool scr = false; bool return_tokens = false; bool return_progress = false; @@ -646,10 +649,17 @@ struct server_prompt { cache_meta meta; + // suffix cache reuse: splices behind this state (0 = exact), and how many leading tokens are still exact. + // a request that did not ask for SCR reuses at most scr_exact tokens and only checkpoints within them + int32_t scr_depth = 0; + int32_t scr_exact = 0; + void clear() { tokens.clear(); checkpoints.clear(); meta = {}; + scr_depth = 0; + scr_exact = 0; } int n_tokens() const { @@ -661,6 +671,8 @@ struct server_prompt { tokens.clone(), checkpoints, meta, + scr_depth, + scr_exact, }; } }; @@ -725,6 +737,15 @@ struct server_prompt_cache { // keep_entry: restore from a RAM entry without consuming it (a "skip" request) bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, bool keep_entry = false); + // restores the RAM entry it into a slot's prompt; the entry is consumed unless keep_entry + bool load_entry(std::list::iterator it, server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, bool keep_entry); + + // the disk entry whose tokens score highest, above n_best (which it raises to that score); empty if none + std::string disk_best(const std::function & score, int32_t & n_best); + + // reads a disk entry into RAM and returns it, or states.end() if the file cannot be read (it is then removed) + std::list::iterator disk_promote(const std::string & path, bool has_mtmd); + void update(); // second tier on disk: an entry the RAM limit evicts is written to dir instead of dropped, and read back