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::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 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 server_prompt_cache::learned_checkpoints(const server_tokens & a, const server_tokens & b, int32_t n_max) const { + std::vector 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 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;