From ec23ac20998427897e3b254c4027b6d143221cb3 Mon Sep 17 00:00:00 2001 From: Jim Routh Date: Fri, 9 Oct 2026 08:28:18 -0500 Subject: [PATCH] server : share context checkpoints with the prompt cache, skip zero-fill of state buffers (port of upstream #30137) Checkpoints become shared_ptr, so saving a slot to the prompt cache no longer deep-copies them; state buffers (checkpoints, cache entries, NVMe reload) use a no-init allocator; prompt_save drops an entry whose state read came back short. Lab sidecar save/restore adapted. Assisted-by: Claude Opus 5.5 --- common/common.h | 19 +++++++- tools/server/server-context.cpp | 85 +++++++++++++++++++-------------- tools/server/server-task.cpp | 16 +++---- tools/server/server-task.h | 10 ++-- 4 files changed, 81 insertions(+), 49 deletions(-) diff --git a/common/common.h b/common/common.h index 19ad5ea17..d45a640f1 100644 --- a/common/common.h +++ b/common/common.h @@ -9,6 +9,7 @@ #include "llama.h" #include +#include #include #include #include @@ -1189,6 +1190,20 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); // prompt utils // +// leave default-inserted elements uninitialized; callers must fill them before use +template +struct common_no_init_allocator : std::allocator { + template struct rebind { using other = common_no_init_allocator; }; + + common_no_init_allocator() = default; + template common_no_init_allocator(const common_no_init_allocator &) {} + + template void construct(U * p) { ::new ((void *) p) U; } +}; + +// sequence state data, filled by llama_state_seq_get_data_ext() +using common_state_data = std::vector>; + struct common_prompt_checkpoint { int64_t n_tokens; @@ -1198,8 +1213,8 @@ struct common_prompt_checkpoint { llama_pos pos_min; llama_pos pos_max; - std::vector data_tgt; - std::vector data_dft; + common_state_data data_tgt; + common_state_data data_dft; // (optional) speculative-decoding implementation state stashed with the checkpoint // (e.g. eagle3's deferred-boundary g_embd row) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 824d00699..8c54f1e45 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -326,9 +326,13 @@ struct server_slot { return false; } - llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE); - if (ctx_dft) { - llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE); + // the buffers are not zero-initialized, so a short read must not reach the cache + const size_t n_tgt = llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t n_dft = ctx_dft ? llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (n_tgt != cur_size_tgt || n_dft != cur_size_dft) { + SLT_ERR(*this, "%s", "failed to save complete prompt state\n"); + prompt_cache.states.pop_back(); + return false; } // buffers are populated: hand the entry to the disk tier so it does not sit in RAM @@ -2428,21 +2432,22 @@ struct server_context_impl { for (auto it = slot.prompt.checkpoints.begin(); slot.prompt.checkpoints.size() + 1 >= (size_t) params_base.n_ctx_checkpoints && it != slot.prompt.checkpoints.end(); ) { - if (it->id_task != id_task && last >= 0 && it->n_tokens <= last + params_base.checkpoint_min_step) { + const auto & cur = **it; + if (cur.id_task != id_task && last >= 0 && cur.n_tokens <= last + params_base.checkpoint_min_step) { SLT_TRC(slot, "erasing context checkpoint too close to an earlier one (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", - it->pos_min, it->pos_max, it->n_tokens, (float) it->size() / 1024 / 1024); + cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); it = slot.prompt.checkpoints.erase(it); continue; } - last = it->n_tokens; + last = cur.n_tokens; ++it; } while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { // make room for the new checkpoint, if needed - const auto & cur = slot.prompt.checkpoints.front(); + const auto & cur = *slot.prompt.checkpoints.front(); SLT_WRN(slot, "erasing old context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); @@ -2454,8 +2459,8 @@ struct server_context_impl { { const int64_t n_tokens_new = slot.prompt.n_tokens() - n_tokens_cur; for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end(); ) { - if (it->n_tokens == n_tokens_new) { - SLT_TRC(slot, "superseding context checkpoint at n_tokens = %" PRId64 "\n", it->n_tokens); + if ((*it)->n_tokens == n_tokens_new) { + SLT_TRC(slot, "superseding context checkpoint at n_tokens = %" PRId64 "\n", (*it)->n_tokens); it = slot.prompt.checkpoints.erase(it); } else { ++it; @@ -2463,7 +2468,8 @@ struct server_context_impl { } } - auto & cur = slot.prompt.checkpoints.emplace_back(); + auto ckpt = std::make_shared(); + auto & cur = *ckpt; cur.id_task = id_task; @@ -2477,6 +2483,8 @@ struct server_context_impl { // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.seq, cur.data_spec); + slot.prompt.checkpoints.push_back(std::move(ckpt)); + SLT_TRC(slot, "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, @@ -2723,7 +2731,8 @@ struct server_context_impl { fckpt.write((const char *) &cnt, sizeof(cnt)); fckpt.write((const char *) &st_size, sizeof(st_size)); fckpt.write((const char *) &st_mtime, sizeof(st_mtime)); - for (const auto & c : slot->prompt.checkpoints) { + for (const auto & ckpt : slot->prompt.checkpoints) { + const auto & c = *ckpt; const int64_t n_tok = c.n_tokens; const int32_t idt = (int32_t) c.id_task; const int32_t pmin = (int32_t) c.pos_min; @@ -2732,14 +2741,16 @@ struct server_context_impl { fckpt.write((const char *) &idt, sizeof(idt)); fckpt.write((const char *) &pmin, sizeof(pmin)); fckpt.write((const char *) &pmax, sizeof(pmax)); - const std::vector * vs[3] = { &c.data_tgt, &c.data_dft, &c.data_spec }; - for (int vi = 0; vi < 3; ++vi) { - const uint64_t n = (uint64_t) vs[vi]->size(); + const auto write_buf = [&](const auto & v) { + const uint64_t n = (uint64_t) v.size(); fckpt.write((const char *) &n, sizeof(n)); if (n) { - fckpt.write((const char *) vs[vi]->data(), n); + fckpt.write((const char *) v.data(), n); } - } + }; + write_buf(c.data_tgt); + write_buf(c.data_dft); + write_buf(c.data_spec); } fckpt.close(); if (fckpt) { @@ -2854,7 +2865,8 @@ struct server_context_impl { cnt, cnt_max); } for (uint32_t i = 0; i < cnt && i < cnt_max && fckpt; ++i) { - common_prompt_checkpoint c; + auto ckpt = std::make_shared(); + auto & c = *ckpt; int64_t n_tok = 0; int32_t idt = 0, pmin = 0, pmax = 0; fckpt.read((char *) &n_tok, sizeof(n_tok)); fckpt.read((char *) &idt, sizeof(idt)); @@ -2870,24 +2882,24 @@ struct server_context_impl { c.id_task = -1; c.pos_min = pmin; c.pos_max = pmax; - std::vector * vs[3] = { &c.data_tgt, &c.data_dft, &c.data_spec }; - bool ok = true; - for (int vi = 0; vi < 3; ++vi) { + const auto read_buf = [&](auto & v) { uint64_t n = 0; fckpt.read((char *) &n, sizeof(n)); // bound the allocation by what the sidecar can actually hold - if (!fckpt || n > ckpt_bytes) { ok = false; break; } - vs[vi]->resize((size_t) n); + if (!fckpt || n > ckpt_bytes) { return false; } + v.resize((size_t) n); if (n) { - fckpt.read((char *) vs[vi]->data(), n); + fckpt.read((char *) v.data(), n); } - } + return true; + }; + bool ok = read_buf(c.data_tgt) && read_buf(c.data_dft) && read_buf(c.data_spec); if (!ok || !fckpt) { SRV_WRN("checkpoint sidecar for %s is truncated; keeping %zu of %u\n", filename.c_str(), slot->prompt.checkpoints.size(), cnt); break; } - slot->prompt.checkpoints.emplace_back(std::move(c)); + slot->prompt.checkpoints.push_back(std::move(ckpt)); } SRV_INF("restored %zu context checkpoint(s) from sidecar\n", slot->prompt.checkpoints.size()); @@ -3624,7 +3636,8 @@ struct server_context_impl { const auto it = std::find_if( slot.prompt.checkpoints.rbegin(), slot.prompt.checkpoints.rend(), - [&](const auto & cur) { + [&](const auto & ckpt) { + const auto & cur = *ckpt; // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS] SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold); // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] @@ -3638,15 +3651,17 @@ struct server_context_impl { bool do_reset = it == slot.prompt.checkpoints.rend(); if (!do_reset) { + const auto & cur = **it; + // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.load_tgt(ctx_tgt, slot.seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.load_dft(ctx_dft, slot.seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state - common_speculative_set_state(spec.get(), slot.seq, it->data_spec); + common_speculative_set_state(spec.get(), slot.seq, cur.data_spec); - pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); - n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); - SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); + pos_next = std::min(pos_next, std::max(cur.pos_min + 1, cur.pos_max)); + n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) cur.n_tokens); + SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_past, (float) cur.size() / 1024 / 1024); } if (do_reset) { @@ -3661,7 +3676,7 @@ struct server_context_impl { { // erase any checkpoints with pos_max > pos_next for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) { - const auto & cur = *it; + const auto & cur = **it; if (cur.pos_max > pos_next) { SLT_TRC(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024); it = slot.prompt.checkpoints.erase(it); @@ -3828,7 +3843,7 @@ struct server_context_impl { // is never restored, and the extra batch it forces costs a full MoE expert sweep (~0.6 s on // Flash-Next). Keep only the min-step breaks then. static const bool user_break = [] { const char * e = getenv("LLAMA_CKPT_USER_BREAK"); return e == nullptr || atoi(e) != 0; }(); - const bool min_step = !checkpoints.empty() && pos > checkpoints.back().n_tokens + params_base.checkpoint_min_step; + const bool min_step = !checkpoints.empty() && pos > checkpoints.back()->n_tokens + params_base.checkpoint_min_step; if (min_step || (user_break && (pos == last_user_pos || checkpoints.empty()))) { break; } @@ -3903,7 +3918,7 @@ struct server_context_impl { do_checkpoint = do_checkpoint && ( slot.prompt.checkpoints.empty() || is_last_user_message || near_prompt_end || - n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); + 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); // note: we create the checkpoint before calling llama_decode(), so the current batch is not diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index caaaebc3d..8b5be0104 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1708,7 +1708,7 @@ size_t server_prompt_cache::ram_size() const { // its checkpoints. Measured 2026-09-09 on a TB2.1 run: 19 spilled entries at // -ctxcp 4 held ~9.7 GiB of anonymous RAM against a --cache-ram budget of 2 GiB. for (const auto & ckpt : state.prompt.checkpoints) { - res += ckpt.size(); + res += ckpt->size(); } } } @@ -1735,7 +1735,7 @@ size_t server_prompt_cache::n_tokens() const { return res; } -static bool write_blob(const std::string & path, const std::vector & data) { +static bool write_blob(const std::string & path, const common_state_data & data) { FILE * f = fopen(path.c_str(), "wb"); if (!f) { return false; @@ -1748,7 +1748,7 @@ static bool write_blob(const std::string & path, const std::vector & da return ok; } -static bool read_blob(const std::string & path, std::vector & data) { +static bool read_blob(const std::string & path, common_state_data & data) { FILE * f = fopen(path.c_str(), "rb"); if (!f) { return false; @@ -1803,7 +1803,7 @@ bool server_prompt_cache::unspill(server_prompt_cache_state & st) { return true; } const int64_t t0 = ggml_time_us(); - std::vector main, drft; + common_state_data main, drft; if (!read_blob(st.spill_main, main) || (!st.spill_drft.empty() && !read_blob(st.spill_drft, drft))) { SRV_WRN(" - prompt cache: reload from disk FAILED (%s)\n", st.spill_main.c_str()); return false; @@ -1854,7 +1854,7 @@ void server_prompt_cache::evict_ram(size_t need) { } size_t freed = 0; for (const auto & c : ck->prompt.checkpoints) { - freed += c.size(); + freed += c->size(); } SRV_WRN(" - prompt cache over --cache-ram, dropping %zu context checkpoint(s) from a spilled entry (%.3f MiB)\n", ck->prompt.checkpoints.size(), freed / (1024.0 * 1024.0)); @@ -1943,7 +1943,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro // calculate checkpoints size to see if it will fit with the prompt size_t checkpoints_size = 0; for (const auto & ckpt : prompt.checkpoints) { - checkpoints_size += ckpt.size(); + checkpoints_size += ckpt->size(); } const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size; @@ -1977,8 +1977,8 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro // make room before allocating the new vectors to avoid breaching the limit (spills to disk when enabled) evict_ram(state_size_new); - std::vector state_data_tgt; - std::vector state_data_dft; + common_state_data state_data_tgt; + common_state_data state_data_dft; // check if we can allocate enough memory for the new state try { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 7741bd328..1d2071f77 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -6,6 +6,7 @@ #include #include #include +#include #include // TODO: prevent including the whole server-common.h as we only use server_tokens @@ -566,7 +567,8 @@ struct server_task_result_apply_lora : server_task_result { struct server_prompt { server_tokens tokens; - std::list checkpoints; + // a checkpoint does not change after it is created, so the prompt cache shares it with the slot instead of a copy + std::list> checkpoints; void clear() { tokens.clear(); @@ -586,8 +588,8 @@ struct server_prompt { }; struct server_prompt_data { - std::vector main; - std::vector drft; + common_state_data main; + common_state_data drft; size_t size() const { return main.size() + drft.size(); @@ -610,7 +612,7 @@ struct server_prompt_cache_state { size_t res = spilled() ? spill_size : data.size(); for (const auto & ckpt : prompt.checkpoints) { - res += ckpt.size(); + res += ckpt->size(); } return res;