diff --git a/common/arg.cpp b/common/arg.cpp
index 63f3a7b0..a036b44e 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1709,6 +1709,21 @@ 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(
+ {"--cache-disk"}, "PATH",
+ "directory for a disk tier under the prompt cache: prompts evicted from RAM are written there and\n"
+ "restored on a prefix match, also after a restart (default: off)",
+ [](common_params & params, const std::string & value) {
+ params.cache_disk_path = value;
+ }
+ ).set_env("LLAMA_ARG_CACHE_DISK").set_examples({LLAMA_EXAMPLE_SERVER}));
+ add_opt(common_arg(
+ {"--cache-disk-size"}, "N",
+ string_format("size limit of the prompt cache's disk tier in MiB (default: %d, 0 - no limit)", params.cache_disk_mib),
+ [](common_params & params, int value) {
+ params.cache_disk_mib = value;
+ }
+ ).set_env("LLAMA_ARG_CACHE_DISK_SIZE").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 d796448d..e70f2a99 100644
--- a/common/common.h
+++ b/common/common.h
@@ -619,6 +619,8 @@ struct common_params {
int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot
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.
+ std::string cache_disk_path; // directory of the prompt cache's disk tier, empty = off
+ int32_t cache_disk_mib = 51200; // disk tier size limit in MiB, 0 = no limit
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 2d42450a..d469635c 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -890,6 +890,14 @@ private:
int64_t t_last_load_progress_ms = 0;
void destroy() {
+ // the slots and the RAM cache go to the disk tier while the contexts can still read them
+ if (prompt_cache) {
+ for (auto & slot : slots) {
+ slot.prompt_save(*prompt_cache);
+ }
+ prompt_cache->flush();
+ }
+
spec.reset();
spec_init.reset();
@@ -1309,6 +1317,19 @@ private:
SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n");
prompt_cache = std::make_unique<server_prompt_cache>(params_base.cache_ram_mib, n_ctx);
+
+ if (!params_base.cache_disk_path.empty()) {
+ // a file is only valid for the same weights (path, size, mtime) and cache types
+ std::error_code ec;
+ const auto mtime = std::filesystem::last_write_time(params_base.model.path, ec).time_since_epoch().count();
+ char desc[256];
+ llama_model_desc(model_tgt, desc, sizeof(desc));
+ const std::string key = string_format("%s|%s|%llu|%llu|%lld|%s|%s|%s|%s", params_base.model.path.c_str(), desc,
+ (unsigned long long) llama_model_size(model_tgt), (unsigned long long) llama_model_n_params(model_tgt), (long long) mtime,
+ ggml_type_name(params_base.cache_type_k), ggml_type_name(params_base.cache_type_v),
+ ggml_type_name(params_base.speculative.draft.cache_type_k), ggml_type_name(params_base.speculative.draft.cache_type_v));
+ prompt_cache->disk_init(params_base.cache_disk_path, params_base.cache_disk_mib, key);
+ }
} else {
SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n");
}
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index 0d3beb31..f858a541 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -10,6 +10,9 @@
#include "speculative.h"
#include "server-common.h"
+#include <algorithm>
+#include <cstdio>
+#include <filesystem>
#include <sstream>
//
@@ -1688,6 +1691,99 @@ json server_task_result_apply_lora::to_json() {
//
// server_prompt_cache
//
+namespace {
+
+constexpr uint32_t DISK_MAGIC = 0x43564b42; // "BKVC"
+constexpr uint32_t DISK_VERSION = 1;
+
+uint64_t fnv1a(const void * data, size_t n, uint64_t h = 1469598103934665603ull) {
+ const auto * p = static_cast<const uint8_t *>(data);
+ for (size_t i = 0; i < n; ++i) {
+ h = (h ^ p[i]) * 1099511628211ull;
+ }
+ return h;
+}
+
+struct disk_file {
+ FILE * f;
+ disk_file(const std::string & path, const char * mode) : f(fopen(path.c_str(), mode)) {}
+ ~disk_file() {
+ if (f) {
+ fclose(f);
+ }
+ }
+};
+
+bool put(FILE * f, const void * p, size_t n) {
+ return n == 0 || fwrite(p, 1, n, f) == n;
+}
+
+template <typename T> bool put(FILE * f, const T & v) {
+ return put(f, &v, sizeof(v));
+}
+
+bool put_bytes(FILE * f, const std::vector<uint8_t> & v) {
+ return put(f, (uint64_t) v.size()) && put(f, v.data(), v.size());
+}
+
+bool get(FILE * f, void * p, size_t n) {
+ return n == 0 || fread(p, 1, n, f) == n;
+}
+
+template <typename T> bool get(FILE * f, T & v) {
+ return get(f, &v, sizeof(v));
+}
+
+bool get_bytes(FILE * f, std::vector<uint8_t> & v) {
+ uint64_t n;
+ if (!get(f, n) || n > (1ull << 40)) {
+ return false;
+ }
+ v.resize(n);
+ return get(f, v.data(), n);
+}
+
+// magic, version, key and the prompt tokens: enough to index a file without reading its state
+bool read_header(FILE * f, uint64_t key, llama_tokens & tokens) {
+ uint32_t magic;
+ uint32_t version;
+ uint64_t k;
+ uint64_t n;
+ if (!get(f, magic) || !get(f, version) || !get(f, k) || !get(f, n)) {
+ return false;
+ }
+ if (magic != DISK_MAGIC || version != DISK_VERSION || k != key || n > (1ull << 26)) {
+ return false;
+ }
+ tokens.resize(n);
+ return get(f, tokens.data(), n * sizeof(llama_token));
+}
+
+size_t common_prefix(const llama_tokens & a, const server_tokens & b) {
+ const size_t n = std::min(a.size(), b.size());
+ size_t i = 0;
+ while (i < n && b[i] != LLAMA_TOKEN_NULL && a[i] == b[i]) {
+ ++i;
+ }
+ return i;
+}
+
+// the token list, media positions as LLAMA_TOKEN_NULL (get_tokens() refuses a multimodal server)
+llama_tokens raw_tokens(const server_tokens & t) {
+ llama_tokens res(t.size());
+ for (size_t i = 0; i < t.size(); ++i) {
+ res[i] = t[i];
+ }
+ return res;
+}
+
+// a is a prefix of b
+bool is_prefix(const llama_tokens & a, const llama_tokens & b) {
+ return a.size() <= b.size() && std::equal(a.begin(), a.end(), b.begin());
+}
+
+} // namespace
+
size_t server_prompt_cache::size() const {
size_t res = 0;
@@ -1753,7 +1849,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
states.front().size() / (1024.0 * 1024.0));
- states.pop_front();
+ evict_front();
}
}
@@ -1822,6 +1918,60 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
}
}
+ // the disk tier, under the same rule: a file wins only over a better RAM entry
+ std::string disk_best;
+ if (disk_thread.joinable()) {
+ std::lock_guard<std::mutex> lock(disk_mutex);
+
+ for (const auto & e : disk) {
+ const int lcp_cur = common_prefix(e.tokens, tokens_new);
+
+ const float f_keep_cur = float(lcp_cur) / e.tokens.size();
+ const float f_sim_cur = float(lcp_cur) / tokens_new.size();
+
+ if (f_keep_cur < 0.25f) {
+ continue;
+ }
+
+ if (f_keep_best < f_keep_cur && f_sim_best < f_sim_cur) {
+ f_keep_best = f_keep_cur;
+ f_sim_best = f_sim_cur;
+
+ disk_best = e.path;
+ it_best = states.end();
+ }
+ }
+ }
+
+ if (!disk_best.empty()) {
+ namespace fs = std::filesystem;
+
+ const int64_t t_start = ggml_time_us();
+ server_prompt_cache_state state;
+ const bool ok = disk_read(disk_best, tokens_new.has_mtmd, state);
+
+ std::lock_guard<std::mutex> lock(disk_mutex);
+ std::error_code ec;
+ auto it = std::find_if(disk.begin(), disk.end(), [&](const disk_entry & e) { return e.path == disk_best; });
+ if (ok) {
+ SRV_INF("prompt cache: read %zu tokens from disk in %.2f s, f_keep = %.3f, f_sim = %.3f\n",
+ state.prompt.tokens.size(), (ggml_time_us() - t_start) / 1e6, f_keep_best, f_sim_best);
+ states.push_back(std::move(state));
+ it_best = std::prev(states.end());
+ // most recently used: last to be evicted, here and after a restart
+ if (it != disk.end()) {
+ disk.splice(disk.end(), disk, it);
+ }
+ fs::last_write_time(disk_best, fs::file_time_type::clock::now(), ec);
+ } else {
+ SRV_WRN("prompt cache: cannot read %s, removing it\n", disk_best.c_str());
+ if (it != disk.end()) {
+ disk.erase(it);
+ }
+ fs::remove(disk_best, ec);
+ }
+ }
+
if (it_best != states.end()) {
SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best);
@@ -1872,7 +2022,7 @@ void server_prompt_cache::update() {
while (!states.empty() && size() > limit_size) {
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
- states.pop_front();
+ evict_front();
}
}
@@ -1887,7 +2037,7 @@ void server_prompt_cache::update() {
SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
- states.pop_front();
+ evict_front();
}
}
@@ -1899,3 +2049,224 @@ void server_prompt_cache::update() {
(const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
}
}
+
+//
+// server_prompt_cache: disk tier
+//
+
+
+server_prompt_cache::~server_prompt_cache() {
+ if (disk_thread.joinable()) {
+ {
+ std::lock_guard<std::mutex> lock(disk_mutex);
+ disk_stop = true;
+ }
+ disk_cv.notify_all();
+ disk_thread.join();
+ }
+}
+
+void server_prompt_cache::disk_init(const std::string & dir, int32_t limit_mib, const std::string & key) {
+ namespace fs = std::filesystem;
+
+ disk_dir = dir;
+ disk_limit = 1024ull*1024ull*(size_t) std::max(limit_mib, 0);
+ disk_key = fnv1a(key.data(), key.size());
+
+ std::error_code ec;
+ fs::create_directories(dir, ec);
+ if (ec) {
+ SRV_ERR("prompt cache: cannot create %s: %s, disk tier off\n", dir.c_str(), ec.message().c_str());
+ return;
+ }
+
+ std::vector<std::pair<fs::file_time_type, disk_entry>> found;
+ for (const auto & de : fs::directory_iterator(dir, ec)) {
+ const auto ext = de.path().extension();
+ if (ext == ".tmp") {
+ fs::remove(de.path(), ec);
+ continue;
+ }
+ if (ext != ".kvc") {
+ continue;
+ }
+ llama_tokens tokens;
+ bool ok;
+ {
+ disk_file f(de.path().string(), "rb");
+ ok = f.f && read_header(f.f, disk_key, tokens);
+ }
+ if (!ok) {
+ // another model, cache layout or format
+ fs::remove(de.path(), ec);
+ continue;
+ }
+ found.push_back({ de.last_write_time(ec), { de.path().string(), std::move(tokens), (size_t) de.file_size(ec) } });
+ }
+ std::sort(found.begin(), found.end(), [](const auto & a, const auto & b) { return a.first < b.first; });
+
+ size_t total = 0;
+ for (auto & [_, e] : found) {
+ total += e.size;
+ disk.push_back(std::move(e));
+ }
+ while (disk_limit > 0 && total > disk_limit && !disk.empty()) {
+ total -= disk.front().size;
+ fs::remove(disk.front().path, ec);
+ disk.pop_front();
+ }
+
+ SRV_INF("prompt cache: disk tier %s, %zu prompts, %.1f GiB (limit %.1f GiB)\n",
+ dir.c_str(), disk.size(), total / (1024.0*1024.0*1024.0), disk_limit / (1024.0*1024.0*1024.0));
+
+ disk_thread = std::thread([this] { disk_write_loop(); });
+}
+
+void server_prompt_cache::evict_front() {
+ if (disk_thread.joinable()) {
+ // text prompts only: media chunks are not serialized
+ const llama_tokens tokens = raw_tokens(states.front().prompt.tokens);
+ if (std::find(tokens.begin(), tokens.end(), LLAMA_TOKEN_NULL) == tokens.end()) {
+ std::lock_guard<std::mutex> lock(disk_mutex);
+ pending.splice(pending.end(), states, states.begin());
+ disk_cv.notify_all();
+ return;
+ }
+ }
+ states.pop_front();
+}
+
+void server_prompt_cache::flush() {
+ if (!disk_thread.joinable()) {
+ return;
+ }
+ while (!states.empty()) {
+ evict_front();
+ }
+ std::unique_lock<std::mutex> lock(disk_mutex);
+ disk_cv.wait(lock, [this] { return pending.empty(); });
+}
+
+void server_prompt_cache::disk_write_loop() {
+ std::unique_lock<std::mutex> lock(disk_mutex);
+ while (true) {
+ disk_cv.wait(lock, [this] { return disk_stop || !pending.empty(); });
+ if (pending.empty()) {
+ return;
+ }
+ auto & state = pending.front();
+ lock.unlock();
+ disk_write(state);
+ lock.lock();
+ pending.pop_front();
+ disk_cv.notify_all();
+ }
+}
+
+void server_prompt_cache::disk_write(server_prompt_cache_state & state) {
+ namespace fs = std::filesystem;
+
+ const llama_tokens tokens = raw_tokens(state.prompt.tokens);
+ {
+ std::lock_guard<std::mutex> lock(disk_mutex);
+ for (const auto & e : disk) {
+ if (is_prefix(tokens, e.tokens)) {
+ return; // a file already holds this prompt or a longer one
+ }
+ }
+ }
+
+ const int64_t t_start = ggml_time_us();
+ const uint64_t h = fnv1a(tokens.data(), tokens.size() * sizeof(llama_token));
+ const std::string path = string_format("%s/%016llx-%zu.kvc", disk_dir.c_str(), (unsigned long long) h, tokens.size());
+ const std::string tmp = path + ".tmp";
+
+ bool ok;
+ {
+ disk_file f(tmp, "wb");
+ ok = f.f != nullptr;
+ if (ok) {
+ setvbuf(f.f, nullptr, _IOFBF, 8u << 20);
+ ok = put(f.f, DISK_MAGIC) && put(f.f, DISK_VERSION) && put(f.f, disk_key) &&
+ put(f.f, (uint64_t) tokens.size()) && put(f.f, tokens.data(), tokens.size() * sizeof(llama_token)) &&
+ put(f.f, (uint64_t) state.prompt.checkpoints.size());
+ for (const auto & ckpt : state.prompt.checkpoints) {
+ ok = ok && put(f.f, ckpt.n_tokens) && put(f.f, ckpt.pos_min) && put(f.f, ckpt.pos_max) &&
+ put_bytes(f.f, ckpt.data_tgt) && put_bytes(f.f, ckpt.data_dft) && put_bytes(f.f, ckpt.data_spec);
+ }
+ ok = ok && put_bytes(f.f, state.data.main) && put_bytes(f.f, state.data.drft) && put(f.f, DISK_MAGIC);
+ ok = fflush(f.f) == 0 && ok;
+ }
+ }
+ std::error_code ec;
+ if (ok) {
+ fs::rename(tmp, path, ec);
+ ok = !ec;
+ }
+ if (!ok) {
+ SRV_ERR("prompt cache: failed to write %s\n", path.c_str());
+ fs::remove(tmp, ec);
+ return;
+ }
+ const size_t size = fs::file_size(path, ec);
+
+ std::lock_guard<std::mutex> lock(disk_mutex);
+
+ // files whose prompt this one extends are obsolete
+ for (auto it = disk.begin(); it != disk.end();) {
+ if (is_prefix(it->tokens, tokens) && it->path != path) {
+ fs::remove(it->path, ec);
+ it = disk.erase(it);
+ } else {
+ ++it;
+ }
+ }
+ disk.remove_if([&](const disk_entry & e) { return e.path == path; });
+ disk.push_back({ path, tokens, size });
+
+ size_t total = 0;
+ for (const auto & e : disk) {
+ total += e.size;
+ }
+ while (disk_limit > 0 && total > disk_limit && disk.size() > 1) {
+ total -= disk.front().size;
+ fs::remove(disk.front().path, ec);
+ disk.pop_front();
+ }
+
+ SRV_INF("prompt cache: wrote %zu tokens to disk, %.1f MiB in %.2f s (%zu prompts, %.1f GiB)\n",
+ tokens.size(), size / (1024.0*1024.0), (ggml_time_us() - t_start) / 1e6, disk.size(), total / (1024.0*1024.0*1024.0));
+}
+
+bool server_prompt_cache::disk_read(const std::string & path, bool has_mtmd, server_prompt_cache_state & out) const {
+ disk_file f(path, "rb");
+ if (!f.f) {
+ return false;
+ }
+ setvbuf(f.f, nullptr, _IOFBF, 8u << 20);
+
+ llama_tokens tokens;
+ uint64_t n_ckpt;
+ if (!read_header(f.f, disk_key, tokens) || !get(f.f, n_ckpt) || n_ckpt > 4096) {
+ return false;
+ }
+ std::list<common_prompt_checkpoint> checkpoints;
+ for (uint64_t i = 0; i < n_ckpt; ++i) {
+ common_prompt_checkpoint ckpt;
+ if (!get(f.f, ckpt.n_tokens) || !get(f.f, ckpt.pos_min) || !get(f.f, ckpt.pos_max) ||
+ !get_bytes(f.f, ckpt.data_tgt) || !get_bytes(f.f, ckpt.data_dft) || !get_bytes(f.f, ckpt.data_spec)) {
+ return false;
+ }
+ checkpoints.push_back(std::move(ckpt));
+ }
+ server_prompt_data data;
+ uint32_t end;
+ if (!get_bytes(f.f, data.main) || !get_bytes(f.f, data.drft) || !get(f.f, end) || end != DISK_MAGIC) {
+ return false;
+ }
+
+ out.prompt.tokens = server_tokens(tokens, has_mtmd);
+ out.prompt.checkpoints = std::move(checkpoints);
+ out.data = std::move(data);
+ return true;
+}
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 9c99143f..b8d0b9ca 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -7,6 +7,9 @@
#include <unordered_set>
#include <list>
#include <map>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
// TODO: prevent including the whole server-common.h as we only use server_tokens
#include "server-common.h"
@@ -615,6 +618,8 @@ struct server_prompt_cache {
this->limit_tokens = limit_tokens;
}
+ ~server_prompt_cache();
+
std::list<server_prompt_cache_state> states;
// in bytes, 0 = no limit
@@ -632,6 +637,38 @@ 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();
+
+ // second tier on disk: an entry the RAM limit evicts is written to dir instead of dropped, and read back
+ // when it is the best prefix match for a new prompt; flush() writes what is still in RAM (on shutdown).
+ // key names the model and cache layout; files written under another key are deleted. text prompts only.
+ void disk_init(const std::string & dir, int32_t limit_mib, const std::string & key);
+
+ void flush();
+
+private:
+ struct disk_entry {
+ std::string path;
+ llama_tokens tokens;
+ size_t size;
+ };
+
+ std::string disk_dir;
+ size_t disk_limit = 0;
+ uint64_t disk_key = 0;
+
+ // oldest first; both lists and stop are guarded by disk_mutex
+ std::list<disk_entry> disk;
+ std::list<server_prompt_cache_state> pending;
+
+ std::mutex disk_mutex;
+ std::condition_variable disk_cv;
+ std::thread disk_thread;
+ bool disk_stop = false;
+
+ void evict_front();
+ void disk_write_loop();
+ void disk_write(server_prompt_cache_state & state);
+ bool disk_read(const std::string & path, bool has_mtmd, server_prompt_cache_state & out) const;
};
// used exclusively by router mode