Commit e152332a1 for llama.cpp
commit e152332a15986f4b15f79396df1980276b984001
Author: cwriter <silvan.niederer@bluewin.ch>
Date: Thu Nov 5 11:02:36 2026 +0100
llama-server: Use reference-counted checkpoints instead of deep-copies (#30137)
* checkpoints: reduce overhead
Assisted-by: Codex
* rebase
Adapt the slot save/restore checkpoint appendix from #26004 to the
shared_ptr checkpoint list and the common_state_data buffers.
Assisted-by: Claude Opus 5.5
---------
Co-authored-by: cwriter <cwriter@localhost>
diff --git a/common/common.h b/common/common.h
index 6ee90c1fe..303e0c293 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1290,6 +1290,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 <typename T>
+struct common_no_init_allocator : std::allocator<T> {
+ template <typename U> struct rebind { using other = common_no_init_allocator<U>; };
+
+ common_no_init_allocator() = default;
+ template <typename U> common_no_init_allocator(const common_no_init_allocator<U> &) {}
+
+ template <typename U> 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<uint8_t, common_no_init_allocator<uint8_t>>;
+
struct common_prompt_checkpoint {
int64_t n_tokens;
@@ -1299,8 +1313,8 @@ struct common_prompt_checkpoint {
llama_pos pos_min;
llama_pos pos_max;
- std::vector<uint8_t> data_tgt;
- std::vector<uint8_t> 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 3a29b99d1..4596e788b 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -358,9 +358,12 @@ struct server_slot {
return false;
}
- llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
- if (ctx_dft) {
- llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);
+ const size_t n_tgt = llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, 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, id, 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;
}
return true;
@@ -2731,21 +2734,22 @@ private:
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);
@@ -2757,8 +2761,8 @@ private:
{
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;
@@ -2766,7 +2770,8 @@ private:
}
}
- auto & cur = slot.prompt.checkpoints.emplace_back();
+ auto ckpt = std::make_shared<common_prompt_checkpoint>();
+ auto & cur = *ckpt;
cur.id_task = id_task;
@@ -2780,6 +2785,8 @@ private:
// stash the draft's speculative state with the checkpoint
common_speculative_get_state(spec.get(), slot.id, 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,
@@ -2799,7 +2806,8 @@ private:
return true;
}
- static bool ckpt_read_buf(std::ifstream & ifs, std::vector<uint8_t> & buf, size_t n_avail, size_t & n_read) {
+ template <typename Buf>
+ static bool ckpt_read_buf(std::ifstream & ifs, Buf & buf, size_t n_avail, size_t & n_read) {
uint64_t n = 0;
// check the size against the bytes left in the file before allocating, the size field may be corrupted
if (!ckpt_read(ifs, &n, sizeof(n), n_read) || n > n_avail - n_read) {
@@ -2814,7 +2822,8 @@ private:
n_written += size;
}
- static void ckpt_write_buf(std::ofstream & ofs, const std::vector<uint8_t> & buf, size_t & n_written) {
+ template <typename Buf>
+ static void ckpt_write_buf(std::ofstream & ofs, const Buf & buf, size_t & n_written) {
const uint64_t n = buf.size();
ckpt_write(ofs, &n, sizeof(n), n_written);
if (n > 0) {
@@ -2839,7 +2848,8 @@ private:
ckpt_write(ofs, &magic, sizeof(magic), n_written);
ckpt_write(ofs, &version, sizeof(version), n_written);
ckpt_write(ofs, &count, sizeof(count), n_written);
- for (const auto & cur : slot.prompt.checkpoints) {
+ for (const auto & ckpt : slot.prompt.checkpoints) {
+ const auto & cur = *ckpt;
ckpt_write(ofs, &cur.n_tokens, sizeof(cur.n_tokens), n_written);
ckpt_write(ofs, &cur.pos_min, sizeof(cur.pos_min), n_written);
ckpt_write(ofs, &cur.pos_max, sizeof(cur.pos_max), n_written);
@@ -2877,9 +2887,10 @@ private:
SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
return 0;
}
- std::list<common_prompt_checkpoint> checkpoints;
+ std::list<std::shared_ptr<common_prompt_checkpoint>> checkpoints;
for (uint32_t i = 0; i < count; ++i) {
- common_prompt_checkpoint cur;
+ auto ckpt = std::make_shared<common_prompt_checkpoint>();
+ auto & cur = *ckpt;
cur.id_task = -1; // not created by a task - marks a checkpoint restored from a slot file
if (!ckpt_read(ifs, &cur.n_tokens, sizeof(cur.n_tokens), n_read) ||
!ckpt_read(ifs, &cur.pos_min, sizeof(cur.pos_min), n_read) ||
@@ -2895,23 +2906,23 @@ private:
SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
return 0;
}
- checkpoints.push_back(std::move(cur));
+ checkpoints.push_back(std::move(ckpt));
if (checkpoints.size() > (size_t) params_base.n_ctx_checkpoints) {
checkpoints.pop_front();
}
}
// the slot file does not check the draft context - test-load one draft checkpoint, drop the draft data if it does not fit
- if (ctx_dft != nullptr && !checkpoints.empty() && !checkpoints.back().data_dft.empty()) {
- const bool ok = checkpoints.back().load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+ if (ctx_dft != nullptr && !checkpoints.empty() && !checkpoints.back()->data_dft.empty()) {
+ const bool ok = checkpoints.back()->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, -1, -1);
if (!ok) {
SLT_WRN(slot, "draft context checkpoint data in '%s' does not match the draft context - dropped\n", filepath.c_str());
for (auto & cur : checkpoints) {
- cur.clear_dft();
+ cur->clear_dft();
}
}
}
- slot.prompt.checkpoints = std::move(checkpoints);
+ slot.prompt.checkpoints.assign(checkpoints.begin(), checkpoints.end());
SLT_INF(slot, "restored %zu context checkpoint(s) from '%s'\n", slot.prompt.checkpoints.size(), filepath.c_str());
return n_read;
}
@@ -4254,7 +4265,8 @@ private:
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]
@@ -4268,10 +4280,12 @@ private:
bool do_reset = it == slot.prompt.checkpoints.rend();
if (!do_reset) {
+ const auto & cur = **it;
+
// restore the context checkpoint
- if (!it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) ||
- !it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) {
- if (it->id_task != -1) {
+ if (!cur.load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) ||
+ !cur.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) {
+ if (cur.id_task != -1) {
GGML_ABORT("failed to restore context checkpoint\n");
}
// restored from a slot file, not guaranteed to load - fall back to full prompt re-processing
@@ -4281,12 +4295,14 @@ private:
}
if (!do_reset) {
+ const auto & cur = **it;
+
// restore the draft's speculative state
- common_speculative_set_state(spec.get(), slot.id, it->data_spec);
+ common_speculative_set_state(spec.get(), slot.id, 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) {
@@ -4301,7 +4317,7 @@ private:
{
// erase any checkpoints with pos_max > pos_next, or that cover tokens not shared with the new prompt
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 || cur.n_tokens > n_prefix) {
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);
@@ -4504,7 +4520,7 @@ private:
const auto pos = slot.prompt.n_tokens();
const auto & checkpoints = slot.prompt.checkpoints;
- if (pos == last_user_pos || checkpoints.empty() || pos > checkpoints.back().n_tokens + params_base.checkpoint_min_step) {
+ if (pos == last_user_pos || checkpoints.empty() || pos > checkpoints.back()->n_tokens + params_base.checkpoint_min_step) {
break;
}
}
@@ -4578,7 +4594,7 @@ private:
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 1a0d434a1..1592330a7 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1767,7 +1767,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;
@@ -1802,8 +1802,8 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
}
}
- std::vector<uint8_t> state_data_tgt;
- std::vector<uint8_t> 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 21b649ac4..8faae1f18 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -686,7 +686,8 @@ struct server_task_result_apply_lora : server_task_result {
struct server_prompt {
server_tokens tokens;
- std::list<common_prompt_checkpoint> 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<std::shared_ptr<const common_prompt_checkpoint>> checkpoints;
void clear() {
tokens.clear();
@@ -706,8 +707,8 @@ struct server_prompt {
};
struct server_prompt_data {
- std::vector<uint8_t> main;
- std::vector<uint8_t> drft;
+ common_state_data main;
+ common_state_data drft;
size_t size() const {
return main.size() + drft.size();
@@ -722,7 +723,7 @@ struct server_prompt_cache_state {
size_t res = data.size();
for (const auto & ckpt : prompt.checkpoints) {
- res += ckpt.size();
+ res += ckpt->size();
}
return res;