diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt
index 5606a966..574c0836 100644
--- a/tools/server/CMakeLists.txt
+++ b/tools/server/CMakeLists.txt
@@ -9,6 +9,8 @@ add_library(${TARGET} STATIC
server-chat.h
server-task.cpp
server-task.h
+ server-kv-share.cpp
+ server-kv-share.h
server-queue.cpp
server-queue.h
server-common.cpp
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 1a3ff43b..ee17a5b6 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -371,6 +371,7 @@ struct server_slot {
if (ctx_dft) {
llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);
}
+ prompt_cache.share(*cur);
return true;
}
@@ -1594,6 +1595,10 @@ private:
SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n");
prompt_cache = std::make_unique<server_prompt_cache>(params_base.cache_ram_mib, n_ctx);
+ {
+ const auto rope = llama_model_rope_type(model_tgt);
+ prompt_cache->share_layout.cell_ext = rope == LLAMA_ROPE_TYPE_MROPE || rope == LLAMA_ROPE_TYPE_IMROPE ? 2*sizeof(llama_pos) : 0;
+ }
{
cache_policy policy;
diff --git a/tools/server/server-kv-share.cpp b/tools/server/server-kv-share.cpp
new file mode 100644
index 00000000..330bbeda
--- /dev/null
+++ b/tools/server/server-kv-share.cpp
@@ -0,0 +1,372 @@
+#include "server-kv-share.h"
+
+#include <cstring>
+#include <unordered_map>
+
+namespace {
+
+constexpr uint32_t DELTA_MAGIC = 0x44564b42; // "BKVD"
+constexpr uint32_t DELTA_VERSION = 1;
+
+constexpr uint8_t OP_COPY = 0;
+constexpr uint8_t OP_LITERAL = 1;
+
+// the parent's identity in a delta: its id, size and a hash of its first and last MiB (a parent never changes, so
+// this catches a delta read against another file with the same id, not edits)
+uint64_t sample_hash(const std::vector<uint8_t> & d) {
+ uint64_t h = 1469598103934665603ull;
+ auto mix = [&](size_t from, size_t to) {
+ for (size_t i = from; i < to; ++i) {
+ h = (h ^ d[i]) * 1099511628211ull;
+ }
+ };
+ const size_t n = std::min<size_t>(d.size(), 1u << 20);
+ mix(0, n);
+ mix(d.size() - n, d.size());
+ return h ^ d.size();
+}
+
+// a forward-only reader over a state blob
+struct reader {
+ const uint8_t * p;
+ size_t n;
+ size_t at = 0;
+
+ bool skip(size_t k) {
+ if (k > n - at) {
+ return false;
+ }
+ at += k;
+ return true;
+ }
+ template <typename T> bool get(T & v) {
+ if (sizeof(T) > n - at) {
+ return false;
+ }
+ std::memcpy(&v, p + at, sizeof(T));
+ at += sizeof(T);
+ return true;
+ }
+};
+
+// where one stream's rows are: per cell its key (position and extra position bytes), and per layer section the
+// offset of row 0 and the row size. a transposed V section has one row per embedding element, each holding one
+// element per cell
+struct kv_section {
+ size_t base; // offset of the first cell's row (element) in the blob
+ size_t stride; // bytes between cells
+ size_t width; // bytes per cell
+ size_t n_rows; // 1, or n_embd_v_gqa for a transposed V
+ size_t row_step; // bytes between those rows
+};
+
+struct kv_stream {
+ std::vector<std::string> keys;
+ std::vector<kv_section> sections;
+};
+
+// the attention KV part of a llama sequence state: magic, seq id, n_stream, then per stream the cell count, the cell
+// metadata and the K and V rows (llama_kv_cache::state_write). whatever follows (the recurrent state) is not parsed
+bool parse_kv(const std::vector<uint8_t> & blob, const kv_share_layout & layout, std::vector<kv_stream> & out) {
+ reader r { blob.data(), blob.size() };
+ uint32_t magic, n_stream;
+ int32_t seq_id;
+ if (!r.get(magic) || !r.get(seq_id) || !r.get(n_stream) || n_stream == 0 || n_stream > 64) {
+ return false;
+ }
+ for (uint32_t s = 0; s < n_stream; ++s) {
+ uint32_t n_cell;
+ if (!r.get(n_cell)) {
+ return false;
+ }
+ kv_stream st;
+ if (n_cell == 0) {
+ out.push_back(std::move(st));
+ continue;
+ }
+ st.keys.reserve(n_cell);
+ for (uint32_t i = 0; i < n_cell; ++i) {
+ const size_t key_at = r.at;
+ int32_t pos;
+ uint32_t n_seq;
+ if (!r.get(pos) || !r.get(n_seq) || n_seq > 256) {
+ return false;
+ }
+ const size_t ext_at = r.at;
+ if (!r.skip(layout.cell_ext) || !r.skip(n_seq * sizeof(int32_t))) {
+ return false;
+ }
+ std::string key(reinterpret_cast<const char *>(blob.data() + key_at), sizeof(pos));
+ key.append(reinterpret_cast<const char *>(blob.data() + ext_at), layout.cell_ext);
+ st.keys.push_back(std::move(key));
+ }
+ uint32_t v_trans, n_layer;
+ if (!r.get(v_trans) || !r.get(n_layer) || n_layer == 0 || n_layer > 4096 || v_trans > 1) {
+ return false;
+ }
+ for (uint32_t l = 0; l < n_layer; ++l) {
+ int32_t type;
+ uint64_t row;
+ if (!r.get(type) || !r.get(row) || row == 0 || row > (1u << 24)) {
+ return false;
+ }
+ st.sections.push_back({ r.at, row, row, 1, 0 });
+ if (!r.skip(row * n_cell)) {
+ return false;
+ }
+ }
+ // every layer of this cache has a V (no MLA): a layer without one would make the rest unparseable, which the
+ // round trip in kv_share_encode then refuses
+ for (uint32_t l = 0; l < n_layer; ++l) {
+ int32_t type;
+ if (!r.get(type)) {
+ return false;
+ }
+ if (!v_trans) {
+ uint64_t row;
+ if (!r.get(row) || row == 0 || row > (1u << 24)) {
+ return false;
+ }
+ st.sections.push_back({ r.at, row, row, 1, 0 });
+ if (!r.skip(row * n_cell)) {
+ return false;
+ }
+ } else {
+ uint32_t el, n_embd;
+ if (!r.get(el) || !r.get(n_embd) || el == 0 || el > 16 || n_embd == 0 || n_embd > (1u << 20)) {
+ return false;
+ }
+ st.sections.push_back({ r.at, el, el, n_embd, (size_t) el * n_cell });
+ if (!r.skip((size_t) el * n_cell * n_embd)) {
+ return false;
+ }
+ }
+ }
+ out.push_back(std::move(st));
+ }
+ return true;
+}
+
+struct delta_writer {
+ std::vector<uint8_t> & out;
+ const uint8_t * child;
+ size_t lit_from = 0; // start of the literal run not yet written
+ size_t copy_src = 0;
+ size_t copy_len = 0;
+ uint32_t n_ops = 0;
+ size_t shared = 0;
+
+ template <typename T> void put(const T & v) {
+ const auto * b = reinterpret_cast<const uint8_t *>(&v);
+ out.insert(out.end(), b, b + sizeof(T));
+ }
+ void flush_copy() {
+ if (copy_len == 0) {
+ return;
+ }
+ put(OP_COPY);
+ put((uint64_t) copy_src);
+ put((uint64_t) copy_len);
+ shared += copy_len;
+ copy_len = 0;
+ n_ops++;
+ }
+ void literal_to(size_t at) {
+ if (at <= lit_from) {
+ return;
+ }
+ flush_copy();
+ put(OP_LITERAL);
+ put((uint64_t) (at - lit_from));
+ out.insert(out.end(), child + lit_from, child + at);
+ lit_from = at;
+ n_ops++;
+ }
+ // the child's bytes [at, at+len) equal the parent's [src, src+len); spans come in child order without overlap
+ bool copy(size_t at, size_t src, size_t len) {
+ if (at < lit_from) {
+ return false;
+ }
+ literal_to(at);
+ if (copy_len > 0 && copy_src + copy_len == src) {
+ copy_len += len;
+ } else {
+ flush_copy();
+ copy_src = src;
+ copy_len = len;
+ }
+ lit_from = at + len;
+ return true;
+ }
+};
+
+} // namespace
+
+bool kv_share_is_delta(const std::vector<uint8_t> & blob) {
+ uint32_t magic = 0;
+ if (blob.size() >= sizeof(magic)) {
+ std::memcpy(&magic, blob.data(), sizeof(magic));
+ }
+ return magic == DELTA_MAGIC;
+}
+
+uint64_t kv_share_parent_id(const std::vector<uint8_t> & blob) {
+ uint64_t id = 0;
+ if (kv_share_is_delta(blob) && blob.size() >= 16) {
+ std::memcpy(&id, blob.data() + 8, sizeof(id));
+ }
+ return id;
+}
+
+bool kv_share_encode(const std::vector<uint8_t> & child, const kv_share_parent & parent, const kv_share_layout & layout,
+ std::vector<uint8_t> & delta, size_t & n_shared) {
+ delta.clear();
+ n_shared = 0;
+ if (parent.data == nullptr || parent.data->empty() || child.empty() || kv_share_is_delta(child)) {
+ return false;
+ }
+ std::vector<kv_stream> cs, ps;
+ if (!parse_kv(child, layout, cs) || !parse_kv(*parent.data, layout, ps) || cs.size() != ps.size()) {
+ return false;
+ }
+
+ // every (child offset, parent offset, length) the two states share, in child order
+ struct span { size_t at, src, len; };
+ std::vector<span> spans;
+ const uint8_t * c = child.data();
+ const uint8_t * p = parent.data->data();
+ for (size_t s = 0; s < cs.size(); ++s) {
+ const kv_stream & a = cs[s];
+ const kv_stream & b = ps[s];
+ if (a.keys.empty() || b.keys.empty() || a.sections.size() != b.sections.size()) {
+ continue;
+ }
+ std::unordered_map<std::string, size_t> where;
+ where.reserve(b.keys.size());
+ for (size_t i = 0; i < b.keys.size(); ++i) {
+ if (!where.emplace(b.keys[i], i).second) {
+ where.clear(); // a position twice (images): no unambiguous match, nothing of this stream is shared
+ break;
+ }
+ }
+ if (where.empty()) {
+ continue;
+ }
+ std::vector<int64_t> match(a.keys.size(), -1);
+ for (size_t i = 0; i < a.keys.size(); ++i) {
+ auto it = where.find(a.keys[i]);
+ if (it != where.end()) {
+ match[i] = (int64_t) it->second;
+ }
+ }
+ for (size_t k = 0; k < a.sections.size(); ++k) {
+ const kv_section & x = a.sections[k];
+ const kv_section & y = b.sections[k];
+ if (x.width != y.width || x.n_rows != y.n_rows) {
+ continue;
+ }
+ for (size_t row = 0; row < x.n_rows; ++row) {
+ for (size_t i = 0; i < a.keys.size(); ++i) {
+ if (match[i] < 0) {
+ continue;
+ }
+ const size_t at = x.base + row * x.row_step + i * x.stride;
+ const size_t src = y.base + row * y.row_step + (size_t) match[i] * y.stride;
+ if (std::memcmp(c + at, p + src, x.width) == 0) {
+ spans.push_back({ at, src, x.width });
+ }
+ }
+ }
+ }
+ }
+
+ size_t total = 0;
+ for (const auto & sp : spans) {
+ total += sp.len;
+ }
+ // a delta has to save at least this much
+ constexpr size_t MIN_SAVED = 16u << 20;
+ if (total < MIN_SAVED) {
+ return false;
+ }
+
+ delta.reserve(child.size() - total + spans.size() / 4 + 64);
+ delta_writer w { delta, c };
+ w.put(DELTA_MAGIC);
+ w.put(DELTA_VERSION);
+ w.put(parent.id);
+ w.put((uint64_t) parent.data->size());
+ w.put(sample_hash(*parent.data));
+ w.put((uint64_t) child.size());
+ const size_t n_ops_at = delta.size();
+ w.put((uint32_t) 0);
+ // every byte of the delta's output is the child's own (a literal) or a parent byte compared equal to the child's
+ // byte at that offset, so decoding gives back exactly this state
+ for (const auto & sp : spans) {
+ if (!w.copy(sp.at, sp.src, sp.len)) {
+ delta.clear();
+ return false;
+ }
+ }
+ w.literal_to(child.size());
+ w.flush_copy();
+ std::memcpy(delta.data() + n_ops_at, &w.n_ops, sizeof(w.n_ops));
+ // shared rows scattered cell by cell cost more in copy records than they save
+ if (delta.size() + MIN_SAVED > child.size()) {
+ delta.clear();
+ return false;
+ }
+ delta.shrink_to_fit();
+ n_shared = w.shared;
+ return true;
+}
+
+bool kv_share_decode(const std::vector<uint8_t> & delta, const kv_share_parent & parent, std::vector<uint8_t> & out,
+ std::string & error) {
+ reader r { delta.data(), delta.size() };
+ uint32_t magic, version, n_ops;
+ uint64_t id, parent_size, parent_hash, size;
+ if (!r.get(magic) || magic != DELTA_MAGIC || !r.get(version) || version != DELTA_VERSION || !r.get(id) ||
+ !r.get(parent_size) || !r.get(parent_hash) || !r.get(size) || !r.get(n_ops)) {
+ error = "not a shared state";
+ return false;
+ }
+ if (parent.data == nullptr || parent.id != id || parent.data->size() != parent_size ||
+ sample_hash(*parent.data) != parent_hash) {
+ error = "its prefix is not in the cache";
+ return false;
+ }
+ const uint8_t * p = parent.data->data();
+ out.clear();
+ out.reserve(size);
+ for (uint32_t i = 0; i < n_ops; ++i) {
+ uint8_t op;
+ uint64_t a, len;
+ if (!r.get(op) || !r.get(a)) {
+ error = "truncated";
+ return false;
+ }
+ if (op == OP_COPY) {
+ if (!r.get(len) || a > parent_size || len > parent_size - a) {
+ error = "copy out of range";
+ return false;
+ }
+ out.insert(out.end(), p + a, p + a + len);
+ } else if (op == OP_LITERAL) {
+ const size_t at = r.at;
+ if (!r.skip(a)) {
+ error = "truncated literal";
+ return false;
+ }
+ out.insert(out.end(), delta.data() + at, delta.data() + at + a);
+ } else {
+ error = "unknown op";
+ return false;
+ }
+ }
+ if (out.size() != size || r.at != delta.size()) {
+ error = "size mismatch";
+ return false;
+ }
+ return true;
+}
diff --git a/tools/server/server-kv-share.h b/tools/server/server-kv-share.h
new file mode 100644
index 00000000..f01e0fec
--- /dev/null
+++ b/tools/server/server-kv-share.h
@@ -0,0 +1,37 @@
+#pragma once
+
+// prefix sharing for prompt cache states: a sequence state (llama_state_seq_get_data) whose KV rows repeat a parent
+// state's rows is stored as a delta, a list of copies from the parent and literal bytes. decoding needs no knowledge
+// of the layout; encoding finds the parent's row for each of the child's KV cells (same position) and copies it only
+// where the bytes are equal, so a decoded delta is the original state byte for byte.
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+struct kv_share_parent {
+ uint64_t id = 0;
+ const std::vector<uint8_t> * data = nullptr;
+};
+
+// the cell metadata carries an extra 2D position (8 bytes) for M-RoPE models
+struct kv_share_layout {
+ size_t cell_ext = 0;
+};
+
+// true if `blob` is a delta written by kv_share_encode
+bool kv_share_is_delta(const std::vector<uint8_t> & blob);
+
+// the parent id a delta needs, 0 if `blob` is not a delta
+uint64_t kv_share_parent_id(const std::vector<uint8_t> & blob);
+
+// a delta of `child` against `parent`; false when the layout is not understood or it would save less than 16 MiB.
+// n_shared receives the bytes copied from the parent
+bool kv_share_encode(const std::vector<uint8_t> & child, const kv_share_parent & parent, const kv_share_layout & layout,
+ std::vector<uint8_t> & delta, size_t & n_shared);
+
+// rebuilds the original state from a delta and its parent; false (with a reason) if the parent is not the one the
+// delta was written against or the delta is malformed
+bool kv_share_decode(const std::vector<uint8_t> & delta, const kv_share_parent & parent, std::vector<uint8_t> & out,
+ std::string & error);
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index b2cf0f40..d6ccd466 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1753,6 +1753,7 @@ constexpr uint32_t DISK_VERSION_MEDIA = 3;
constexpr size_t META_OFFSET = 16;
constexpr size_t META_SIZE = 512;
constexpr size_t META_FIXED = 74;
+constexpr size_t META_PARENT = META_SIZE - 8; // the learned prefix a shared state needs
constexpr size_t FAMILY_TOKENS = 256; // cap on the system block that is hashed
constexpr size_t FAMILY_USER_TOKENS = 32; // without a system block: the start of the first user turn
@@ -1939,7 +1940,9 @@ std::vector<uint8_t> meta_encode(const cache_meta & m) {
put_le<uint64_t>(b, 48, m.hits);
put_le<uint64_t>(b, 56, m.family);
put_le<int64_t> (b, 64, m.n_input);
- const std::string ex = clean_utf8(m.excerpt, META_SIZE - META_FIXED);
+ put_le<uint8_t> (b, 12, m.parent != 0);
+ put_le<uint64_t>(b, META_PARENT, m.parent);
+ const std::string ex = clean_utf8(m.excerpt, META_PARENT - META_FIXED);
put_le<uint16_t>(b, 72, (uint16_t) ex.size());
memcpy(b.data() + META_FIXED, ex.data(), ex.size());
return b;
@@ -1961,7 +1964,10 @@ bool meta_decode(const std::vector<uint8_t> & b, cache_meta & m) {
m.hits = get_le<uint64_t>(b, 48);
m.family = get_le<uint64_t>(b, 56);
m.n_input = get_le<int64_t>(b, 64);
- const size_t n = std::min<size_t>(get_le<uint16_t>(b, 72), META_SIZE - META_FIXED);
+ // written before entries could share a prefix, the excerpt may run to the end of the block
+ const bool shared = get_le<uint8_t>(b, 12) != 0;
+ m.parent = shared ? get_le<uint64_t>(b, META_PARENT) : 0;
+ const size_t n = std::min<size_t>(get_le<uint16_t>(b, 72), (shared ? META_PARENT : META_SIZE) - META_FIXED);
m.excerpt.assign((const char *) b.data() + META_FIXED, n);
return true;
}
@@ -2156,6 +2162,7 @@ json entry_json(const cache_meta & m, size_t n_tokens, size_t n_images, size_t b
{ "expires_in", m.expires_at == 0 ? json(nullptr) : json(m.expires_at - now) },
{ "family", hex_id(m.family) },
{ "excerpt", m.excerpt },
+ { "shared_from", m.parent == 0 ? json(nullptr) : json(hex_id(m.parent)) },
};
}
@@ -2261,6 +2268,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
}
cache_meta meta = prompt.meta;
+ meta.parent = 0; // the new state is stored whole until share() finds its prefix
// remove any cached prompts that are fully contained in the current prompt
for (auto it = states.begin(); it != states.end();) {
@@ -2503,6 +2511,62 @@ std::string server_prompt_cache::disk_best(const std::function<int32_t(const ser
return best;
}
+const server_prompt_cache_state * server_prompt_cache::learned_state(uint64_t id) const {
+ for (const auto & s : states) {
+ if (s.prompt.meta.id == id && is_learned(s.prompt.meta) && !s.data.main.empty()) {
+ return &s;
+ }
+ }
+ return nullptr;
+}
+
+bool server_prompt_cache::unshare(const std::vector<uint8_t> & data, bool draft, std::vector<uint8_t> & out, std::string & err) const {
+ if (!kv_share_is_delta(data)) {
+ return true;
+ }
+ const uint64_t id = kv_share_parent_id(data);
+ const auto * parent = learned_state(id);
+ if (parent == nullptr) {
+ err = "its learned prefix " + hex_id(id) + " is gone";
+ return false;
+ }
+ return kv_share_decode(data, { id, draft ? &parent->data.drft : &parent->data.main }, out, err);
+}
+
+void server_prompt_cache::share(server_prompt_cache_state & state) {
+ if (is_learned(state.prompt.meta) || state.data.main.empty() || kv_share_is_delta(state.data.main)) {
+ return;
+ }
+ const server_prompt_cache_state * parent = nullptr;
+ for (const auto & s : states) {
+ if (&s == &state || !is_learned(s.prompt.meta) || s.data.main.empty() || s.prompt.tokens.size() >= state.prompt.tokens.size()) {
+ continue;
+ }
+ if (is_prefix(s.prompt.tokens, state.prompt.tokens) && (parent == nullptr || s.prompt.tokens.size() > parent->prompt.tokens.size())) {
+ parent = &s;
+ }
+ }
+ if (parent == nullptr) {
+ return;
+ }
+ const int64_t t_start = ggml_time_us();
+ const uint64_t id = parent->prompt.meta.id;
+ const size_t before = state.data.size();
+ std::vector<uint8_t> delta;
+ size_t n_shared = 0;
+ if (!kv_share_encode(state.data.main, { id, &parent->data.main }, share_layout, delta, n_shared)) {
+ return;
+ }
+ state.data.main = std::move(delta);
+ if (kv_share_encode(state.data.drft, { id, &parent->data.drft }, share_layout, delta, n_shared)) {
+ state.data.drft = std::move(delta);
+ }
+ state.prompt.meta.parent = id;
+ SRV_INF("prompt cache: %d tokens stored against the learned %d-token prefix %s, %.1f MiB instead of %.1f MiB (%.2f s)\n",
+ state.prompt.n_tokens(), parent->prompt.n_tokens(), hex_id(id).c_str(), state.data.size() / (1024.0*1024.0),
+ before / (1024.0*1024.0), (ggml_time_us() - t_start) / 1e6);
+}
+
bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::iterator it, server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, bool keep_entry) {
const bool learned = is_learned(it->prompt.meta);
const bool consume = !keep_entry && !learned;
@@ -2515,8 +2579,22 @@ bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::itera
return false;
}
+ // a shared state is rebuilt from its learned prefix; one whose prefix is gone can never be restored
+ std::vector<uint8_t> main_full, drft_full;
+ {
+ std::string err;
+ const bool ok = unshare(it->data.main, false, main_full, err) && unshare(it->data.drft, true, drft_full, err);
+ if (!ok) {
+ SRV_WRN("prompt cache: dropping entry %s (%d tokens): %s\n", hex_id(it->prompt.meta.id).c_str(), it->prompt.n_tokens(), err.c_str());
+ delete_entry(it->prompt.meta.id);
+ return false;
+ }
+ }
+ const auto & state_main = main_full.empty() ? it->data.main : main_full;
+ const auto & state_drft = drft_full.empty() ? it->data.drft : drft_full;
+
{
- const auto & data = it->data.main;
+ const auto & data = state_main;
const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), data.size(), id_slot, 0);
if (n != data.size()) {
SRV_ERR("failed to restore state with size %zu\n", data.size());
@@ -2525,7 +2603,7 @@ bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::itera
}
{
- const auto & data = it->data.drft;
+ const auto & data = state_drft;
if (!data.empty()) {
GGML_ASSERT(ctx_dft);
@@ -2545,6 +2623,7 @@ bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::itera
} else {
prompt = it->prompt.clone();
}
+ prompt.meta.parent = 0; // the slot holds the whole state
if (learned) {
it->prompt.meta.hits++;
it->prompt.meta.last_used = unix_now();
@@ -2979,6 +3058,32 @@ bool server_prompt_cache::set_entry(uint64_t id, cache_class cls, int64_t ttl) {
}
bool server_prompt_cache::delete_entry(uint64_t id) {
+ // the entries stored against it are rebuilt whole while it is still here (RAM), or go with it (disk)
+ for (auto & s : states) {
+ if (s.prompt.meta.parent != id) {
+ continue;
+ }
+ std::vector<uint8_t> main, drft;
+ std::string err;
+ if (unshare(s.data.main, false, main, err) && unshare(s.data.drft, true, drft, err)) {
+ if (!main.empty()) {
+ s.data.main = std::move(main);
+ }
+ if (!drft.empty()) {
+ s.data.drft = std::move(drft);
+ }
+ s.prompt.meta.parent = 0;
+ }
+ }
+ std::vector<uint64_t> orphans;
+ if (disk_thread.joinable()) {
+ std::lock_guard<std::mutex> lock(disk_mutex);
+ for (const auto & e : disk) {
+ if (e.meta.parent == id) {
+ orphans.push_back(e.meta.id);
+ }
+ }
+ }
bool found = false;
for (auto it = states.begin(); it != states.end();) {
if (it->prompt.meta.id == id) {
@@ -3002,6 +3107,9 @@ bool server_prompt_cache::delete_entry(uint64_t id) {
}
}
forget_classify(id);
+ for (uint64_t child : orphans) {
+ delete_entry(child);
+ }
update_ram_next_expiry();
return found;
}
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 2f8c4032..4665ad27 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -2,6 +2,7 @@
#include "common.h"
#include "llama.h"
+#include "server-kv-share.h"
#include <string>
#include <unordered_set>
@@ -641,6 +642,7 @@ struct cache_meta {
bool family_saved = false; // counted in its family's saved/reused stats
bool family_reused = false;
std::string excerpt; // start of the first user message
+ uint64_t parent = 0; // the learned prefix this entry's state is stored against (0 = stored whole)
bool classified() const {
return source != CACHE_SOURCE_DEFAULT;
@@ -739,6 +741,20 @@ struct server_prompt_cache {
server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
+ // the layout of a state's KV cell metadata (M-RoPE models carry an extra position), set from the model
+ kv_share_layout share_layout;
+
+ // stores a state that continues a learned prefix as a delta against it: only the KV rows that differ from the
+ // prefix's and the rest of the state (the recurrent part) are kept
+ void share(server_prompt_cache_state & state);
+
+ // the learned entry with this id that holds a state, or null
+ const server_prompt_cache_state * learned_state(uint64_t id) const;
+
+ // a shared state (main, or the draft's) rebuilt whole into out; out stays empty for a state stored whole.
+ // false if its learned prefix is gone
+ bool unshare(const std::vector<uint8_t> & data, bool draft, std::vector<uint8_t> & out, std::string & err) const;
+
// 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);