This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568diff --git a/common/arg.cpp b/common/arg.cppindex 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.hindex 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 = ""; // NOLINTdiff --git a/tools/server/server-context.cpp b/tools/server/server-context.cppindex 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.cppindex 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.hindex 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