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(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 +#include +#include #include // @@ -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(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 bool put(FILE * f, const T & v) { + return put(f, &v, sizeof(v)); +} + +bool put_bytes(FILE * f, const std::vector & 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 bool get(FILE * f, T & v) { + return get(f, &v, sizeof(v)); +} + +bool get_bytes(FILE * f, std::vector & 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 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 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 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> 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 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 lock(disk_mutex); + disk_cv.wait(lock, [this] { return pending.empty(); }); +} + +void server_prompt_cache::disk_write_loop() { + std::unique_lock 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 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 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 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 #include #include +#include +#include +#include // 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 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; + std::list 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