diff --git a/common/arg.cpp b/common/arg.cpp index 6d5edcb7f2d5..1b42c7fa6cc0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1736,6 +1736,47 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_ram_mib = value; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-kvcd", "--kv-cache-dir"}, "PATH", + "directory holding the on-disk KV cache. Conversations are fingerprinted and their KV state is\n" + "persisted there, so a returning conversation is restored instead of re-prefilled (default: disabled)", + [](common_params & params, const std::string & value) { + params.kv_cache_dir = value; + } + ).set_env("LLAMA_ARG_KV_CACHE_DIR").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--kv-cache-max"}, "N", + string_format("maximum size of the on-disk KV cache in MiB, least-recently-used entries are evicted " + "(default: %d, 0 = no limit)", params.kv_cache_max_mib), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("kv-cache-max must be non-negative"); + } + params.kv_cache_max_mib = value; + } + ).set_env("LLAMA_ARG_KV_CACHE_MAX").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--kv-cache-min-tokens"}, "N", + string_format("do not persist prompts shorter than this to the on-disk KV cache (default: %d)", + params.kv_cache_min_toks), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("kv-cache-min-tokens must be non-negative"); + } + params.kv_cache_min_toks = value; + } + ).set_env("LLAMA_ARG_KV_CACHE_MIN_TOKENS").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--kv-cache-ttl"}, "N", + string_format("drop on-disk KV cache entries not used for N seconds, checked before every " + "lookup (default: %d = 3 days, 0 = never)", params.kv_cache_ttl_s), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("kv-cache-ttl must be non-negative"); + } + params.kv_cache_ttl_s = value; + } + ).set_env("LLAMA_ARG_KV_CACHE_TTL").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 13bebcb52c54..0f729489f7f9 100644 --- a/common/common.h +++ b/common/common.h @@ -657,6 +657,13 @@ struct common_params { int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + // on-disk KV cache: conversations are fingerprinted (sha1 of model tag + prompt tokens) and their + // KV state persisted, so a returning conversation is restored from disk instead of re-prefilled + std::string kv_cache_dir = ""; // "" = disabled + int32_t kv_cache_max_mib = 0; // LRU budget on disk, 0 = no limit + int32_t kv_cache_min_toks = 256; // do not persist prompts shorter than this + int32_t kv_cache_ttl_s = 259200; // drop entries older than this (3 days), 0 = never + std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT std::string api_prefix = ""; // NOLINT diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 280bd9e19dca..3036ca7d33de 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -8,6 +8,7 @@ add_library(${TARGET} STATIC server-chat.cpp server-chat.h server-task.cpp + server-kv-disk.cpp server-task.h server-queue.cpp server-queue.h diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f41650ecb658..3d777bb06a52 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1,6 +1,7 @@ #include "server-context.h" #include "server-chat.h" #include "server-common.h" +#include "server-kv-disk.h" #include "server-http.h" #include "server-task.h" #include "server-queue.h" @@ -307,6 +308,10 @@ struct server_slot { return false; } + if (!prompt_cache.ram_enabled) { + return false; + } + const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; @@ -337,6 +342,10 @@ struct server_slot { return res; } + // conversation key whose KV this slot currently holds, so a returning turn that lands on + // the same slot is not re-read from disk over the top of an identical copy [TAG_KV_DISK_RESIDENT] + std::string kv_conv_key; + void prompt_clear() { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); @@ -346,6 +355,8 @@ struct server_slot { } prompt.clear(); + + kv_conv_key.clear(); } std::vector lora; @@ -1547,11 +1558,15 @@ struct server_context_impl { batch.init(std::max(n_batch, params_base.n_parallel), n_embd); } - if (params_base.cache_ram_mib != 0) { + // the on-disk tier works on its own, so the cache object is also needed when the RAM tier + // is switched off with `--cache-ram 0` + if (params_base.cache_ram_mib != 0 || !params_base.kv_cache_dir.empty()) { if (params_base.cache_ram_mib < 0) { SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); - } else { + } else if (params_base.cache_ram_mib > 0) { SRV_TRC("prompt cache is enabled, size limit: %d MiB\n", params_base.cache_ram_mib); + } else { + SRV_TRC("%s", "in-memory prompt cache is disabled\n"); } SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n"); @@ -1559,6 +1574,25 @@ struct server_context_impl { } else { SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); } + + if (prompt_cache && !params_base.kv_cache_dir.empty()) { + // The fingerprint is bound to the model, so two models sharing a directory never collide. + // Keyed on the weights rather than on --alias: renaming an alias must not invalidate a + // cache, and two models sharing one must not be able to read each other's state. + char buf_desc[128] = {0}; + llama_model_desc(model_tgt, buf_desc, sizeof(buf_desc)); + + const std::string model_tag = params_base.model.path + "|" + buf_desc; + + prompt_cache->disk_min_tokens = (size_t) params_base.kv_cache_min_toks; + + if (!prompt_cache->disk_init(params_base.kv_cache_dir, model_tag, + 1024ull*1024ull*(size_t) params_base.kv_cache_max_mib, + (int64_t) params_base.kv_cache_ttl_s, + mctx != nullptr)) { + SRV_WRN("%s", "on-disk KV cache could not be initialised, continuing without it\n"); + } + } SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); if (params_base.n_ctx_checkpoints > 0) { @@ -1617,8 +1651,8 @@ struct server_context_impl { metrics.init(); if (params_base.cache_idle_slots) { - if (params_base.cache_ram_mib == 0) { - SRV_WRN("%s", "--cache-idle-slots requires --cache-ram, disabling\n"); + if (params_base.cache_ram_mib == 0 && params_base.kv_cache_dir.empty()) { + SRV_WRN("%s", "--cache-idle-slots requires --cache-ram or --kv-cache-dir, disabling\n"); params_base.cache_idle_slots = false; } else { if (params_base.kv_unified) { @@ -1844,6 +1878,23 @@ struct server_context_impl { } } + // On-disk KV cache. The key is an exact sha1 of the conversation truncated at its last + // assistant message, so a hit means this slot now holds precisely the state the previous + // turn ended with, and only the new user message has to be prefilled. [TAG_KV_DISK_LOAD] + if (ret && prompt_cache && prompt_cache->disk_enabled() && + task.type == SERVER_TASK_TYPE_COMPLETION && !task.params.kv_conv_key.empty()) { + if (ret->kv_conv_key == task.params.kv_conv_key) { + // the slot already holds exactly this conversation. Reading it back would install a + // byte-identical copy at the cost of a multi-hundred-MiB read and would throw away + // the checkpoints built since. Let the ordinary prefix logic reuse what is resident. + SLT_DBG(*ret, "%s", "conversation already resident, skipping disk restore\n"); + } else if (prompt_cache->disk_load(ret->prompt, task.params.kv_conv_key, ret->ctx_tgt, ret->ctx_dft, ret->id)) { + ret->kv_conv_key = task.params.kv_conv_key; + } else { + ret->kv_conv_key.clear(); + } + } + return ret; } @@ -2275,6 +2326,35 @@ struct server_context_impl { } void send_final_response(server_slot & slot) { + // Persist the finished conversation while its KV is still in the slot, and before + // generated_text is moved into the response below. The key is this turn's conversation + // extended with the reply, which is exactly what the next turn will look itself up by. + // [TAG_KV_DISK_SAVE] + if (prompt_cache && prompt_cache->disk_enabled() && slot.task && + !slot.task->is_child() && slot.stats.n_gen > 0 && + !slot.task->params.kv_conv_base.empty()) { + sha1_state base; + + if (sha1_state_from_hex(slot.task->params.kv_conv_base, base)) { + // hash the reply as the client will echo it back, i.e. without reasoning + std::string reply = slot.generated_text; + try { + reply = common_chat_parse(slot.generated_text, false, slot.task->params.chat_parser_params).content; + } catch (const std::exception &) { + // not a chat request, or unparseable: the raw text is the best key available + } + + server_kv_add_message(base, "assistant", reply); + + const std::string key_store = sha1_hex(base); + + prompt_cache->disk_store(slot.prompt, key_store, slot.ctx_tgt, slot.ctx_dft, slot.id); + + // the next turn of this conversation looks itself up by exactly this key + slot.kv_conv_key = key_store; + } + } + auto res = std::make_unique(); res->id = slot.task->id; @@ -4435,6 +4515,55 @@ void server_context::set_state_callback(server_state_callback_t callback) { // server_routes // +// Build the on-disk KV cache keys for a chat request. +// +// The lookup key hashes the conversation truncated at its last assistant message, because that is +// byte-identical to what the previous turn stored. Hashing the whole arriving conversation would +// never hit: every request carries one more message than anything already saved. +// +// The base digest covers all arriving messages and is finished off with the reply once it exists, +// producing the key this turn stores under - which is what the next turn will look up. +static void kv_attach_conv_keys(const json & body, json & data) { + if (!body.contains("messages") || !body.at("messages").is_array()) { + return; + } + + const auto & msgs = body.at("messages"); + + int i_last_assistant = -1; + for (size_t i = 0; i < msgs.size(); i++) { + if (json_value(msgs[i], "role", std::string()) == "assistant") { + i_last_assistant = (int) i; + } + } + + sha1_state s; + std::string key_load; + + for (size_t i = 0; i < msgs.size(); i++) { + const std::string role = json_value(msgs[i], "role", std::string()); + + std::string content; + if (msgs[i].contains("content") && !msgs[i].at("content").is_null()) { + const auto & c = msgs[i].at("content"); + content = c.is_string() ? c.get() : c.dump(); + } + + server_kv_add_message(s, role, content); + + if ((int) i == i_last_assistant) { + key_load = sha1_hex(s); + } + } + + // no assistant message yet means a first turn: nothing can have been stored for it + if (!key_load.empty()) { + data["__kv_conv_key"] = key_load; + } + + data["__kv_conv_base"] = sha1_state_to_hex(s); +} + std::unique_ptr server_routes::handle_completions_impl( const server_http_req & req, server_task_type type, @@ -4504,6 +4633,10 @@ std::unique_ptr server_routes::handle_completions_impl( task.id_slot = json_value(data, "id_slot", -1); sse_ping_interval = task.params.sse_ping_interval; + // on-disk KV cache keys, computed by the chat route where the messages are still visible + task.params.kv_conv_key = json_value(data, "__kv_conv_key", std::string()); + task.params.kv_conv_base = json_value(data, "__kv_conv_base", std::string()); + // OAI-compat task.params.res_type = res_type; task.params.oaicompat_cmpl_id = completion_id; @@ -5116,6 +5249,7 @@ void server_routes::init_routes() { body, meta->chat_params, files); + kv_attach_conv_keys(body, body_parsed); return handle_completions_impl( req, SERVER_TASK_TYPE_COMPLETION, diff --git a/tools/server/server-kv-disk.cpp b/tools/server/server-kv-disk.cpp new file mode 100644 index 000000000000..d322c56b1de0 --- /dev/null +++ b/tools/server/server-kv-disk.cpp @@ -0,0 +1,682 @@ +// On-disk KV cache (--kv-cache-dir). See server-kv-disk.h for the keying scheme. +// +// Prefill on this class of machine costs minutes where the same KV state reads back from NVMe in +// seconds, so it is worth spending disk to never prefill a conversation twice. Measured on +// gemma-4-26B-A4B at q8_0 KV: 15.4 KiB/token, so a 22.9k-token conversation is 344 MiB that writes +// in 45 ms and reads in 47 ms, against 18.4 s to re-prefill it. +// +// An entry is three files: +// +// .kv target context state (llama_state_seq_save_file, carries the token list) +// .dft draft context state, when speculative decoding is on +// .idx index record, written last and deleted first, so it is the commit marker +// +// llama_state_seq_save_file / _load_file are the same primitives the /slots endpoints expose, but +// nothing here goes through HTTP: they are called in-process, on the inference thread, driven by +// the conversation key rather than by hand. + +#include "server-kv-disk.h" +#include "server-task.h" + +#include "common.h" +#include "llama.h" +#include "server-common.h" + +#include +#include +#include +#include +#include +#include +#include + +// +// sha1 +// + +namespace { + +inline uint32_t sha1_rol(uint32_t v, int b) { + return (v << b) | (v >> (32 - b)); +} + +void sha1_compress(sha1_state & s, const uint8_t * p) { + uint32_t w[80]; + + for (int i = 0; i < 16; i++) { + w[i] = (uint32_t(p[4*i + 0]) << 24) | (uint32_t(p[4*i + 1]) << 16) | + (uint32_t(p[4*i + 2]) << 8) | (uint32_t(p[4*i + 3])); + } + for (int i = 16; i < 80; i++) { + w[i] = sha1_rol(w[i-3] ^ w[i-8] ^ w[i-14] ^ w[i-16], 1); + } + + uint32_t a = s.h[0], b = s.h[1], c = s.h[2], d = s.h[3], e = s.h[4]; + + for (int i = 0; i < 80; i++) { + uint32_t f, k; + + if (i < 20) { + f = (b & c) | ((~b) & d); k = 0x5A827999; + } else if (i < 40) { + f = b ^ c ^ d; k = 0x6ED9EBA1; + } else if (i < 60) { + f = (b & c) | (b & d) | (c & d); k = 0x8F1BBCDC; + } else { + f = b ^ c ^ d; k = 0xCA62C1D6; + } + + const uint32_t t = sha1_rol(a, 5) + f + e + k + w[i]; + + e = d; d = c; c = sha1_rol(b, 30); b = a; a = t; + } + + s.h[0] += a; s.h[1] += b; s.h[2] += c; s.h[3] += d; s.h[4] += e; +} + +const char * HEX = "0123456789abcdef"; + +} // namespace + +void sha1_update(sha1_state & s, const void * data, size_t n) { + const uint8_t * p = (const uint8_t *) data; + + s.n_len += n; + + while (n > 0) { + const size_t take = std::min(n, sizeof(s.buf) - s.n_buf); + + memcpy(s.buf + s.n_buf, p, take); + + s.n_buf += (uint32_t) take; + p += take; + n -= take; + + if (s.n_buf == sizeof(s.buf)) { + sha1_compress(s, s.buf); + s.n_buf = 0; + } + } +} + +std::string sha1_hex(sha1_state s) { + const uint64_t n_bits = s.n_len * 8; + + const uint8_t pad = 0x80; + sha1_update(s, &pad, 1); + + const uint8_t zero = 0x00; + while (s.n_buf != 56) { + sha1_update(s, &zero, 1); + } + + uint8_t len[8]; + for (int i = 0; i < 8; i++) { + len[i] = (uint8_t) (n_bits >> (56 - 8*i)); + } + sha1_update(s, len, 8); + + char out[41]; + for (int i = 0; i < 5; i++) { + snprintf(out + 8*i, 9, "%08x", s.h[i]); + } + + return std::string(out, 40); +} + +std::string sha1_state_to_hex(const sha1_state & s) { + uint8_t raw[sizeof(sha1_state)]; + memcpy(raw, &s, sizeof(s)); + + std::string out; + out.reserve(sizeof(raw) * 2); + + for (size_t i = 0; i < sizeof(raw); i++) { + out += HEX[raw[i] >> 4]; + out += HEX[raw[i] & 0xf]; + } + + return out; +} + +bool sha1_state_from_hex(const std::string & hex, sha1_state & s) { + if (hex.size() != sizeof(sha1_state) * 2) { + return false; + } + + uint8_t raw[sizeof(sha1_state)]; + + for (size_t i = 0; i < sizeof(raw); i++) { + const char hi = hex[2*i], lo = hex[2*i + 1]; + + const auto nib = [](char ch) -> int { + if (ch >= '0' && ch <= '9') return ch - '0'; + if (ch >= 'a' && ch <= 'f') return ch - 'a' + 10; + return -1; + }; + + const int a = nib(hi), b = nib(lo); + if (a < 0 || b < 0) { + return false; + } + + raw[i] = (uint8_t) ((a << 4) | b); + } + + memcpy(&s, raw, sizeof(s)); + + if (s.n_buf >= sizeof(s.buf)) { + return false; + } + + return true; +} + +void server_kv_add_message(sha1_state & s, const std::string & role, const std::string & content) { + sha1_update(s, role.data(), role.size()); + sha1_update(s, "\0", 1); + sha1_update(s, content.data(), content.size()); + sha1_update(s, "\x01", 1); +} + +// +// on-disk store +// + +namespace { + +constexpr char KV_IDX_MAGIC[8] = { 'K','V','C','A','C','H','E','2' }; +constexpr uint32_t KV_IDX_VERSION = 2; + +int64_t kv_now() { + return std::filesystem::file_time_type::clock::now().time_since_epoch().count(); +} + +// file_time_type ticks are implementation defined; derive the tick rate rather than assume it +int64_t kv_ticks_per_sec() { + using clock = std::filesystem::file_time_type::clock; + return std::chrono::duration_cast(std::chrono::seconds(1)).count(); +} + +bool kv_write_idx(const std::string & path, uint64_t n_packed, uint64_t bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + if (!f) { + return false; + } + + f.write(KV_IDX_MAGIC, sizeof(KV_IDX_MAGIC)); + f.write((const char *) &KV_IDX_VERSION, sizeof(KV_IDX_VERSION)); + f.write((const char *) &n_packed, sizeof(n_packed)); + f.write((const char *) &bytes, sizeof(bytes)); + f.close(); + + return f.good(); +} + +bool kv_read_idx(const std::string & path, uint64_t & n_packed, uint64_t & bytes) { + std::ifstream f(path, std::ios::binary); + if (!f) { + return false; + } + + char magic[8]; + uint32_t version = 0; + + f.read(magic, sizeof(magic)); + f.read((char *) &version, sizeof(version)); + f.read((char *) &n_packed, sizeof(n_packed)); + f.read((char *) &bytes, sizeof(bytes)); + + if (!f || memcmp(magic, KV_IDX_MAGIC, sizeof(magic)) != 0 || version != KV_IDX_VERSION) { + return false; + } + + return n_packed <= (1ull << 32); +} + +llama_tokens kv_packed_to_tokens(const std::vector & packed) { + llama_tokens out(packed.size() / sizeof(llama_token)); + + if (!out.empty()) { + memcpy(out.data(), packed.data(), out.size() * sizeof(llama_token)); + } + + return out; +} + +// Context checkpoints have to travel with the state. +// +// On a sliding-window model the sequence state only carries the SWA window, so after a restore +// llama_memory_seq_pos_min sits above pos_min_thold and the server will only reuse the prompt if a +// checkpoint reaches further back. Without these the whole restored prefix is discarded and the +// conversation is re-prefilled - which is exactly the state the in-memory tier avoids by copying +// prompt.checkpoints alongside the blob. [TAG_KV_DISK_CKPT] +void kv_write_vec(std::ofstream & f, const std::vector & v) { + const uint64_t n = v.size(); + f.write((const char *) &n, sizeof(n)); + if (n) { + f.write((const char *) v.data(), n); + } +} + +bool kv_read_vec(std::ifstream & f, std::vector & v) { + uint64_t n = 0; + f.read((char *) &n, sizeof(n)); + if (!f || n > (1ull << 34)) { + return false; + } + v.resize(n); + if (n) { + f.read((char *) v.data(), n); + } + return (bool) f; +} + +uint64_t kv_write_checkpoints(const std::string & path, const std::list & ckpts) { + if (ckpts.empty()) { + return 0; + } + + std::ofstream f(path, std::ios::binary | std::ios::trunc); + if (!f) { + return 0; + } + + const uint32_t n = (uint32_t) ckpts.size(); + f.write((const char *) &n, sizeof(n)); + + for (const auto & c : ckpts) { + f.write((const char *) &c.n_tokens, sizeof(c.n_tokens)); + f.write((const char *) &c.pos_min, sizeof(c.pos_min)); + f.write((const char *) &c.pos_max, sizeof(c.pos_max)); + + kv_write_vec(f, c.data_tgt); + kv_write_vec(f, c.data_dft); + kv_write_vec(f, c.data_spec); + } + + const uint64_t bytes = (uint64_t) f.tellp(); + f.close(); + + return f.good() ? bytes : 0; +} + +bool kv_read_checkpoints(const std::string & path, std::list & out) { + std::ifstream f(path, std::ios::binary); + if (!f) { + return false; // no checkpoints stored, not an error + } + + uint32_t n = 0; + f.read((char *) &n, sizeof(n)); + if (!f || n > 1024) { + return false; + } + + for (uint32_t i = 0; i < n; i++) { + common_prompt_checkpoint c; + + f.read((char *) &c.n_tokens, sizeof(c.n_tokens)); + f.read((char *) &c.pos_min, sizeof(c.pos_min)); + f.read((char *) &c.pos_max, sizeof(c.pos_max)); + + if (!f || !kv_read_vec(f, c.data_tgt) || !kv_read_vec(f, c.data_dft) || !kv_read_vec(f, c.data_spec)) { + out.clear(); + return false; + } + + // the task that made it is long gone; -1 keeps the eviction rule in create_checkpoint honest + c.id_task = -1; + + out.push_back(std::move(c)); + } + + return true; +} + +// a failed restore must leave the slot genuinely empty in every context, otherwise the next request +// would measure a prefix against cells that were never installed +void kv_reset_slot(server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { + llama_memory_seq_rm(llama_get_memory(ctx_tgt), id_slot, -1, -1); + + if (ctx_dft) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), id_slot, -1, -1); + } + + prompt.clear(); +} + +} // namespace + +std::string server_prompt_cache::disk_key(const std::string & conv_key) const { + // bind the conversation to the weights, so two models sharing a directory cannot collide + sha1_state s; + + sha1_update(s, model_tag.data(), model_tag.size()); + sha1_update(s, "\0", 1); + sha1_update(s, conv_key.data(), conv_key.size()); + + return sha1_hex(s); +} + +std::string server_prompt_cache::disk_path(const std::string & key) const { + return (std::filesystem::path(dir) / key).string(); +} + +size_t server_prompt_cache::disk_size() const { + size_t res = 0; + + for (const auto & [_, e] : disk) { + res += e.bytes; + } + + return res; +} + +void server_prompt_cache::disk_erase(const std::string & key) { + std::error_code ec; + + // .idx first: an entry stops existing the moment its commit marker is gone + std::filesystem::remove(disk_path(key) + ".idx", ec); + std::filesystem::remove(disk_path(key) + ".kv", ec); + std::filesystem::remove(disk_path(key) + ".dft", ec); + std::filesystem::remove(disk_path(key) + ".ckpt", ec); + + disk.erase(key); +} + +bool server_prompt_cache::disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, int64_t ttl_s, bool has_mtmd) { + this->dir = dir; + this->model_tag = model_tag; + this->disk_limit_size = limit_size; + this->disk_ttl_s = ttl_s; + this->disk_has_mtmd = has_mtmd; + + std::error_code ec; + + std::filesystem::create_directories(dir, ec); + if (ec) { + SRV_WRN("kv cache: cannot create directory '%s': %s\n", dir.c_str(), ec.message().c_str()); + this->dir.clear(); + return false; + } + + std::vector orphans; + + for (const auto & de : std::filesystem::directory_iterator(dir, ec)) { + const auto & path = de.path(); + + if (path.extension() == ".tmp") { + orphans.push_back(path); // an interrupted write + continue; + } + + if (path.extension() != ".idx") { + continue; + } + + const std::string key = path.stem().string(); + + uint64_t n_packed = 0; + uint64_t bytes = 0; + + if (!kv_read_idx(path.string(), n_packed, bytes) || + !std::filesystem::exists(disk_path(key) + ".kv", ec)) { + SRV_WRN("kv cache: dropping incomplete entry %s\n", key.substr(0, 12).c_str()); + orphans.push_back(path); + continue; + } + + server_prompt_cache_disk_entry entry; + + entry.n_packed = n_packed; + entry.bytes = bytes; + + const auto mtime = std::filesystem::last_write_time(path, ec); + entry.t_last = ec ? kv_now() : mtime.time_since_epoch().count(); + + disk.emplace(key, entry); + } + + for (const auto & path : orphans) { + const std::string base = (path.parent_path() / path.stem()).string(); + + std::filesystem::remove(path, ec); + std::filesystem::remove(base + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + std::filesystem::remove(base + ".ckpt", ec); + } + + SRV_INF("kv cache: '%s' holds %zu conversation(s), %.3f GiB (limit %.3f GiB, ttl %lld h)\n", + dir.c_str(), disk.size(), disk_size() / (1024.0*1024.0*1024.0), + disk_limit_size / (1024.0*1024.0*1024.0), (long long) (disk_ttl_s / 3600)); + + disk_prune(); + + return true; +} + +void server_prompt_cache::disk_prune() { + if (dir.empty()) { + return; + } + + // age first: a stale conversation is dropped whether or not the store is over its limit + if (disk_ttl_s > 0) { + const int64_t cutoff = kv_now() - disk_ttl_s * kv_ticks_per_sec(); + + for (auto it = disk.begin(); it != disk.end();) { + if (it->second.t_last < cutoff) { + SRV_INF("kv cache: expiring %s (%.3f GiB, older than %lld h)\n", + it->first.substr(0, 12).c_str(), it->second.bytes / (1024.0*1024.0*1024.0), + (long long) (disk_ttl_s / 3600)); + + const std::string key = it->first; + ++it; + disk_erase(key); + } else { + ++it; + } + } + } + + if (disk_limit_size == 0) { + return; + } + + while (!disk.empty() && disk_size() > disk_limit_size) { + auto lru = std::min_element(disk.begin(), disk.end(), + [](const auto & a, const auto & b) { return a.second.t_last < b.second.t_last; }); + + SRV_INF("kv cache: evicting %s (%.3f GiB), store over its %.3f GiB limit\n", + lru->first.substr(0, 12).c_str(), lru->second.bytes / (1024.0*1024.0*1024.0), + disk_limit_size / (1024.0*1024.0*1024.0)); + + disk_erase(lru->first); + } +} + +bool server_prompt_cache::disk_store(const server_prompt & prompt, const std::string & conv_key, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { + if (dir.empty() || conv_key.empty() || prompt.tokens.size() < disk_min_tokens) { + return false; + } + + const std::string key = disk_key(conv_key); + + if (disk.count(key)) { + disk[key].t_last = kv_now(); // already stored, e.g. a regenerated turn + return false; + } + + std::vector packed; + try { + packed = prompt.tokens.serialize(); + } catch (const std::exception & err) { + SRV_WRN("kv cache: cannot serialize prompt: %s\n", err.what()); + return false; + } + + const std::string base = disk_path(key); + + std::error_code ec; + uint64_t bytes = 0; + + const int64_t t_start = ggml_time_us(); + + // written under .tmp and renamed, so a crash can never leave a half-written state behind a + // valid index. The token list travels inside the state file itself. + const size_t n_tgt = llama_state_seq_save_file(ctx_tgt, (base + ".kv.tmp").c_str(), id_slot, + (const llama_token *) packed.data(), packed.size() / sizeof(llama_token)); + + if (n_tgt == 0) { + SRV_WRN("kv cache: failed to write state for %s\n", key.substr(0, 12).c_str()); + std::filesystem::remove(base + ".kv.tmp", ec); + return false; + } + + bytes += n_tgt; + + if (ctx_dft) { + const size_t n_dft = llama_state_seq_save_file(ctx_dft, (base + ".dft.tmp").c_str(), id_slot, + (const llama_token *) packed.data(), packed.size() / sizeof(llama_token)); + + if (n_dft == 0) { + SRV_WRN("kv cache: failed to write draft state for %s\n", key.substr(0, 12).c_str()); + std::filesystem::remove(base + ".kv.tmp", ec); + std::filesystem::remove(base + ".dft.tmp", ec); + return false; + } + + bytes += n_dft; + } + + std::filesystem::rename(base + ".kv.tmp", base + ".kv", ec); + if (ec) { + SRV_WRN("kv cache: cannot commit state for %s: %s\n", key.substr(0, 12).c_str(), ec.message().c_str()); + std::filesystem::remove(base + ".kv.tmp", ec); + std::filesystem::remove(base + ".dft.tmp", ec); + return false; + } + + if (ctx_dft) { + std::filesystem::rename(base + ".dft.tmp", base + ".dft", ec); + } + + // checkpoints go with the state: on an SWA model the prefix is unusable without them + bytes += kv_write_checkpoints(base + ".ckpt", prompt.checkpoints); + + const uint64_t n_packed = packed.size() / sizeof(llama_token); + + // the index is the commit marker, so it goes last + if (!kv_write_idx(base + ".idx", n_packed, bytes)) { + SRV_WRN("kv cache: cannot write index for %s\n", key.substr(0, 12).c_str()); + std::filesystem::remove(base + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + std::filesystem::remove(base + ".ckpt", ec); + return false; + } + + server_prompt_cache_disk_entry entry; + + entry.n_packed = n_packed; + entry.bytes = bytes; + entry.t_last = kv_now(); + + disk[key] = entry; + + SRV_INF("kv cache: stored %s, %zu tokens, %.3f GiB in %.1f ms (%zu entries, %.3f GiB total)\n", + key.substr(0, 12).c_str(), prompt.tokens.size(), bytes / (1024.0*1024.0*1024.0), + (ggml_time_us() - t_start) / 1000.0, disk.size(), disk_size() / (1024.0*1024.0*1024.0)); + + disk_prune(); + + return true; +} + +bool server_prompt_cache::disk_load(server_prompt & prompt, const std::string & conv_key, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { + if (dir.empty() || conv_key.empty()) { + return false; + } + + // prune before the read, so an expired conversation is never resurrected + disk_prune(); + + const std::string key = disk_key(conv_key); + + auto it = disk.find(key); + if (it == disk.end()) { + return false; + } + + const std::string base = disk_path(key); + + const int64_t t_start = ggml_time_us(); + + llama_tokens packed(std::max(1, it->second.n_packed)); + size_t n_packed = 0; + + // llama_state_seq_load_file clears the destination sequence before installing cells, so a + // failure here cannot leave another conversation's KV behind under these tokens + const size_t n_read = llama_state_seq_load_file(ctx_tgt, (base + ".kv").c_str(), id_slot, + packed.data(), packed.size(), &n_packed); + + if (n_read == 0) { + SRV_WRN("kv cache: failed to restore %s, dropping it\n", key.substr(0, 12).c_str()); + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + disk_erase(key); + return false; + } + + packed.resize(n_packed); + + if (ctx_dft) { + llama_tokens packed_dft(std::max(1, it->second.n_packed)); + size_t n_dft = 0; + + if (llama_state_seq_load_file(ctx_dft, (base + ".dft").c_str(), id_slot, + packed_dft.data(), packed_dft.size(), &n_dft) == 0) { + SRV_WRN("kv cache: no draft state for %s, it will be rebuilt\n", key.substr(0, 12).c_str()); + } + } + + server_tokens restored; + try { + restored = server_tokens::deserialize(packed, disk_has_mtmd); + } catch (const std::exception & err) { + SRV_WRN("kv cache: cannot deserialize tokens for %s: %s\n", key.substr(0, 12).c_str(), err.what()); + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + disk_erase(key); + return false; + } + + if (!restored.validate(ctx_tgt)) { + SRV_WRN("kv cache: %s holds tokens this model cannot represent, dropping it\n", key.substr(0, 12).c_str()); + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + disk_erase(key); + return false; + } + + // The slot's own checkpoints describe the state it used to hold; replace them with the ones + // stored alongside this conversation. On an SWA model the restored sequence only carries the + // window, so without these the server discards the whole prefix and re-prefills. [TAG_KV_DISK_CKPT] + std::list ckpts; + if (!kv_read_checkpoints(base + ".ckpt", ckpts)) { + ckpts.clear(); + } + + const size_t n_ckpt = ckpts.size(); + + SRV_INF("kv cache: restored %s into slot %d, %zu tokens, %zu checkpoint(s), %.3f GiB in %.1f ms\n", + key.substr(0, 12).c_str(), id_slot, restored.size(), n_ckpt, + it->second.bytes / (1024.0*1024.0*1024.0), (ggml_time_us() - t_start) / 1000.0); + + prompt.checkpoints = std::move(ckpts); + prompt.tokens = std::move(restored); + + it->second.t_last = kv_now(); + + std::error_code ec; + std::filesystem::last_write_time(base + ".idx", std::filesystem::file_time_type::clock::now(), ec); + + return true; +} diff --git a/tools/server/server-kv-disk.h b/tools/server/server-kv-disk.h new file mode 100644 index 000000000000..561ad4f52bbd --- /dev/null +++ b/tools/server/server-kv-disk.h @@ -0,0 +1,46 @@ +#pragma once + +// On-disk KV cache: fingerprint a conversation, restore its KV state instead of re-prefilling it. +// +// The key is an exact sha1 of the conversation, not a prefix match. That works because the lookup +// hashes the conversation *truncated at the last assistant message*, which is byte-identical to +// what the previous turn stored: +// +// turn 1 arrives [sys, u1] -> no assistant message yet, no lookup +// generate A1, store under key([sys, u1, A1]) +// turn 2 arrives [sys, u1, A1, u2] -> truncate at last assistant -> key([sys, u1, A1]) HIT +// generate A2, store under key([sys, u1, A1, u2, A2]) +// +// Hashing the whole arriving conversation instead would never hit: every request carries one more +// message than anything already stored. + +#include +#include +#include + +// +// sha1 +// + +struct sha1_state { + uint32_t h[5] = { 0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0 }; + uint64_t n_len = 0; + uint8_t buf[64] = {0}; + uint32_t n_buf = 0; +}; + +void sha1_update(sha1_state & s, const void * data, size_t n); +std::string sha1_hex(sha1_state s); // by value: finalising must not consume the running state + +// the running state travels from the HTTP layer (where the messages are) to the slot (where the +// reply is), so it has to survive a trip through the task's json +std::string sha1_state_to_hex(const sha1_state & s); +bool sha1_state_from_hex(const std::string & hex, sha1_state & s); + +// +// conversation keys +// + +// Feeds one message into a running conversation digest. Deliberately a plain concatenation, so a +// digest can be extended with the assistant's reply once it exists. +void server_kv_add_message(sha1_state & s, const std::string & role, const std::string & content); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..5b9138217d1e 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1709,6 +1709,11 @@ size_t server_prompt_cache::n_tokens() const { } server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) { + // `--cache-ram 0` leaves limit_size at 0, which otherwise reads as "no limit" here + if (!ram_enabled) { + return nullptr; + } + // first check if the current state is contained fully in the cache for (auto it = states.begin(); it != states.end(); ++it) { const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..84bdc578ca6b 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -91,6 +91,12 @@ struct task_params { // per-request parameters for chat parsing common_chat_parser_params chat_parser_params; + // on-disk KV cache (--kv-cache-dir). kv_conv_key is the exact key of the conversation as it + // arrived, truncated at its last assistant message; kv_conv_base is the running sha1 over the + // full arriving conversation, extended with the reply to form the key this turn stores under. + std::string kv_conv_key; + std::string kv_conv_base; + // message spans for checkpointing common_chat_msg_spans message_spans; @@ -609,10 +615,20 @@ struct server_prompt_cache_state { } }; +// one conversation persisted under --kv-cache-dir, keyed by an exact sha1 of the conversation +struct server_prompt_cache_disk_entry { + uint64_t n_packed = 0; // length of the serialized token blob, sizes the restore buffer + uint64_t bytes = 0; + int64_t t_last = 0; // LRU clock, seeded from the file mtime at startup +}; + struct server_prompt_cache { server_prompt_cache(int32_t limit_size_mib, size_t limit_tokens) { this->limit_size = 1024ull*1024ull*(limit_size_mib < 0 ? 0 : limit_size_mib); this->limit_tokens = limit_tokens; + + // cache_ram_mib == 0 disables the in-memory tier, but the on-disk tier can still be used + this->ram_enabled = limit_size_mib != 0; } std::list states; @@ -623,6 +639,8 @@ struct server_prompt_cache { // in tokens, 0 = no limit size_t limit_tokens = 0; + bool ram_enabled = true; + size_t size() const; size_t n_tokens() const; @@ -632,6 +650,42 @@ struct server_prompt_cache { bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); void update(); + + // + // on-disk tier (--kv-cache-dir) + // + // Lookup is an exact key, not a prefix search: the caller hashes the conversation truncated at + // its last assistant message, which is byte-identical to what the previous turn stored. + // + + std::string dir; // "" = disabled + std::string model_tag; // mixed into the file name so two models never share an entry + + size_t disk_limit_size = 0; // in bytes, 0 = no limit + size_t disk_min_tokens = 0; + int64_t disk_ttl_s = 0; // drop entries older than this, 0 = never + bool disk_has_mtmd = false; + + std::map disk; + + bool disk_enabled() const { return !dir.empty(); } + + bool disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, int64_t ttl_s, bool has_mtmd); + + bool disk_store(const server_prompt & prompt, const std::string & conv_key, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + + bool disk_load(server_prompt & prompt, const std::string & conv_key, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + + // age first, then least-recently-used down to the size limit. Called before every read, so an + // expired conversation is never resurrected. + void disk_prune(); + + void disk_erase(const std::string & key); + + std::string disk_key (const std::string & conv_key) const; + std::string disk_path(const std::string & key) const; + + size_t disk_size() const; }; // used exclusively by router mode