diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index eea5165f..b2cf0f40 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1764,9 +1764,9 @@ constexpr size_t CLASSIFY_QUEUE_MAX = 64; constexpr size_t EXCERPT_CHARS = 80; constexpr size_t INFO_ENTRIES = 200; -constexpr int32_t LEARN_MIN_N = 3; // distinct prompts that must share a prefix before it is pinned +constexpr int32_t LEARN_MIN_N = 2; // distinct prompts that must share a prefix before it is pinned constexpr int32_t LEARN_MIN_TOKENS = 32; -constexpr size_t LEARN_MAX = 8; // learned prefixes kept; the least recently used goes first +constexpr size_t LEARN_MAX = 16; // learned prefixes kept; the least recently used goes first constexpr size_t LEARN_SEEN_MAX = 32; // recent distinct prompts compared constexpr size_t LEARN_HEAD_MAX = 32768; // tokens of each kept for the comparison constexpr size_t LEARN_KEY_USER_TOKENS = 64; // of the first user message: what tells two conversations apart @@ -2490,6 +2490,10 @@ std::string server_prompt_cache::disk_best(const std::function n_best) { n_best = n; @@ -2555,6 +2559,10 @@ std::list::iterator server_prompt_cache::ram_victim(b auto best = states.end(); std::tuple best_key; for (auto it = states.begin(); it != states.end(); ++it) { + // learned prefixes stay in RAM: splices find their checkpoints only there, and LEARN_MAX bounds them + if (is_learned(it->prompt.meta)) { + continue; + } std::tuple key; if (evict_key(it->prompt.meta, now, allow_pin, key) && (best == states.end() || key < best_key)) { best = it; @@ -3060,8 +3068,56 @@ void server_prompt_cache::families_load() { } } +// a count, then per prompt its key, head length and head tokens +void server_prompt_cache::seen_load() { + disk_file f(disk_dir + "/prompts-seen.bin", "rb"); + uint32_t n = 0; + if (!f.f || !get(f.f, n)) { + return; + } + for (uint32_t i = 0; i < n && i < LEARN_SEEN_MAX; ++i) { + prompt_seen p; + uint32_t n_head = 0; + if (!get(f.f, p.key) || !get(f.f, n_head) || n_head > LEARN_HEAD_MAX) { + prompts_seen.clear(); + return; + } + p.head.resize(n_head); + if (n_head > 0 && fread(p.head.data(), sizeof(llama_token), n_head, f.f) != n_head) { + prompts_seen.clear(); + return; + } + prompts_seen.push_back(std::move(p)); + } +} + +void server_prompt_cache::seen_save() { + if (!seen_dirty || disk_dir.empty() || !disk_thread.joinable()) { + return; + } + const std::string path = disk_dir + "/prompts-seen.bin"; + bool ok; + { + disk_file f(path + ".tmp", "wb"); + ok = f.f && put(f.f, (uint32_t) prompts_seen.size()); + for (const auto & p : prompts_seen) { + ok = ok && put(f.f, p.key) && put(f.f, (uint32_t) p.head.size()) && + (p.head.empty() || fwrite(p.head.data(), sizeof(llama_token), p.head.size(), f.f) == p.head.size()); + } + ok = ok && fflush(f.f) == 0; + } + std::error_code ec; + if (ok) { + std::filesystem::rename(path + ".tmp", path, ec); + seen_dirty = false; + } else { + std::filesystem::remove(path + ".tmp", ec); + } +} + // keeps the FAMILIES_MAX most recently seen families void server_prompt_cache::families_save() { + seen_save(); if (families.size() > FAMILIES_MAX) { std::vector seen; for (const auto & [_, f] : families) { @@ -3131,6 +3187,7 @@ void server_prompt_cache::disk_init(const std::string & dir, int32_t limit_mib, } families_load(); + seen_load(); size_t n_v1 = 0; std::vector> found; @@ -3197,12 +3254,29 @@ void server_prompt_cache::disk_init(const std::string & dir, int32_t limit_mib, dir.c_str(), disk.size(), n_v1, n_media, total / (1024.0*1024.0*1024.0), disk_limit / (1024.0*1024.0*1024.0)); disk_thread = std::thread([this] { disk_write_loop(); }); + + // learned prefixes live in RAM; their files are the copy that outlasts a restart + std::vector learned_files; + { + std::lock_guard lock(disk_mutex); + for (const auto & e : disk) { + if (is_learned(e.meta)) { + learned_files.push_back(e.path); + } + } + } + for (const auto & path : learned_files) { + disk_promote(path); + } + if (!learned_files.empty()) { + SRV_INF("prompt cache: read %zu learned prefixes back into RAM\n", learned_files.size()); + } } void server_prompt_cache::evict(std::list::iterator it) { - // a checkpoint without its state cannot be restored on its own - const bool whole = !it->data.main.empty(); - if (disk_thread.joinable() && whole && !expired(it->prompt.meta, unix_now())) { + // a checkpoint without its state is kept on disk only for a learned prefix, which is read back at startup + const bool keep = !it->data.main.empty() || is_learned(it->prompt.meta); + if (disk_thread.joinable() && keep && !expired(it->prompt.meta, unix_now())) { // a prompt with images needs the media key its file is read back under if (disk_media != 0 || !has_media(it->prompt.tokens)) { std::lock_guard lock(disk_mutex); @@ -3564,6 +3638,7 @@ int32_t server_prompt_cache::observe_prefix(const server_tokens & tokens, int32_ if (prompts_seen.size() > LEARN_SEEN_MAX) { prompts_seen.pop_front(); } + seen_dirty = true; } if ((int32_t) shared.size() < LEARN_MIN_N - 1) { return -1; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 39499f9d..2f8c4032 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -871,6 +871,11 @@ private: llama_tokens head; }; std::deque prompts_seen; + bool seen_dirty = false; + + // prompts_seen outlasts restarts in prompts-seen.bin beside families.json + void seen_load(); + void seen_save(); // a learned entry of at least n tokens is a prefix of tokens; marks it used bool learned_covers(const server_tokens & tokens, int32_t n);