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(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 +#include + +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 & 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(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 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 keys; + std::vector 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 & blob, const kv_share_layout & layout, std::vector & 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(blob.data() + key_at), sizeof(pos)); + key.append(reinterpret_cast(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 & 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 void put(const T & v) { + const auto * b = reinterpret_cast(&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 & 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 & 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 & child, const kv_share_parent & parent, const kv_share_layout & layout, + std::vector & 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 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 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 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 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 & delta, const kv_share_parent & parent, std::vector & 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 +#include +#include +#include + +struct kv_share_parent { + uint64_t id = 0; + const std::vector * 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 & blob); + +// the parent id a delta needs, 0 if `blob` is not a delta +uint64_t kv_share_parent_id(const std::vector & 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 & child, const kv_share_parent & parent, const kv_share_layout & layout, + std::vector & 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 & delta, const kv_share_parent & parent, std::vector & 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 meta_encode(const cache_meta & m) { put_le(b, 48, m.hits); put_le(b, 56, m.family); put_le (b, 64, m.n_input); - const std::string ex = clean_utf8(m.excerpt, META_SIZE - META_FIXED); + put_le (b, 12, m.parent != 0); + put_le(b, META_PARENT, m.parent); + const std::string ex = clean_utf8(m.excerpt, META_PARENT - META_FIXED); put_le(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 & b, cache_meta & m) { m.hits = get_le(b, 48); m.family = get_le(b, 56); m.n_input = get_le(b, 64); - const size_t n = std::min(get_le(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(b, 12) != 0; + m.parent = shared ? get_le(b, META_PARENT) : 0; + const size_t n = std::min(get_le(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 & data, bool draft, std::vector & 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 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::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::itera return false; } + // a shared state is rebuilt from its learned prefix; one whose prefix is gone can never be restored + std::vector 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::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::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 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 orphans; + if (disk_thread.joinable()) { + std::lock_guard 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 #include @@ -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 & data, bool draft, std::vector & 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);