diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index eae1becd..a4713a88 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -305,7 +305,7 @@ struct server_slot {
// the request that produced the prompt asked for "cache": "skip"
bool cache_skip = false;
- // prompt position of a repeated system prefix to store for the prompt cache, -1 = none
+ // length of a shared system prefix to store for the prompt cache, -1 = none
int32_t learn_at = -1;
// a splice in progress (suffix cache reuse): the plan, the next run to take over, and the recurrent state the
@@ -2058,9 +2058,8 @@ private:
prompt_cache->apply_request(slot.prompt.meta, continuation, task.params, task.tokens.size());
slot.cache_skip = task.params.cache_explicit && task.params.cache_cls == CACHE_CLASS_SKIP;
slot.learn_at = -1;
- const int32_t n_system = first_user_pos(task.params.message_spans);
- if (!slot.cache_skip && prompt_cache->observe_prefix(task.tokens, n_system)) {
- slot.learn_at = n_system;
+ if (!slot.cache_skip) {
+ slot.learn_at = prompt_cache->observe_prefix(task.tokens, first_user_pos(task.params.message_spans));
}
SLT_INF(slot, "cache: %s, class %s (%s), hits %llu%s\n", continuation ? "continues the slot prompt" : "new prompt",
cache_class_name(meta.cls), cache_source_name(meta.source), (unsigned long long) meta.hits, slot.cache_skip ? ", skip" : "");
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index 10207afb..68cdcc48 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1764,10 +1764,12 @@ 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; // requests that must repeat a prefix before it is pinned
+constexpr int32_t LEARN_MIN_N = 3; // 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_SEEN_MAX = 256; // prefixes counted at once
+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
int64_t unix_now() {
return (int64_t) std::time(nullptr);
@@ -3475,11 +3477,14 @@ void server_prompt_cache::disk_write(server_prompt_cache_state & state) {
disk.size(), total / (1024.0*1024.0*1024.0));
}
-bool server_prompt_cache::learned_has(uint64_t hash) {
+bool server_prompt_cache::learned_covers(const server_tokens & tokens, int32_t n) {
const int64_t now = unix_now();
+ auto covers = [&](const server_tokens & t) {
+ return (int32_t) t.size() >= n && t.size() <= tokens.size() && t.get_common_prefix(tokens) == t.size();
+ };
bool found = false;
for (auto & s : states) {
- if (is_learned(s.prompt.meta) && s.prompt.meta.family == hash) {
+ if (is_learned(s.prompt.meta) && covers(s.prompt.tokens)) {
s.prompt.meta.last_used = now;
found = true;
}
@@ -3487,10 +3492,10 @@ bool server_prompt_cache::learned_has(uint64_t hash) {
if (disk_thread.joinable()) {
std::lock_guard<std::mutex> lock(disk_mutex);
for (auto & s : pending) {
- found = found || (is_learned(s.prompt.meta) && s.prompt.meta.family == hash);
+ found = found || (is_learned(s.prompt.meta) && covers(s.prompt.tokens));
}
for (auto & e : disk) {
- if (is_learned(e.meta) && e.meta.family == hash) {
+ if (is_learned(e.meta) && covers(e.tokens)) {
e.meta.last_used = now;
found = true;
}
@@ -3531,34 +3536,53 @@ void server_prompt_cache::learned_cap(uint64_t keep) {
}
}
-bool server_prompt_cache::observe_prefix(const server_tokens & tokens, int32_t n_prefix) {
- if (n_prefix < LEARN_MIN_TOKENS || n_prefix >= (int32_t) tokens.size()) {
- return false;
+int32_t server_prompt_cache::observe_prefix(const server_tokens & tokens, int32_t n_system) {
+ if (n_system < LEARN_MIN_TOKENS || n_system >= (int32_t) tokens.size()) {
+ return -1;
}
- for (int32_t i = 0; i < n_prefix; ++i) {
- if (tokens[i] == LLAMA_TOKEN_NULL) {
- return false;
+ const size_t n_head = std::min<size_t>(n_system, LEARN_HEAD_MAX);
+ llama_tokens head(n_head);
+ for (size_t i = 0; i < n_head; ++i) {
+ head[i] = tokens[i];
+ if (head[i] == LLAMA_TOKEN_NULL) {
+ return -1;
}
}
- const uint64_t h = prefix_hash(tokens, n_prefix);
- if (learned_has(h)) {
- return false;
+ // requests of one conversation repeat its first user message; only other prompts count
+ const uint64_t key = prefix_hash(tokens, std::min<size_t>(tokens.size(), n_system + LEARN_KEY_USER_TOKENS));
+
+ // the longest prefix this prompt shares with at least LEARN_MIN_N - 1 other prompts
+ std::vector<int32_t> shared;
+ bool seen = false;
+ for (const auto & p : prompts_seen) {
+ if (p.key == key) {
+ seen = true;
+ continue;
+ }
+ const size_t n = std::min(head.size(), p.head.size());
+ const size_t lcp = std::mismatch(head.begin(), head.begin() + n, p.head.begin()).first - head.begin();
+ shared.push_back((int32_t) lcp);
}
- const int64_t now = unix_now();
- auto & seen = prefix_seen[h];
- seen.first++;
- seen.second = now;
- const bool ready = seen.first >= LEARN_MIN_N;
- if (prefix_seen.size() > LEARN_SEEN_MAX) {
- auto oldest = std::min_element(prefix_seen.begin(), prefix_seen.end(), [](const auto & a, const auto & b) { return a.second.second < b.second.second; });
- prefix_seen.erase(oldest);
+ if (!seen) {
+ prompts_seen.push_back({ key, std::move(head) });
+ if (prompts_seen.size() > LEARN_SEEN_MAX) {
+ prompts_seen.pop_front();
+ }
}
- return ready;
+ if ((int32_t) shared.size() < LEARN_MIN_N - 1) {
+ return -1;
+ }
+ std::nth_element(shared.begin(), shared.begin() + (LEARN_MIN_N - 2), shared.end(), std::greater<int32_t>());
+ const int32_t n = shared[LEARN_MIN_N - 2];
+ if (n < LEARN_MIN_TOKENS || learned_covers(tokens, n)) {
+ return -1;
+ }
+ return n;
}
server_prompt_cache_state * server_prompt_cache::learn_prefix(const server_tokens & tokens, int32_t n_prefix, size_t size_tgt, size_t size_dft) {
const uint64_t h = prefix_hash(tokens, n_prefix);
- if (learned_has(h)) {
+ if (learned_covers(tokens, n_prefix)) {
return nullptr;
}
const size_t size_new = size_tgt + size_dft;
@@ -3606,7 +3630,6 @@ server_prompt_cache_state * server_prompt_cache::learn_prefix(const server_token
const uint64_t id = m.id;
states.push_back({ std::move(prompt), std::move(data) });
- prefix_seen.erase(h);
SRV_INF("prompt cache: learned a %d-token prefix (%s), %.1f MiB\n", n_prefix, hex_id(id).c_str(), size_new / (1024.0 * 1024.0));
auto * res = &states.back();
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index a0241a7a..68e6f22a 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -795,10 +795,10 @@ struct server_prompt_cache {
const cache_policy & policy() const { return pol; }
- // learned prefixes: the tokens before the first user turn, once enough requests repeated them exactly, are
- // pinned as an entry of their own (never the conversation after them). observe_prefix counts a request and
- // returns true when its slot should store the prefix state with learn_prefix at that boundary
- bool observe_prefix(const server_tokens & tokens, int32_t n_prefix);
+ // learned prefixes: the longest start of the system block that three distinct prompts share is pinned as an
+ // entry of its own (never the conversation after it). observe_prefix records a prompt whose first user message
+ // starts at n_system and returns the length to store with learn_prefix, or -1
+ int32_t observe_prefix(const server_tokens & tokens, int32_t n_system);
server_prompt_cache_state * learn_prefix(const server_tokens & tokens, int32_t n_prefix, size_t size_tgt, size_t size_dft);
// telemetry: disk entries read into RAM so far; the disk tier's bytes and limit, false if it is off or busy
@@ -860,9 +860,15 @@ private:
// server loop only
std::map<uint64_t, family_stat> families;
- // requests seen per exact prefix hash: count and last time
- std::map<uint64_t, std::pair<int32_t, int64_t>> prefix_seen;
- bool learned_has(uint64_t hash);
+ // recent distinct prompts: a key that tells conversations apart, and the start of the system block
+ struct prompt_seen {
+ uint64_t key;
+ llama_tokens head;
+ };
+ std::deque<prompt_seen> prompts_seen;
+
+ // 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);
void learned_cap(uint64_t keep); // drops the least recently used learned prefixes beyond the cap, never keep
bool families_dirty = false;