Skip to content

server : shared context checkpoints, no zero-fill of state buffers (port of upstream #30137) #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
routhjim merged 1 commit into from
Oct 9, 2026
Merged

server : shared context checkpoints, no zero-fill of state buffers (port of upstream #30137) #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "llama.h"

#include <list>
#include <memory>
#include <set>
#include <sstream>
#include <string>
Expand Down Expand Up @@ -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 <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;

Expand All @@ -1198,8 +1213,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)
Expand Down
85 changes: 50 additions & 35 deletions tools/server/server-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -2454,16 +2459,17 @@ 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;
}
}
}

auto & cur = slot.prompt.checkpoints.emplace_back();
auto ckpt = std::make_shared<common_prompt_checkpoint>();
auto & cur = *ckpt;

cur.id_task = id_task;

Expand All @@ -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,
Expand Down Expand Up @@ -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;
Expand All @@ -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<uint8_t> * 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) {
Expand Down Expand Up @@ -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<common_prompt_checkpoint>();
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));
Expand All @@ -2870,24 +2882,24 @@ struct server_context_impl {
c.id_task = -1;
c.pos_min = pmin;
c.pos_max = pmax;
std::vector<uint8_t> * 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());
Expand Down Expand Up @@ -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]
Expand All @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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
Expand Down
16 changes: 8 additions & 8 deletions tools/server/server-task.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
}
Expand All @@ -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<uint8_t> & 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;
Expand All @@ -1748,7 +1748,7 @@ static bool write_blob(const std::string & path, const std::vector<uint8_t> & da
return ok;
}

static bool read_blob(const std::string & path, std::vector<uint8_t> & data) {
static bool read_blob(const std::string & path, common_state_data & data) {
FILE * f = fopen(path.c_str(), "rb");
if (!f) {
return false;
Expand Down Expand Up @@ -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<uint8_t> 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;
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<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 {
Expand Down
10 changes: 6 additions & 4 deletions tools/server/server-task.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include <string>
#include <unordered_set>
#include <list>
#include <memory>
#include <map>

// TODO: prevent including the whole server-common.h as we only use server_tokens
Expand Down Expand Up @@ -566,7 +567,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();
Expand All @@ -586,8 +588,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();
Expand All @@ -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;
Expand Down