diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c2f29094..cbfbff6b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -68,6 +68,16 @@ enum slot_state { struct server_slot; // forward declaration +// where the first user message starts, -1 without one +static int32_t first_user_pos(const common_chat_msg_spans & spans) { + for (const auto & span : spans.spans) { + if (span.role == COMMON_CHAT_ROLE_USER) { + return (int32_t) span.pos; + } + } + return -1; +} + // while a prompt is being processed beside another active request, each step runs either the decoding slots // or one prompt chunk of one slot, never both; the choice goes to the member owed the most GPU time // (equal processor sharing: every request weighs 1, a decoding one --fair-decode-weight) @@ -292,6 +302,9 @@ 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 + 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 // last run continues from. while set, the cache holds the moved runs beyond the prefill position struct scr_splice { @@ -2041,6 +2054,11 @@ private: : 4*lcp >= (int64_t) slot.prompt.tokens.size()); 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; + } 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" : ""); } @@ -2909,6 +2927,28 @@ private: return true; } + // stores the slot's exact state of its first n tokens (decoded, the batch after them not yet) as a learned prefix + void learn_prefix(server_slot & slot, int32_t n) { + slot.learn_at = -1; + // a splice's moved runs or approximate state must not become a pinned prefix + if (slot.scr || (slot.prompt.scr_depth > 0 && slot.prompt.scr_exact < n)) { + return; + } + if (slot.can_speculate() && !common_speculative_flush(spec.get(), slot.id)) { + return; + } + const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + auto * state = prompt_cache->learn_prefix(slot.prompt.tokens, n, size_tgt, size_dft); + if (state == nullptr) { + return; + } + llama_state_seq_get_data_ext(ctx_tgt, state->data.main.data(), size_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE); + if (ctx_dft) { + llama_state_seq_get_data_ext(ctx_dft, state->data.drft.data(), size_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE); + } + } + // n_tokens_cur: the number of tokens added to the batch for the current slot void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) { const int id_task = slot.task->id; @@ -4439,6 +4479,11 @@ private: /* is_prompt = */ true); slot.prompt.tokens.push_back(cur_tok); + // a learned prefix is stored before the batch that starts after it + if (slot.prompt.n_tokens() == slot.learn_at) { + break; + } + // break at the last user message, or at user messages at least min step past the last checkpoint. // a last user message only a few tokens past the last checkpoint is cheaper to re-process than to // split the batch for (a short follow-up turn would take three small batches instead of two) @@ -4534,6 +4579,10 @@ private: n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max); + if (n_tokens_start == slot.learn_at) { + learn_prefix(slot, n_tokens_start); + } + // note: we create the checkpoint before calling llama_decode(), so the current batch is not // yet processed and therefore it is not part of the checkpoint. if (do_checkpoint) { diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index fac61f1b..10207afb 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1726,6 +1726,7 @@ const char * cache_source_name(cache_source src) { case CACHE_SOURCE_CLASSIFIER: return "classifier"; case CACHE_SOURCE_FAMILY: return "family"; case CACHE_SOURCE_EXPLICIT: return "explicit"; + case CACHE_SOURCE_LEARNED: return "learned"; } return "default"; } @@ -1763,6 +1764,11 @@ 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_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 + int64_t unix_now() { return (int64_t) std::time(nullptr); } @@ -1938,7 +1944,7 @@ std::vector meta_encode(const cache_meta & m) { } bool meta_decode(const std::vector & b, cache_meta & m) { - if (b.size() != META_SIZE || get_le(b, 8) > CACHE_CLASS_PIN || get_le(b, 9) > CACHE_SOURCE_EXPLICIT) { + if (b.size() != META_SIZE || get_le(b, 8) > CACHE_CLASS_PIN || get_le(b, 9) > CACHE_SOURCE_LEARNED) { return false; } m.id = get_le(b, 0); @@ -2094,6 +2100,18 @@ llama_tokens raw_tokens(const server_tokens & t) { return res; } +bool is_learned(const cache_meta & m) { + return m.source == CACHE_SOURCE_LEARNED; +} + +uint64_t prefix_hash(const server_tokens & t, size_t n) { + llama_tokens raw(n); + for (size_t i = 0; i < n; ++i) { + raw[i] = t[i]; + } + return fnv1a(raw.data(), n * sizeof(llama_token)); +} + bool has_media(const server_tokens & t) { for (size_t i = 0; i < t.size(); ++i) { if (t[i] == LLAMA_TOKEN_NULL) { @@ -2211,6 +2229,9 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro // first check if the current state is contained fully in the cache for (auto it = states.begin(); it != states.end(); ++it) { + if (is_learned(it->prompt.meta)) { + continue; + } const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens); if (cur_lcp_len == (int) prompt.tokens.size()) { @@ -2243,7 +2264,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro for (auto it = states.begin(); it != states.end();) { const int len = it->prompt.tokens.get_common_prefix(prompt.tokens); - if (len == (int) it->prompt.tokens.size()) { + if (len == (int) it->prompt.tokens.size() && !is_learned(it->prompt.meta)) { SRV_TRC(" - removing obsolete cached prompt with length %d\n", len); merge_meta(meta, it->prompt.meta); @@ -2317,10 +2338,23 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok auto it_best = states.end(); + // a learned prefix is a fallback: used only when it covers more of the prompt than any conversation + auto it_learned = states.end(); + std::string disk_learned; + int n_learned = 0; + // find the most similar cached prompt, that would also preserve the most context for (auto it = states.begin(); it != states.end(); ++it) { const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new); + if (is_learned(it->prompt.meta)) { + if (lcp_cur == (int) it->prompt.tokens.size() && lcp_cur > n_learned) { + it_learned = it; + n_learned = lcp_cur; + } + continue; + } + const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size(); const float f_sim_cur = float(lcp_cur) / tokens_new.size(); @@ -2349,6 +2383,15 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok for (const auto & e : disk) { const int lcp_cur = e.tokens.get_common_prefix(tokens_new); + if (is_learned(e.meta)) { + if (lcp_cur == (int) e.tokens.size() && lcp_cur > n_learned) { + it_learned = states.end(); + disk_learned = e.path; + n_learned = lcp_cur; + } + continue; + } + const float f_keep_cur = float(lcp_cur) / e.tokens.size(); const float f_sim_cur = float(lcp_cur) / tokens_new.size(); @@ -2367,6 +2410,12 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok } } + if (n_learned > (int) (f_sim_best * tokens_new.size() + 0.5f)) { + SRV_INF("prompt cache: using a learned %d-token prefix\n", n_learned); + it_best = it_learned; + disk_best = it_learned == states.end() ? disk_learned : std::string(); + } + bool from_disk = false; if (!disk_best.empty()) { const auto it = disk_promote(disk_best); @@ -2449,7 +2498,8 @@ std::string server_prompt_cache::disk_best(const std::function::iterator it, server_prompt & prompt, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, bool keep_entry) { - const bool consume = !keep_entry; + const bool learned = is_learned(it->prompt.meta); + const bool consume = !keep_entry && !learned; { auto & data = it->data.main; @@ -2496,6 +2546,11 @@ bool server_prompt_cache::load_entry(std::list::itera } else { prompt = it->prompt.clone(); } + if (learned) { + it->prompt.meta.hits++; + it->prompt.meta.last_used = unix_now(); + prompt.meta = {}; + } return true; } @@ -3350,10 +3405,17 @@ void server_prompt_cache::disk_write(server_prompt_cache_state & state) { const server_tokens & tokens = state.prompt.tokens; const bool media = has_media(tokens); + const bool learned = is_learned(state.prompt.meta); cache_meta meta; { std::lock_guard lock(disk_mutex); for (auto & e : disk) { + if (learned || is_learned(e.meta)) { + if (learned && is_learned(e.meta) && e.meta.family == state.prompt.meta.family) { + return; // already on disk + } + continue; + } if (is_prefix(tokens, e.tokens)) { // a file already holds this prompt or a longer one if (merge_meta(e.meta, state.prompt.meta)) { @@ -3370,7 +3432,7 @@ void server_prompt_cache::disk_write(server_prompt_cache_state & state) { // the name hashes the media chunks too: prompts with equal text and other images get files of their own const llama_tokens words = disk_tokens(tokens); const uint64_t h = fnv1a(words.data(), words.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 path = string_format("%s/%016llx-%zu%s.kvc", disk_dir.c_str(), (unsigned long long) h, tokens.size(), learned ? "-learned" : ""); if (!write_file(path, disk_key, media ? disk_media : 0, words, state, meta)) { SRV_ERR("prompt cache: failed to write %s\n", path.c_str()); @@ -3386,7 +3448,7 @@ void server_prompt_cache::disk_write(server_prompt_cache_state & state) { // 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) { + if (!learned && !is_learned(it->meta) && is_prefix(it->tokens, tokens) && it->path != path) { merge_meta(latest, it->meta); fs::remove(it->path, ec); it = disk.erase(it); @@ -3413,6 +3475,146 @@ 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) { + const int64_t now = unix_now(); + bool found = false; + for (auto & s : states) { + if (is_learned(s.prompt.meta) && s.prompt.meta.family == hash) { + s.prompt.meta.last_used = now; + found = true; + } + } + 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); + } + for (auto & e : disk) { + if (is_learned(e.meta) && e.meta.family == hash) { + e.meta.last_used = now; + found = true; + } + } + } + return found; +} + +void server_prompt_cache::learned_cap(uint64_t keep) { + // a learned prefix can be in RAM and on disk at once: one record per id, its latest use + std::map used; + auto add = [&](const cache_meta & m) { + if (is_learned(m) && m.id != keep) { + used[m.id] = std::max(used[m.id], m.last_used); + } + }; + for (const auto & s : states) { + add(s.prompt.meta); + } + if (disk_thread.joinable()) { + std::lock_guard lock(disk_mutex); + for (const auto & s : pending) { + add(s.prompt.meta); + } + for (const auto & e : disk) { + add(e.meta); + } + } + std::vector> learned; + for (const auto & [id, last_used] : used) { + learned.push_back({ last_used, id }); + } + std::sort(learned.begin(), learned.end()); + // keep is the newest; the cap counts it + for (size_t i = 0; i + LEARN_MAX <= learned.size(); ++i) { + SRV_INF("prompt cache: dropping learned prefix %s, the least recently used of %zu\n", hex_id(learned[i].second).c_str(), learned.size()); + delete_entry(learned[i].second); + } +} + +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; + } + for (int32_t i = 0; i < n_prefix; ++i) { + if (tokens[i] == LLAMA_TOKEN_NULL) { + return false; + } + } + const uint64_t h = prefix_hash(tokens, n_prefix); + if (learned_has(h)) { + return false; + } + 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); + } + return ready; +} + +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)) { + return nullptr; + } + const size_t size_new = size_tgt + size_dft; + if (limit_size > 0 && size_new > limit_size) { + return nullptr; + } + while (limit_size > 0 && !states.empty() && size() + size_new > limit_size) { + auto victim = ram_victim(disk_thread.joinable()); + if (victim == states.end()) { + return nullptr; + } + evict(victim); + } + + server_prompt_data data; + try { + data.main.resize(size_tgt); + data.drft.resize(size_dft); + } catch (const std::bad_alloc &) { + return nullptr; + } + + server_prompt prompt; + prompt.tokens = tokens.clone(); + prompt.tokens.keep_first(n_prefix); + + const int64_t now = unix_now(); + auto & m = prompt.meta; + m.id = new_entry_id(); + m.cls = CACHE_CLASS_PIN; + m.source = CACHE_SOURCE_LEARNED; + m.created = now; + m.last_used = now; + m.family = h; + m.n_input = n_prefix; + if (detokenize) { + // the system text after its role header + llama_tokens head = raw_tokens(prompt.tokens); + head.resize(std::min(head.size(), 128)); + std::string text = detokenize(head); + const size_t nl = text.rfind("<|", 0) == 0 ? text.find('\n') : std::string::npos; + m.excerpt = user_excerpt(nl == std::string::npos ? text : text.substr(nl + 1)); + } + set_expiry(m); + + 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(); + learned_cap(id); + update_ram_next_expiry(); + return res; +} + json server_prompt_cache::info() { const int64_t now = unix_now(); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index e428a5a4..a0241a7a 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -66,6 +66,7 @@ enum cache_source : uint8_t { CACHE_SOURCE_CLASSIFIER, CACHE_SOURCE_FAMILY, CACHE_SOURCE_EXPLICIT, + CACHE_SOURCE_LEARNED, // a system prefix the cache pinned itself after repeated requests }; const char * cache_class_name(cache_class cls); @@ -794,6 +795,12 @@ 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); + 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 uint64_t n_disk_reads = 0; bool disk_usage(size_t & bytes, size_t & limit) const; @@ -852,6 +859,11 @@ 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); + void learned_cap(uint64_t keep); // drops the least recently used learned prefixes beyond the cap, never keep bool families_dirty = false; // guarded by policy_mutex