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<int32_t(const ser
}
expire_disk_locked();
for (const auto & e : disk) {
+ // learned prefixes are in RAM; their files are only the copy that outlasts a restart
+ if (is_learned(e.meta)) {
+ continue;
+ }
const int32_t n = score(e.tokens);
if (n > n_best) {
n_best = n;
@@ -2555,6 +2559,10 @@ std::list<server_prompt_cache_state>::iterator server_prompt_cache::ram_victim(b
auto best = states.end();
std::tuple<int, int, int64_t> 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<int, int, int64_t> 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<int64_t> 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<std::pair<fs::file_time_type, disk_entry>> 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<std::string> learned_files;
+ {
+ std::lock_guard<std::mutex> 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<server_prompt_cache_state>::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<std::mutex> 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<prompt_seen> 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);