From 9ab027238801de1a91b4da9155c13a39616b2df4 Mon Sep 17 00:00:00 2001 From: Daniele Zannotti Date: Thu, 3 Sep 2026 10:39:41 +0100 Subject: [PATCH 1/4] server : on-disk KV cache (--kv-cache-dir) Prefill on a Strix Halo costs minutes where the same KV state reads back from NVMe in seconds, so it is worth spending disk to never prefill a prefix 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 saves in 45 ms and restores in 47 ms, against 18.4 s to re-prefill it. A full 262k-token slot is 4.13 GB: 1.9 s to restore, ~3.5 min to re-prefill. Each conversation is fingerprinted with sha1(model tag + serialized prompt tokens) and its sequence state written under that name; a returning conversation is restored into a free slot instead of being reprocessed. The model tag is the weights path plus the arch description rather than --alias, so renaming an alias does not invalidate a cache and two models sharing an alias cannot read each other's state. 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 State blobs are multi-GiB, so .idx keeps the tokens resident: a lookup measures the longest common prefix against every stored conversation without reading any state. That matters because an exact whole-conversation hash can never hit on turn 2 - the next request appends a new user message - so lookup is prefix-anchored while the key stays a sha1 of the conversation. Selection reuses the in-memory tier's rule (only move if it both keeps more of the stored context and covers more of the incoming prompt), and skips entries below f_keep 0.25, since restoring gigabytes to reuse a sliver is slower than prefilling it. Writing an entry supersedes any stored prompt that is a prefix of it, so a conversation does not leave one file per turn. Eviction is least-recently-used against --kv-cache-max. Storing happens on slot reset, while the KV is still resident, so a finished turn is persisted immediately rather than when the next request happens to arrive. It runs on the inference thread because reading sequence state has to. --cache-ram 0 previously left no cache object at all; it now disables only the in-memory tier, and alloc() returns early rather than reading limit_size 0 as "no limit". Claude-Session: https://claude.ai/code/session_01JDKcT3SjBYaJKmpRJqTPmq --- common/arg.cpp | 30 ++ common/common.h | 6 + tools/server/CMakeLists.txt | 1 + tools/server/server-context.cpp | 57 ++- tools/server/server-kv-disk.cpp | 595 ++++++++++++++++++++++++++++++++ tools/server/server-task.cpp | 13 + tools/server/server-task.h | 49 +++ 7 files changed, 747 insertions(+), 4 deletions(-) create mode 100644 tools/server/server-kv-disk.cpp diff --git a/common/arg.cpp b/common/arg.cpp index 6d5edcb7f2d5..d2fa9c0ab2ce 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1736,6 +1736,36 @@ 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( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index 13bebcb52c54..bda190b55ee8 100644 --- a/common/common.h +++ b/common/common.h @@ -657,6 +657,12 @@ 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 + 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..1aa3b2c1de8b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -307,6 +307,14 @@ struct server_slot { return false; } + // persist before the in-memory tier so the conversation survives a restart even when the + // RAM tier is disabled or full + const bool stored = prompt_cache.disk_store(prompt, ctx_tgt, ctx_dft, id); + + if (!prompt_cache.ram_enabled) { + return stored; + } + 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; @@ -1507,6 +1515,18 @@ struct server_context_impl { if (slot.stats.n_gen > 0) { metrics_on_prediction(slot); } + + // persist the finished conversation while its KV is still in the slot, so the next + // turn restores it instead of prefilling it again. This runs on the inference thread + // because reading sequence state has to; it costs one sequential write of the state + // (~70 ms for a typical conversation, ~0.8 s for a full context) once per turn. + // [TAG_KV_DISK_SAVE] + if (prompt_cache && prompt_cache->disk_enabled() && + slot.stats.n_gen > 0 && slot.task && + !slot.task->is_child() && + slot.task->type == SERVER_TASK_TYPE_COMPLETION) { + prompt_cache->disk_store(slot.prompt, slot.ctx_tgt, slot.ctx_dft, slot.id); + } }; slot.reset(); @@ -1547,11 +1567,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 +1583,23 @@ 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; + + if (!prompt_cache->disk_init(params_base.kv_cache_dir, model_tag, + 1024ull*1024ull*(size_t) params_base.kv_cache_max_mib, + (size_t) params_base.kv_cache_min_toks, + 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 +1658,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) { @@ -1827,6 +1868,14 @@ struct server_context_impl { // cache prompts only for completion tasks update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; + // the on-disk store holds far more conversations than there are slots, so it is always + // worth a look - even when this slot was picked for already sharing a prefix. disk_load + // only moves if it beats what the slot holds, so consulting it cannot make things worse. + if (!update_cache && prompt_cache && prompt_cache->disk_enabled() && + task.type == SERVER_TASK_TYPE_COMPLETION) { + update_cache = true; + } + if (update_cache) { SRV_TRC("%s", "updating prompt cache\n"); diff --git a/tools/server/server-kv-disk.cpp b/tools/server/server-kv-disk.cpp new file mode 100644 index 000000000000..9a603b326728 --- /dev/null +++ b/tools/server/server-kv-disk.cpp @@ -0,0 +1,595 @@ +// On-disk KV cache (--kv-cache-dir). +// +// Prefill on some machines costs minutes while the KV state of the same conversation reads back from +// an NVMe in a couple of seconds, so it is worth spending disk to never prefill the same prefix twice. +// Each conversation is fingerprinted with sha1(model tag + serialized prompt tokens) and its sequence +// state written under that name; a returning conversation is restored into a free slot instead of +// being reprocessed. Entries are evicted least-recently-used once the store exceeds --kv-cache-max. +// +// Three files make up one entry: +// +// .kv target context sequence state (llama_state_seq_save_file, carries the token list) +// .dft draft context sequence state, only when speculative decoding is enabled +// .idx small index record, written last and deleted first, so it doubles as a commit marker +// +// The state blobs are multi-GiB, so the index keeps the tokens resident: a lookup can measure the +// longest common prefix against every stored conversation without reading a single byte of state. + +#include "server-task.h" + +#include "common.h" +#include "llama.h" +#include "server-common.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// +// sha1 +// + +struct sha1_state { + uint32_t h[5] = { 0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0 }; + uint64_t n_len = 0; + uint8_t buf[64]; + size_t n_buf = 0; +}; + +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; +} + +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 += 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); +} + +// the fingerprint the user asked for: the conversation, bound to the model that produced it +std::string kv_fingerprint(const std::string & model_tag, const std::vector & packed) { + sha1_state s; + + sha1_update(s, model_tag.data(), model_tag.size()); + + const uint8_t sep = 0; + sha1_update(s, &sep, 1); + + sha1_update(s, packed.data(), packed.size()); + + return sha1_hex(s); +} + +// +// index record +// + +constexpr char KV_IDX_MAGIC[8] = { 'K','V','C','A','C','H','E','1' }; +constexpr uint32_t KV_IDX_VERSION = 1; + +int64_t kv_now() { + return std::filesystem::file_time_type::clock::now().time_since_epoch().count(); +} + +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; +} + +bool kv_write_idx(const std::string & path, const std::vector & packed, uint64_t bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + if (!f) { + return false; + } + + const uint64_t n_packed = packed.size(); + + 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.write(packed.data(), packed.size()); + + f.close(); + + return f.good(); +} + +bool kv_read_idx(const std::string & path, std::vector & packed, uint64_t & bytes) { + std::ifstream f(path, std::ios::binary); + if (!f) { + return false; + } + + char magic[8]; + uint32_t version = 0; + uint64_t n_packed = 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; + } + + // a truncated index is indistinguishable from a corrupt one; both are simply dropped + if (n_packed > (1ull << 34) || n_packed % sizeof(llama_token) != 0) { + return false; + } + + packed.resize(n_packed); + f.read(packed.data(), n_packed); + + return f.good() || (size_t) f.gcount() == n_packed; +} + +} // namespace + +namespace { + +// a failed restore must leave the slot genuinely empty in every context, otherwise the next request +// would compute a common 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_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; +} + +bool server_prompt_cache::disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, size_t min_tokens, bool has_mtmd) { + this->dir = dir; + this->model_tag = model_tag; + this->disk_limit_size = limit_size; + this->disk_min_tokens = min_tokens; + 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; + } + + // an entry is committed by its .idx; anything else left over is from an interrupted write + 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); + continue; + } + + if (path.extension() != ".idx") { + continue; + } + + const std::string key = path.stem().string(); + + std::vector packed; + uint64_t bytes = 0; + + if (!kv_read_idx(path.string(), packed, bytes)) { + SRV_WRN("kv cache: dropping unreadable index %s\n", key.c_str()); + orphans.push_back(path); + continue; + } + + if (!std::filesystem::exists(disk_path(key) + ".kv", ec)) { + SRV_WRN("kv cache: dropping index %s with no state file\n", key.c_str()); + orphans.push_back(path); + continue; + } + + server_prompt_cache_disk_entry entry; + + entry.key = key; + entry.tokens = server_tokens::deserialize(kv_packed_to_tokens(packed), has_mtmd); + entry.n_packed = packed.size() / sizeof(llama_token); + entry.bytes = bytes; + + const auto mtime = std::filesystem::last_write_time(path, ec); + entry.t_last = ec ? 0 : mtime.time_since_epoch().count(); + + disk.push_back(std::move(entry)); + } + + if (ec) { + SRV_WRN("kv cache: cannot scan directory '%s': %s\n", dir.c_str(), ec.message().c_str()); + this->dir.clear(); + return false; + } + + for (const auto & path : orphans) { + std::filesystem::remove(path, ec); + + // .tmp files carry the full state, their siblings are removed with them + const std::string base = (path.parent_path() / path.stem()).string(); + std::filesystem::remove(base + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + } + + std::sort(disk.begin(), disk.end(), [](const auto & a, const auto & b) { return a.t_last < b.t_last; }); + + SRV_INF("kv cache: '%s' holds %zu conversation(s), %.3f GiB (limit %.3f GiB)\n", + dir.c_str(), disk.size(), disk_size() / (1024.0*1024.0*1024.0), + disk_limit_size / (1024.0*1024.0*1024.0)); + + disk_prune(); + + return true; +} + +void server_prompt_cache::disk_prune() { + if (dir.empty() || disk_limit_size == 0) { + return; + } + + std::error_code ec; + + while (!disk.empty() && disk_size() > disk_limit_size) { + // disk is kept ordered by t_last, so the front is the least recently used + auto it = std::min_element(disk.begin(), disk.end(), + [](const auto & a, const auto & b) { return a.t_last < b.t_last; }); + + SRV_WRN("kv cache: evicting %s (%zu tokens, %.3f GiB), store over limit\n", + it->key.substr(0, 12).c_str(), it->tokens.size(), it->bytes / (1024.0*1024.0*1024.0)); + + // .idx first: an entry stops existing the moment its commit marker is gone + std::filesystem::remove(disk_path(it->key) + ".idx", ec); + std::filesystem::remove(disk_path(it->key) + ".kv", ec); + std::filesystem::remove(disk_path(it->key) + ".dft", ec); + + disk.erase(it); + } +} + +bool server_prompt_cache::disk_store(const server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { + if (dir.empty() || prompt.tokens.size() < disk_min_tokens) { + 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 key = kv_fingerprint(model_tag, packed); + + // already stored: this is the common case when a slot is saved twice without generating + for (auto & e : disk) { + if (e.key == key) { + e.t_last = kv_now(); + 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 = 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 == 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; + } + + if (ctx_dft) { + const size_t n = 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 == 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; + } + + 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); + } + + // the index is the commit marker, so it goes last + if (!kv_write_idx(base + ".idx", 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); + return false; + } + + const double t_ms = (ggml_time_us() - t_start) / 1000.0; + + // supersede earlier turns of this same conversation: any stored prompt that is a prefix of the + // one just written is reachable through it, so keeping it only costs disk + for (auto it = disk.begin(); it != disk.end();) { + if (it->tokens.size() < prompt.tokens.size() && + it->tokens.get_common_prefix(prompt.tokens) == it->tokens.size()) { + SRV_TRC("kv cache: superseding %s (%zu tokens)\n", it->key.substr(0, 12).c_str(), it->tokens.size()); + + std::filesystem::remove(disk_path(it->key) + ".idx", ec); + std::filesystem::remove(disk_path(it->key) + ".kv", ec); + std::filesystem::remove(disk_path(it->key) + ".dft", ec); + + it = disk.erase(it); + } else { + ++it; + } + } + + server_prompt_cache_disk_entry entry; + + entry.key = key; + entry.tokens = prompt.tokens.clone(); + entry.n_packed = packed.size() / sizeof(llama_token); + entry.bytes = bytes; + entry.t_last = kv_now(); + + disk.push_back(std::move(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), t_ms, + disk.size(), disk_size() / (1024.0*1024.0*1024.0)); + + disk_prune(); + + return true; +} + +bool server_prompt_cache::disk_load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { + if (dir.empty() || tokens_new.empty()) { + return false; + } + + const int lcp_cur = prompt.tokens.get_common_prefix(tokens_new); + + // same rule as the in-memory tier: only move if it both keeps more of the stored context and + // covers more of the incoming prompt than what the slot already holds + float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_cur) / prompt.tokens.size() : -1.0f; + float f_sim_best = float(lcp_cur) / tokens_new.size(); + + auto it_best = disk.end(); + + for (auto it = disk.begin(); it != disk.end(); ++it) { + if (it->tokens.empty()) { + continue; + } + + const int lcp = it->tokens.get_common_prefix(tokens_new); + + const float f_keep = float(lcp) / it->tokens.size(); + const float f_sim = float(lcp) / tokens_new.size(); + + // restoring gigabytes to reuse a sliver of them is slower than prefilling the sliver + if (f_keep < 0.25f) { + continue; + } + + if (f_keep_best < f_keep && f_sim_best < f_sim) { + f_keep_best = f_keep; + f_sim_best = f_sim; + + it_best = it; + } + } + + if (it_best == disk.end()) { + return false; + } + + const std::string base = disk_path(it_best->key); + + const int64_t t_start = ggml_time_us(); + + llama_tokens packed(std::max(1, it_best->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", it_best->key.substr(0, 12).c_str()); + + std::error_code ec; + std::filesystem::remove(base + ".idx", ec); + std::filesystem::remove(base + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + + disk.erase(it_best); + + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + + return false; + } + + packed.resize(n_packed); + + if (ctx_dft) { + llama_tokens packed_dft(std::max(1, it_best->n_packed)); + size_t n_packed_dft = 0; + + if (llama_state_seq_load_file(ctx_dft, (base + ".dft").c_str(), id_slot, + packed_dft.data(), packed_dft.size(), &n_packed_dft) == 0) { + SRV_WRN("kv cache: no draft state for %s, it will be rebuilt\n", it_best->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", it_best->key.substr(0, 12).c_str(), err.what()); + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + return false; + } + + if (!restored.validate(ctx_tgt)) { + SRV_WRN("kv cache: %s holds tokens this model cannot represent, dropping it\n", it_best->key.substr(0, 12).c_str()); + + std::error_code ec; + std::filesystem::remove(base + ".idx", ec); + std::filesystem::remove(base + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + + disk.erase(it_best); + + kv_reset_slot(prompt, ctx_tgt, ctx_dft, id_slot); + + return false; + } + + const double t_ms = (ggml_time_us() - t_start) / 1000.0; + + SRV_INF("kv cache: restored %s into slot %d, %zu tokens, %.3f GiB in %.1f ms (f_keep = %.3f, f_sim = %.3f)\n", + it_best->key.substr(0, 12).c_str(), id_slot, restored.size(), + it_best->bytes / (1024.0*1024.0*1024.0), t_ms, f_keep_best, f_sim_best); + + // checkpoints describe the state this slot used to hold, not the one just installed + prompt.checkpoints.clear(); + prompt.tokens = std::move(restored); + + it_best->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-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..eb5fda5b3b10 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); @@ -1862,6 +1867,14 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok prompt = std::move(it_best->prompt); states.erase(it_best); + + return true; + } + + // nothing better in memory. The on-disk tier holds far more conversations than RAM can, and on + // this class of machine reading a state back beats re-prefilling it by two orders of magnitude. + if (disk_enabled()) { + disk_load(prompt, tokens_new, ctx_tgt, ctx_dft, id_slot); } return true; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..bca680509220 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -609,10 +609,23 @@ struct server_prompt_cache_state { } }; +// one conversation persisted under --kv-cache-dir. The tokens are kept resident so a lookup can +// measure the longest common prefix without touching the (multi-GiB) state blob on disk. +struct server_prompt_cache_disk_entry { + std::string key; // sha1(model tag + packed prompt tokens), also the file name + server_tokens tokens; + size_t n_packed = 0; // length of the serialized token blob, sizes the restore buffer + size_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 +636,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 +647,40 @@ 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) + // + + std::string dir; // "" = disabled + std::string model_tag; // mixed into the fingerprint so two models never share an entry + + size_t disk_limit_size = 0; // in bytes, 0 = no limit + size_t disk_min_tokens = 0; + bool disk_has_mtmd = false; + + std::vector disk; + + bool disk_enabled() const { return !dir.empty(); } + + // scan dir and build the in-memory index; returns false if the directory is unusable + bool disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, size_t min_tokens, bool has_mtmd); + + // persist prompt's KV state, replacing any stored entry that is a prefix of it (an earlier turn + // of the same conversation). Returns true if something was written. + bool disk_store(const server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + + // restore the stored entry sharing the longest prefix with tokens_new, if it beats what the slot + // already holds. Mirrors the in-memory load() selection rule. + bool disk_load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + + // evict least-recently-used entries until the store fits in disk_limit_size + void disk_prune(); + + std::string disk_path(const std::string & key) const; + + // stats, for logging + size_t disk_size() const; }; // used exclusively by router mode From 80bf148eb74e0e3164a5440ff0f3380ce436c7d8 Mon Sep 17 00:00:00 2001 From: Daniele Zannotti Date: Thu, 3 Sep 2026 10:57:09 +0100 Subject: [PATCH 2/4] server : key the on-disk KV cache by exact conversation hash Replaces the token-prefix index with the exact-hash scheme this was meant to be. An exact sha1 of the *arriving* conversation can never hit: every request carries one more message than anything already stored. So the lookup hashes the conversation truncated at its last assistant message, which is byte-identical to what the previous turn saved: turn 1 [sys, u1] -> no assistant message, no lookup generate A1, store under key([sys, u1, A1]) turn 2 [sys, u1, A1, u2] -> truncate -> key([sys, u1, A1]) HIT generate A2, store under key([sys, u1, A1, u2, A2]) One exact key, an O(1) map lookup, no resident token index and no longest-common-prefix scan. The keys are built at the chat route, where the messages are still visible, and travel to the slot on the task. The reply is hashed as the client will echo it back (parsed content, reasoning stripped), since that is what the next turn will send. Storing moved to send_final_response, because it has to happen before generated_text is moved into the response - callback_on_reset sees it already empty. Eviction now runs before every lookup, and drops entries by age (--kv-cache-ttl, default 3 days) before falling back to least-recently-used against --kv-cache-max. Two conversations that really are identical collapse onto one entry and share it, which is correct: they restore the same prefix and then prefill their own continuations. Claude-Session: https://claude.ai/code/session_01JDKcT3SjBYaJKmpRJqTPmq --- common/arg.cpp | 11 + common/common.h | 1 + tools/server/server-context.cpp | 117 +++++++-- tools/server/server-kv-disk.cpp | 449 +++++++++++++++----------------- tools/server/server-kv-disk.h | 46 ++++ tools/server/server-task.cpp | 8 - tools/server/server-task.h | 49 ++-- 7 files changed, 391 insertions(+), 290 deletions(-) create mode 100644 tools/server/server-kv-disk.h diff --git a/common/arg.cpp b/common/arg.cpp index d2fa9c0ab2ce..1b42c7fa6cc0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1766,6 +1766,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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 bda190b55ee8..0f729489f7f9 100644 --- a/common/common.h +++ b/common/common.h @@ -662,6 +662,7 @@ struct common_params { 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 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1aa3b2c1de8b..91b067b5dc88 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,12 +308,8 @@ struct server_slot { return false; } - // persist before the in-memory tier so the conversation survives a restart even when the - // RAM tier is disabled or full - const bool stored = prompt_cache.disk_store(prompt, ctx_tgt, ctx_dft, id); - if (!prompt_cache.ram_enabled) { - return stored; + return false; } const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); @@ -1515,18 +1512,6 @@ struct server_context_impl { if (slot.stats.n_gen > 0) { metrics_on_prediction(slot); } - - // persist the finished conversation while its KV is still in the slot, so the next - // turn restores it instead of prefilling it again. This runs on the inference thread - // because reading sequence state has to; it costs one sequential write of the state - // (~70 ms for a typical conversation, ~0.8 s for a full context) once per turn. - // [TAG_KV_DISK_SAVE] - if (prompt_cache && prompt_cache->disk_enabled() && - slot.stats.n_gen > 0 && slot.task && - !slot.task->is_child() && - slot.task->type == SERVER_TASK_TYPE_COMPLETION) { - prompt_cache->disk_store(slot.prompt, slot.ctx_tgt, slot.ctx_dft, slot.id); - } }; slot.reset(); @@ -1593,9 +1578,11 @@ struct server_context_impl { 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, - (size_t) params_base.kv_cache_min_toks, + (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"); } @@ -1868,14 +1855,6 @@ struct server_context_impl { // cache prompts only for completion tasks update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; - // the on-disk store holds far more conversations than there are slots, so it is always - // worth a look - even when this slot was picked for already sharing a prefix. disk_load - // only moves if it beats what the slot holds, so consulting it cannot make things worse. - if (!update_cache && prompt_cache && prompt_cache->disk_enabled() && - task.type == SERVER_TASK_TYPE_COMPLETION) { - update_cache = true; - } - if (update_cache) { SRV_TRC("%s", "updating prompt cache\n"); @@ -1893,6 +1872,14 @@ 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()) { + prompt_cache->disk_load(ret->prompt, task.params.kv_conv_key, ret->ctx_tgt, ret->ctx_dft, ret->id); + } + return ret; } @@ -2324,6 +2311,30 @@ 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); + + prompt_cache->disk_store(slot.prompt, sha1_hex(base), slot.ctx_tgt, slot.ctx_dft, slot.id); + } + } + auto res = std::make_unique(); res->id = slot.task->id; @@ -4484,6 +4495,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, @@ -4553,6 +4613,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; @@ -5165,6 +5229,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 index 9a603b326728..cdac102d5d0d 100644 --- a/tools/server/server-kv-disk.cpp +++ b/tools/server/server-kv-disk.cpp @@ -1,20 +1,21 @@ -// On-disk KV cache (--kv-cache-dir). +// On-disk KV cache (--kv-cache-dir). See server-kv-disk.h for the keying scheme. // -// Prefill on some machines costs minutes while the KV state of the same conversation reads back from -// an NVMe in a couple of seconds, so it is worth spending disk to never prefill the same prefix twice. -// Each conversation is fingerprinted with sha1(model tag + serialized prompt tokens) and its sequence -// state written under that name; a returning conversation is restored into a free slot instead of -// being reprocessed. Entries are evicted least-recently-used once the store exceeds --kv-cache-max. +// 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. // -// Three files make up one entry: +// An entry is three files: // -// .kv target context sequence state (llama_state_seq_save_file, carries the token list) -// .dft draft context sequence state, only when speculative decoding is enabled -// .idx small index record, written last and deleted first, so it doubles as a commit marker +// .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 // -// The state blobs are multi-GiB, so the index keeps the tokens resident: a lookup can measure the -// longest common prefix against every stored conversation without reading a single byte of state. +// 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" @@ -22,25 +23,17 @@ #include "server-common.h" #include -#include +#include #include #include #include -#include #include -namespace { - // // sha1 // -struct sha1_state { - uint32_t h[5] = { 0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0 }; - uint64_t n_len = 0; - uint8_t buf[64]; - size_t n_buf = 0; -}; +namespace { inline uint32_t sha1_rol(uint32_t v, int b) { return (v << b) | (v >> (32 - b)); @@ -63,27 +56,27 @@ void sha1_compress(sha1_state & s, const uint8_t * p) { uint32_t f, k; if (i < 20) { - f = (b & c) | ((~b) & d); k = 0x5A827999; + f = (b & c) | ((~b) & d); k = 0x5A827999; } else if (i < 40) { - f = b ^ c ^ d; k = 0x6ED9EBA1; + f = b ^ c ^ d; k = 0x6ED9EBA1; } else if (i < 60) { - f = (b & c) | (b & d) | (c & d); k = 0x8F1BBCDC; + f = (b & c) | (b & d) | (c & d); k = 0x8F1BBCDC; } else { - f = b ^ c ^ d; k = 0xCA62C1D6; + 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; + 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; @@ -94,7 +87,7 @@ void sha1_update(sha1_state & s, const void * data, size_t n) { memcpy(s.buf + s.n_buf, p, take); - s.n_buf += take; + s.n_buf += (uint32_t) take; p += take; n -= take; @@ -105,7 +98,7 @@ void sha1_update(sha1_state & s, const void * data, size_t n) { } } -std::string sha1_hex(sha1_state & s) { +std::string sha1_hex(sha1_state s) { const uint64_t n_bits = s.n_len * 8; const uint8_t pad = 0x80; @@ -130,69 +123,103 @@ std::string sha1_hex(sha1_state & s) { return std::string(out, 40); } -// the fingerprint the user asked for: the conversation, bound to the model that produced it -std::string kv_fingerprint(const std::string & model_tag, const std::vector & packed) { - sha1_state s; +std::string sha1_state_to_hex(const sha1_state & s) { + uint8_t raw[sizeof(sha1_state)]; + memcpy(raw, &s, sizeof(s)); - sha1_update(s, model_tag.data(), model_tag.size()); + 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]; + } - const uint8_t sep = 0; - sha1_update(s, &sep, 1); + return out; +} - sha1_update(s, packed.data(), packed.size()); +bool sha1_state_from_hex(const std::string & hex, sha1_state & s) { + if (hex.size() != sizeof(sha1_state) * 2) { + return false; + } - return sha1_hex(s); + 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); } // -// index record +// on-disk store // -constexpr char KV_IDX_MAGIC[8] = { 'K','V','C','A','C','H','E','1' }; -constexpr uint32_t KV_IDX_VERSION = 1; +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(); } -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; +// 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, const std::vector & packed, uint64_t bytes) { +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; } - const uint64_t n_packed = packed.size(); - 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.write(packed.data(), packed.size()); - f.close(); return f.good(); } -bool kv_read_idx(const std::string & path, std::vector & packed, uint64_t & bytes) { +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; - uint64_t n_packed = 0; + uint32_t version = 0; f.read(magic, sizeof(magic)); f.read((char *) &version, sizeof(version)); @@ -203,23 +230,21 @@ bool kv_read_idx(const std::string & path, std::vector & packed, uint64_t return false; } - // a truncated index is indistinguishable from a corrupt one; both are simply dropped - if (n_packed > (1ull << 34) || n_packed % sizeof(llama_token) != 0) { - return false; - } - - packed.resize(n_packed); - f.read(packed.data(), n_packed); - - return f.good() || (size_t) f.gcount() == n_packed; + return n_packed <= (1ull << 32); } -} // namespace +llama_tokens kv_packed_to_tokens(const std::vector & packed) { + llama_tokens out(packed.size() / sizeof(llama_token)); -namespace { + if (!out.empty()) { + memcpy(out.data(), packed.data(), out.size() * sizeof(llama_token)); + } + + return out; +} // a failed restore must leave the slot genuinely empty in every context, otherwise the next request -// would compute a common prefix against cells that were never installed +// 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); @@ -232,6 +257,17 @@ void kv_reset_slot(server_prompt & prompt, llama_context * ctx_tgt, llama_contex } // 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(); } @@ -239,18 +275,29 @@ std::string server_prompt_cache::disk_path(const std::string & key) const { size_t server_prompt_cache::disk_size() const { size_t res = 0; - for (const auto & e : disk) { + for (const auto & [_, e] : disk) { res += e.bytes; } return res; } -bool server_prompt_cache::disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, size_t min_tokens, bool has_mtmd) { +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); + + 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_min_tokens = min_tokens; + this->disk_ttl_s = ttl_s; this->disk_has_mtmd = has_mtmd; std::error_code ec; @@ -262,14 +309,13 @@ bool server_prompt_cache::disk_init(const std::string & dir, const std::string & return false; } - // an entry is committed by its .idx; anything else left over is from an interrupted write 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); + orphans.push_back(path); // an interrupted write continue; } @@ -279,54 +325,38 @@ bool server_prompt_cache::disk_init(const std::string & dir, const std::string & const std::string key = path.stem().string(); - std::vector packed; - uint64_t bytes = 0; - - if (!kv_read_idx(path.string(), packed, bytes)) { - SRV_WRN("kv cache: dropping unreadable index %s\n", key.c_str()); - orphans.push_back(path); - continue; - } + uint64_t n_packed = 0; + uint64_t bytes = 0; - if (!std::filesystem::exists(disk_path(key) + ".kv", ec)) { - SRV_WRN("kv cache: dropping index %s with no state file\n", key.c_str()); + 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.key = key; - entry.tokens = server_tokens::deserialize(kv_packed_to_tokens(packed), has_mtmd); - entry.n_packed = packed.size() / sizeof(llama_token); + entry.n_packed = n_packed; entry.bytes = bytes; const auto mtime = std::filesystem::last_write_time(path, ec); - entry.t_last = ec ? 0 : mtime.time_since_epoch().count(); + entry.t_last = ec ? kv_now() : mtime.time_since_epoch().count(); - disk.push_back(std::move(entry)); - } - - if (ec) { - SRV_WRN("kv cache: cannot scan directory '%s': %s\n", dir.c_str(), ec.message().c_str()); - this->dir.clear(); - return false; + disk.emplace(key, entry); } for (const auto & path : orphans) { - std::filesystem::remove(path, ec); - - // .tmp files carry the full state, their siblings are removed with them 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::sort(disk.begin(), disk.end(), [](const auto & a, const auto & b) { return a.t_last < b.t_last; }); - - SRV_INF("kv cache: '%s' holds %zu conversation(s), %.3f GiB (limit %.3f GiB)\n", + 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)); + disk_limit_size / (1024.0*1024.0*1024.0), (long long) (disk_ttl_s / 3600)); disk_prune(); @@ -334,31 +364,54 @@ bool server_prompt_cache::disk_init(const std::string & dir, const std::string & } void server_prompt_cache::disk_prune() { - if (dir.empty() || disk_limit_size == 0) { + if (dir.empty()) { return; } - std::error_code ec; + // 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(); - while (!disk.empty() && disk_size() > disk_limit_size) { - // disk is kept ordered by t_last, so the front is the least recently used - auto it = std::min_element(disk.begin(), disk.end(), - [](const auto & a, const auto & b) { return a.t_last < b.t_last; }); + 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; + } - SRV_WRN("kv cache: evicting %s (%zu tokens, %.3f GiB), store over limit\n", - it->key.substr(0, 12).c_str(), it->tokens.size(), it->bytes / (1024.0*1024.0*1024.0)); + 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; }); - // .idx first: an entry stops existing the moment its commit marker is gone - std::filesystem::remove(disk_path(it->key) + ".idx", ec); - std::filesystem::remove(disk_path(it->key) + ".kv", ec); - std::filesystem::remove(disk_path(it->key) + ".dft", ec); + 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(it); + disk_erase(lru->first); } } -bool server_prompt_cache::disk_store(const server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { - if (dir.empty() || prompt.tokens.size() < disk_min_tokens) { +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; } @@ -370,16 +423,6 @@ bool server_prompt_cache::disk_store(const server_prompt & prompt, llama_context return false; } - const std::string key = kv_fingerprint(model_tag, packed); - - // already stored: this is the common case when a slot is saved twice without generating - for (auto & e : disk) { - if (e.key == key) { - e.t_last = kv_now(); - return false; - } - } - const std::string base = disk_path(key); std::error_code ec; @@ -389,29 +432,29 @@ bool server_prompt_cache::disk_store(const server_prompt & prompt, llama_context // 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 = 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 == 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; - } + 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)); - bytes += n; + 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 = llama_state_seq_save_file(ctx_dft, (base + ".dft.tmp").c_str(), id_slot, + 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 == 0) { + + 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; + bytes += n_dft; } std::filesystem::rename(base + ".kv.tmp", base + ".kv", ec); @@ -426,98 +469,53 @@ bool server_prompt_cache::disk_store(const server_prompt & prompt, llama_context std::filesystem::rename(base + ".dft.tmp", base + ".dft", ec); } + 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", packed, bytes)) { + 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); return false; } - const double t_ms = (ggml_time_us() - t_start) / 1000.0; - - // supersede earlier turns of this same conversation: any stored prompt that is a prefix of the - // one just written is reachable through it, so keeping it only costs disk - for (auto it = disk.begin(); it != disk.end();) { - if (it->tokens.size() < prompt.tokens.size() && - it->tokens.get_common_prefix(prompt.tokens) == it->tokens.size()) { - SRV_TRC("kv cache: superseding %s (%zu tokens)\n", it->key.substr(0, 12).c_str(), it->tokens.size()); - - std::filesystem::remove(disk_path(it->key) + ".idx", ec); - std::filesystem::remove(disk_path(it->key) + ".kv", ec); - std::filesystem::remove(disk_path(it->key) + ".dft", ec); - - it = disk.erase(it); - } else { - ++it; - } - } - server_prompt_cache_disk_entry entry; - entry.key = key; - entry.tokens = prompt.tokens.clone(); - entry.n_packed = packed.size() / sizeof(llama_token); + entry.n_packed = n_packed; entry.bytes = bytes; entry.t_last = kv_now(); - disk.push_back(std::move(entry)); + 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), t_ms, - disk.size(), disk_size() / (1024.0*1024.0*1024.0)); + 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 server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { - if (dir.empty() || tokens_new.empty()) { +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; } - const int lcp_cur = prompt.tokens.get_common_prefix(tokens_new); - - // same rule as the in-memory tier: only move if it both keeps more of the stored context and - // covers more of the incoming prompt than what the slot already holds - float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_cur) / prompt.tokens.size() : -1.0f; - float f_sim_best = float(lcp_cur) / tokens_new.size(); - - auto it_best = disk.end(); - - for (auto it = disk.begin(); it != disk.end(); ++it) { - if (it->tokens.empty()) { - continue; - } - - const int lcp = it->tokens.get_common_prefix(tokens_new); - - const float f_keep = float(lcp) / it->tokens.size(); - const float f_sim = float(lcp) / tokens_new.size(); - - // restoring gigabytes to reuse a sliver of them is slower than prefilling the sliver - if (f_keep < 0.25f) { - continue; - } - - if (f_keep_best < f_keep && f_sim_best < f_sim) { - f_keep_best = f_keep; - f_sim_best = f_sim; + // prune before the read, so an expired conversation is never resurrected + disk_prune(); - it_best = it; - } - } + const std::string key = disk_key(conv_key); - if (it_best == disk.end()) { + auto it = disk.find(key); + if (it == disk.end()) { return false; } - const std::string base = disk_path(it_best->key); + const std::string base = disk_path(key); const int64_t t_start = ggml_time_us(); - llama_tokens packed(std::max(1, it_best->n_packed)); + 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 @@ -526,29 +524,21 @@ bool server_prompt_cache::disk_load(server_prompt & prompt, const server_tokens packed.data(), packed.size(), &n_packed); if (n_read == 0) { - SRV_WRN("kv cache: failed to restore %s, dropping it\n", it_best->key.substr(0, 12).c_str()); - - std::error_code ec; - std::filesystem::remove(base + ".idx", ec); - std::filesystem::remove(base + ".kv", ec); - std::filesystem::remove(base + ".dft", ec); - - disk.erase(it_best); - + 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_best->n_packed)); - size_t n_packed_dft = 0; + 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_packed_dft) == 0) { - SRV_WRN("kv cache: no draft state for %s, it will be rebuilt\n", it_best->key.substr(0, 12).c_str()); + 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()); } } @@ -556,37 +546,28 @@ bool server_prompt_cache::disk_load(server_prompt & prompt, const server_tokens 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", it_best->key.substr(0, 12).c_str(), err.what()); + 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", it_best->key.substr(0, 12).c_str()); - - std::error_code ec; - std::filesystem::remove(base + ".idx", ec); - std::filesystem::remove(base + ".kv", ec); - std::filesystem::remove(base + ".dft", ec); - - disk.erase(it_best); - + 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; } - const double t_ms = (ggml_time_us() - t_start) / 1000.0; - - SRV_INF("kv cache: restored %s into slot %d, %zu tokens, %.3f GiB in %.1f ms (f_keep = %.3f, f_sim = %.3f)\n", - it_best->key.substr(0, 12).c_str(), id_slot, restored.size(), - it_best->bytes / (1024.0*1024.0*1024.0), t_ms, f_keep_best, f_sim_best); + SRV_INF("kv cache: restored %s into slot %d, %zu tokens, %.3f GiB in %.1f ms\n", + key.substr(0, 12).c_str(), id_slot, restored.size(), + it->second.bytes / (1024.0*1024.0*1024.0), (ggml_time_us() - t_start) / 1000.0); // checkpoints describe the state this slot used to hold, not the one just installed prompt.checkpoints.clear(); prompt.tokens = std::move(restored); - it_best->t_last = kv_now(); + 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); 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 eb5fda5b3b10..5b9138217d1e 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1867,14 +1867,6 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok prompt = std::move(it_best->prompt); states.erase(it_best); - - return true; - } - - // nothing better in memory. The on-disk tier holds far more conversations than RAM can, and on - // this class of machine reading a state back beats re-prefilling it by two orders of magnitude. - if (disk_enabled()) { - disk_load(prompt, tokens_new, ctx_tgt, ctx_dft, id_slot); } return true; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index bca680509220..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,14 +615,11 @@ struct server_prompt_cache_state { } }; -// one conversation persisted under --kv-cache-dir. The tokens are kept resident so a lookup can -// measure the longest common prefix without touching the (multi-GiB) state blob on disk. +// one conversation persisted under --kv-cache-dir, keyed by an exact sha1 of the conversation struct server_prompt_cache_disk_entry { - std::string key; // sha1(model tag + packed prompt tokens), also the file name - server_tokens tokens; - size_t n_packed = 0; // length of the serialized token blob, sizes the restore buffer - size_t bytes = 0; - int64_t t_last = 0; // LRU clock, seeded from the file mtime at startup + 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 { @@ -651,35 +654,37 @@ struct server_prompt_cache { // // 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 fingerprint so two models never share an entry + 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; - bool disk_has_mtmd = false; + 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::vector disk; + std::map disk; bool disk_enabled() const { return !dir.empty(); } - // scan dir and build the in-memory index; returns false if the directory is unusable - bool disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, size_t min_tokens, bool has_mtmd); + bool disk_init(const std::string & dir, const std::string & model_tag, size_t limit_size, int64_t ttl_s, bool has_mtmd); - // persist prompt's KV state, replacing any stored entry that is a prefix of it (an earlier turn - // of the same conversation). Returns true if something was written. - bool disk_store(const server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + bool disk_store(const server_prompt & prompt, const std::string & conv_key, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); - // restore the stored entry sharing the longest prefix with tokens_new, if it beats what the slot - // already holds. Mirrors the in-memory load() selection rule. - bool disk_load(server_prompt & prompt, const server_tokens & tokens_new, 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); - // evict least-recently-used entries until the store fits in disk_limit_size + // 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; - // stats, for logging size_t disk_size() const; }; From 1396ae94d943056f34f277f03bb44647e63bf456 Mon Sep 17 00:00:00 2001 From: Daniele Zannotti Date: Thu, 3 Sep 2026 11:07:28 +0100 Subject: [PATCH 3/4] server : persist context checkpoints with the on-disk KV state Restoring a conversation was measurably worse than not caching at all on gemma4: turn 2 went from cache_n 1224 / prompt_n 13 with the cache off, to cache_n 0 with it on. gemma-4 is a sliding-window model. llama_state_seq_save_file only persists the SWA window, so after a restore llama_memory_seq_pos_min sits above pos_min_thold and update_slots will only reuse the prefix if a context checkpoint reaches further back. disk_load cleared prompt.checkpoints and the state file carried none, so the checkpoint search found nothing, do_reset fired and the entire restored prefix was thrown away. This is the same reason the in-memory tier works: alloc() copies prompt.checkpoints alongside the state blob. Store them in a fourth file, .ckpt, and load them back. Claude-Session: https://claude.ai/code/session_01JDKcT3SjBYaJKmpRJqTPmq --- tools/server/server-kv-disk.cpp | 130 +++++++++++++++++++++++++++++--- 1 file changed, 118 insertions(+), 12 deletions(-) diff --git a/tools/server/server-kv-disk.cpp b/tools/server/server-kv-disk.cpp index cdac102d5d0d..d322c56b1de0 100644 --- a/tools/server/server-kv-disk.cpp +++ b/tools/server/server-kv-disk.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include // @@ -243,6 +244,96 @@ llama_tokens kv_packed_to_tokens(const std::vector & packed) { 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) { @@ -286,9 +377,10 @@ 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) + ".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); } @@ -350,8 +442,9 @@ bool server_prompt_cache::disk_init(const std::string & dir, const std::string & 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 + ".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", @@ -469,13 +562,17 @@ bool server_prompt_cache::disk_store(const server_prompt & prompt, const std::st 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 + ".kv", ec); + std::filesystem::remove(base + ".dft", ec); + std::filesystem::remove(base + ".ckpt", ec); return false; } @@ -559,13 +656,22 @@ bool server_prompt_cache::disk_load(server_prompt & prompt, const std::string & return false; } - SRV_INF("kv cache: restored %s into slot %d, %zu tokens, %.3f GiB in %.1f ms\n", - key.substr(0, 12).c_str(), id_slot, restored.size(), + // 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); - // checkpoints describe the state this slot used to hold, not the one just installed - prompt.checkpoints.clear(); - prompt.tokens = std::move(restored); + prompt.checkpoints = std::move(ckpts); + prompt.tokens = std::move(restored); it->second.t_last = kv_now(); From 6d0f67ec060b7ef342e1960d2cbb54fc628c8c2a Mon Sep 17 00:00:00 2001 From: Daniele Zannotti Date: Thu, 3 Sep 2026 20:53:36 +0100 Subject: [PATCH 4/4] server : do not restore a conversation the slot already holds disk_load restored unconditionally whenever the key was found on disk, without ever looking at what the slot already contained. On a returning turn that lands back on the slot holding its own state -- which --slot-prompt-similarity makes likely -- that reads hundreds of MiB off disk to install a byte-identical copy, and drops the checkpoints built since the last restore on the way through. On qwen38-27b, whose entries are about 2 GiB, that is 2 GiB read to change nothing. Track the conversation key whose KV a slot currently holds, set it both when a turn is stored and when one is restored, clear it in prompt_clear(), and skip the restore when it already matches. The key stored after generation is exactly the key the next turn looks itself up by, so the comparison is direct. Both directions of a stale key are safe: a missed skip costs one needless read, and a missed restore falls back to the ordinary common-prefix path and prefills more. Neither can produce wrong output. The existing 1,2,3,4 / 1,2,4,3 verification cannot catch this, because four conversations through two slots evict between every turn, so the slot never already holds the one being asked for. Added a 1,2,1,2 case, where both conversations stay resident. The broker this replaced had the same guard (`previous == key`), and it was dropped in the port. Claude-Session: https://claude.ai/code/session_01JDKcT3SjBYaJKmpRJqTPmq --- tools/server/server-context.cpp | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 91b067b5dc88..3d777bb06a52 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -342,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()); @@ -351,6 +355,8 @@ struct server_slot { } prompt.clear(); + + kv_conv_key.clear(); } std::vector lora; @@ -1877,7 +1883,16 @@ struct server_context_impl { // 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()) { - prompt_cache->disk_load(ret->prompt, task.params.kv_conv_key, ret->ctx_tgt, ret->ctx_dft, ret->id); + 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; @@ -2331,7 +2346,12 @@ struct server_context_impl { server_kv_add_message(base, "assistant", reply); - prompt_cache->disk_store(slot.prompt, sha1_hex(base), slot.ctx_tgt, slot.ctx_dft, slot.id); + 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; } }