diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index a4713a88..1a3ff43b 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -2663,6 +2663,9 @@ private:
snapshots.push_back((int32_t) ckpt.n_tokens);
}
}
+ for (const auto * ckpt : prompt_cache->learned_checkpoints(prompt.tokens, tokens, n_state)) {
+ snapshots.push_back((int32_t) ckpt->n_tokens);
+ }
std::sort(snapshots.begin(), snapshots.end());
common_scr_params params;
@@ -2825,6 +2828,13 @@ private:
ckpt_adopt = &ckpt;
}
}
+ if (plan.restore > 0 && ckpt_restore == nullptr) {
+ for (const auto * ckpt : prompt_cache->learned_checkpoints(slot.prompt.tokens, tokens, n_state)) {
+ if (ckpt->n_tokens == plan.restore) {
+ ckpt_restore = ckpt;
+ }
+ }
+ }
if (plan.restore > 0 && ckpt_restore == nullptr) {
return -1;
}
@@ -2929,20 +2939,29 @@ 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
+ // stores the slot's state after its first n tokens as a learned prefix: the recurrent state as a checkpoint
+ // (a restore point for splices), and the whole state unless a splice has moved other text into the cache
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)) {
+ // a splice's tokens are exact only from its restore point to its first run
+ const bool splicing = slot.scr != nullptr;
+ if (splicing ? n < slot.scr->plan.restore || n > slot.scr->plan.runs.front().new_start
+ : 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) {
+ common_prompt_checkpoint ckpt;
+ ckpt.update_pos(n, 0, n - 1);
+ ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+ ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+ common_speculative_get_state(spec.get(), slot.id, ckpt.data_spec);
+
+ const size_t size_tgt = splicing ? 0 : llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE);
+ const size_t size_dft = splicing || !ctx_dft ? 0 : llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE);
+ auto * state = prompt_cache->learn_prefix(slot.prompt.tokens, n, size_tgt, size_dft, std::move(ckpt));
+ if (state == nullptr || splicing) {
return;
}
llama_state_seq_get_data_ext(ctx_tgt, state->data.main.data(), size_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE);
@@ -4454,6 +4473,11 @@ private:
const auto & spans = slot.task->params.message_spans;
const auto last_user_pos = spans.last_user_message_pos();
+ // the previous batch ended at a prefix to learn: store it before anything follows it
+ if (slot.prompt.n_tokens() == slot.learn_at) {
+ learn_prefix(slot, slot.learn_at);
+ }
+
// add prompt tokens for processing in the current batch
while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) {
// a splice takes its next run over once the prefill in front of it is decoded
@@ -4593,10 +4617,6 @@ 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 68cdcc48..fbec0049 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -2350,7 +2350,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
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) {
+ if (!it->data.main.empty() && lcp_cur == (int) it->prompt.tokens.size() && lcp_cur > n_learned) {
it_learned = it;
n_learned = lcp_cur;
}
@@ -3206,7 +3206,9 @@ void server_prompt_cache::disk_init(const std::string & dir, int32_t limit_mib,
}
void server_prompt_cache::evict(std::list<server_prompt_cache_state>::iterator it) {
- if (disk_thread.joinable() && !expired(it->prompt.meta, unix_now())) {
+ // 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 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);
@@ -3580,12 +3582,26 @@ int32_t server_prompt_cache::observe_prefix(const server_tokens & tokens, int32_
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) {
+std::vector<const common_prompt_checkpoint *> server_prompt_cache::learned_checkpoints(const server_tokens & a, const server_tokens & b, int32_t n_max) const {
+ std::vector<const common_prompt_checkpoint *> res;
+ for (const auto & s : states) {
+ const auto & t = s.prompt.tokens;
+ if (!is_learned(s.prompt.meta) || s.prompt.checkpoints.empty() || (int32_t) t.size() > n_max ||
+ t.size() > a.size() || t.size() > b.size() || t.get_common_prefix(a) != t.size() || t.get_common_prefix(b) != t.size()) {
+ continue;
+ }
+ res.push_back(&s.prompt.checkpoints.back());
+ }
+ return res;
+}
+
+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,
+ common_prompt_checkpoint ckpt) {
const uint64_t h = prefix_hash(tokens, n_prefix);
if (learned_covers(tokens, n_prefix)) {
return nullptr;
}
- const size_t size_new = size_tgt + size_dft;
+ const size_t size_new = size_tgt + size_dft + ckpt.size();
if (limit_size > 0 && size_new > limit_size) {
return nullptr;
}
@@ -3608,6 +3624,7 @@ server_prompt_cache_state * server_prompt_cache::learn_prefix(const server_token
server_prompt prompt;
prompt.tokens = tokens.clone();
prompt.tokens.keep_first(n_prefix);
+ prompt.checkpoints.push_back(std::move(ckpt));
const int64_t now = unix_now();
auto & m = prompt.meta;
@@ -3630,7 +3647,8 @@ 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) });
- 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));
+ SRV_INF("prompt cache: learned a %d-token prefix (%s, %s), %.1f MiB\n", n_prefix, hex_id(id).c_str(),
+ size_tgt > 0 ? "state and checkpoint" : "checkpoint only", size_new / (1024.0 * 1024.0));
auto * res = &states.back();
learned_cap(id);
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 68e6f22a..39499f9d 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -799,7 +799,12 @@ struct server_prompt_cache {
// 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);
+ // size_tgt == 0 stores only ckpt, the recurrent state at the prefix end: a restore point for splices
+ server_prompt_cache_state * learn_prefix(const server_tokens & tokens, int32_t n_prefix, size_t size_tgt, size_t size_dft,
+ common_prompt_checkpoint ckpt);
+
+ // checkpoints of learned prefixes (in RAM) that start both a and b, at most n_max tokens long
+ std::vector<const common_prompt_checkpoint *> learned_checkpoints(const server_tokens & a, const server_tokens & b, int32_t n_max) const;
// 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;