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 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(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(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 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()); + 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 families; - // requests seen per exact prefix hash: count and last time - std::map> 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 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;