server: on-disk KV cache (--kv-cache-dir) — restore a conversation instead of re-prefilling it

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"},

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

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

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<common_adapter_lora_info> lora;
@@ -1547,18 +1558,41 @@ 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");
 
             prompt_cache = std::make_unique<server_prompt_cache>(params_base.cache_ram_mib, n_ctx);
         } 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<server_task_result_cmpl_final>();
 
         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<std::string>() : 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_res_generator> server_routes::handle_completions_impl(
             const server_http_req & req,
             server_task_type type,
@@ -4504,6 +4633,10 @@ std::unique_ptr<server_res_generator> 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,