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<uint8_t> meta_encode(const cache_meta & m) {
}
bool meta_decode(const std::vector<uint8_t> & b, cache_meta & m) {
- if (b.size() != META_SIZE || get_le<uint8_t>(b, 8) > CACHE_CLASS_PIN || get_le<uint8_t>(b, 9) > CACHE_SOURCE_EXPLICIT) {
+ if (b.size() != META_SIZE || get_le<uint8_t>(b, 8) > CACHE_CLASS_PIN || get_le<uint8_t>(b, 9) > CACHE_SOURCE_LEARNED) {
return false;
}
m.id = get_le<uint64_t>(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<int32_t(const ser
}
bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::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<server_prompt_cache_state>::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<std::mutex> 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<std::mutex> 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<uint64_t, int64_t> 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<std::mutex> lock(disk_mutex);
+ for (const auto & s : pending) {
+ add(s.prompt.meta);
+ }
+ for (const auto & e : disk) {
+ add(e.meta);
+ }
+ }
+ std::vector<std::pair<int64_t, uint64_t>> 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<size_t>(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<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);
+ 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