wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 28 Packages Projects Releases Wiki Activity

server : preserve context checkpoints across slot save/restore (#26004)

Browse Source 
* server : preserve context checkpoints across slot save/restore

Append the checkpoints after the packed server_tokens payload added in #26640
and count them in n_written / n_read, so a restored slot can still roll back to
a checkpoint instead of re-processing the whole prompt.

* server : drop draft checkpoint data that does not match the draft context

Restoring a slot saved with a different draft KV cache type aborted in
load_dft(). Test-load one draft checkpoint on restore and drop the draft
data if it does not fit, instead of crashing. Adds a regression test.

Co-authored-by: Igor Okulist <okigan@gmail.com>

* server : harden the checkpoint appendix of slot save files

Bound each blob size by the bytes left in the file before allocating, open the
file with UTF-8 paths on Windows like the llama state payload, fall back to full
prompt re-processing when a checkpoint restored from a slot file fails to load,
and replace the 1024 count cap by keeping the last n_ctx_checkpoints while reading.

* server : report an incomplete checkpoint appendix as a failed slot save

Return an error to the client when the appendix cannot be written, like a
failed payload write, and make the oversized-blob test declare a size that
cannot be allocated, so an unbounded allocation fails the test.

* server : reject an empty target state in the checkpoint appendix

A saved checkpoint always holds a target state, an empty blob would roll back
without restoring anything. Also log with the slot id, and load the draft test
model from the HF cache instead of a second download.

* common : return bool from checkpoint load_tgt / load_dft

A checkpoint restored from a slot file falls back to full prompt re-processing
when it fails to load, a checkpoint created in memory still aborts.

---------

Co-authored-by: Igor Okulist <okigan@gmail.com>
This commit is contained in:
Leebr Data ConsultingandIgor Okulist authored and GitHub committed 2026-10-08 13:18:20 +03:00
1 parent ff5888f999
commit 033df86b69
6 files changed
+419 -24
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tools/server/server-context.cpp
+156 -7
View File
@@ -2583,6 +2583,136 @@ private:
cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
}
// checkpoints are appended to the slot save file, after the llama state payload
// they cannot be recreated from the final state alone (a recurrent state cannot be rewound)
static constexpr uint32_t SLOT_CKPT_MAGIC = 0x504b4353; // "SCKP"
static constexpr uint32_t SLOT_CKPT_VERSION = 1;
static bool ckpt_read(std::ifstream & ifs, void * dst, size_t size, size_t & n_read) {
if (!ifs.read((char *) dst, size)) {
return false;
}
n_read += size;
return true;
}
static bool ckpt_read_buf(std::ifstream & ifs, std::vector<uint8_t> & 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) {
return false;
}
buf.resize(n);
return n == 0 || ckpt_read(ifs, buf.data(), n, n_read);
}
static void ckpt_write(std::ofstream & ofs, const void * src, size_t size, size_t & n_written) {
ofs.write((const char *) src, size);
n_written += size;
}
static void ckpt_write_buf(std::ofstream & ofs, const std::vector<uint8_t> & buf, size_t & n_written) {
const uint64_t n = buf.size();
ckpt_write(ofs, &n, sizeof(n), n_written);
if (n > 0) {
ckpt_write(ofs, buf.data(), n, n_written);
}
}
// returns false if the appendix could not be written completely
bool save_slot_checkpoints(const std::string & filepath, const server_slot & slot, size_t & n_written) const {
n_written = 0;
if (slot.prompt.checkpoints.empty()) {
return true;
}
std::ofstream ofs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::app);
if (!ofs) {
SLT_WRN(slot, "failed to append context checkpoints to '%s'\n", filepath.c_str());
return false;
}
const uint32_t magic = SLOT_CKPT_MAGIC;
const uint32_t version = SLOT_CKPT_VERSION;
const uint32_t count = (uint32_t) slot.prompt.checkpoints.size();
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) {
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);
ckpt_write_buf(ofs, cur.data_tgt, n_written);
ckpt_write_buf(ofs, cur.data_dft, n_written);
ckpt_write_buf(ofs, cur.data_spec, n_written);
}
ofs.flush();
if (!ofs) {
SLT_WRN(slot, "failed to append context checkpoints to '%s' - the appendix is incomplete\n", filepath.c_str());
return false;
}
SLT_INF(slot, "appended %u context checkpoint(s) (%.3f MiB) to '%s'\n",
count, (float) n_written / 1024 / 1024, filepath.c_str());
return true;
}
// returns the number of bytes consumed, 0 if there is no usable appendix
size_t load_slot_checkpoints(const std::string & filepath, size_t offset, server_slot & slot) const {
std::ifstream ifs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::ate);
const size_t file_size = ifs ? (size_t) ifs.tellg() : 0;
if (!ifs || file_size < offset || !ifs.seekg(offset)) {
return 0;
}
const size_t n_avail = file_size - offset; // bytes after the llama state payload
size_t n_read = 0;
uint32_t magic = 0;
uint32_t version = 0;
uint32_t count = 0;
if (!ckpt_read(ifs, &magic, sizeof(magic), n_read) || magic != SLOT_CKPT_MAGIC) {
return 0;
}
if (!ckpt_read(ifs, &version, sizeof(version), n_read) || version != SLOT_CKPT_VERSION ||
!ckpt_read(ifs, &count, sizeof(count), n_read)) {
SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
return 0;
}
std::list<common_prompt_checkpoint> checkpoints;
for (uint32_t i = 0; i < count; ++i) {
common_prompt_checkpoint cur;
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) ||
!ckpt_read(ifs, &cur.pos_max, sizeof(cur.pos_max), n_read) ||
!ckpt_read_buf(ifs, cur.data_tgt, n_avail, n_read) ||
!ckpt_read_buf(ifs, cur.data_dft, n_avail, n_read) ||
!ckpt_read_buf(ifs, cur.data_spec, n_avail, n_read)) {
SLT_WRN(slot, "truncated context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
return 0;
}
// a saved checkpoint always holds a target state - an empty blob would roll back without restoring anything
if (cur.data_tgt.empty()) {
SLT_WRN(slot, "invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str());
return 0;
}
checkpoints.push_back(std::move(cur));
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);
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();
}
}
}
slot.prompt.checkpoints = std::move(checkpoints);
SLT_INF(slot, "restored %zu context checkpoint(s) from '%s'\n", slot.prompt.checkpoints.size(), filepath.c_str());
return n_read;
}
// returns false to decline the task, it is offered again after the decode is done
bool process_single_task(server_task && task, bool is_yielding) {
// while yielding, an encode / decode is running and only reading the server state is safe
@@ -2792,6 +2922,12 @@ private:
break;
}
size_t nwrite_ckpt = 0;
if (!save_slot_checkpoints(filepath, *slot, nwrite_ckpt)) {
send_error(task, "Unable to save slot: incomplete context checkpoints", ERROR_TYPE_SERVER);
break;
}
const int64_t t_end = ggml_time_us();
const double t_save_ms = (t_end - t_start) / 1000.0;
@@ -2801,7 +2937,7 @@ private:
res->filename = filename;
res->is_save = true;
res->n_tokens = slot->prompt.tokens.size();
res->n_bytes = nwrite;
res->n_bytes = nwrite + nwrite_ckpt;
res->t_ms = t_save_ms;
queue_results.send(std::move(res));
} break;
@@ -2857,6 +2993,9 @@ private:
break;
}
// nread is the end offset of the llama state payload within the file
const size_t nread_ckpt = load_slot_checkpoints(filepath, nread, *slot);
const int64_t t_end = ggml_time_us();
const double t_restore_ms = (t_end - t_start) / 1000.0;
@@ -2866,7 +3005,7 @@ private:
res->filename = filename;
res->is_save = false;
res->n_tokens = slot->prompt.tokens.size();
res->n_bytes = nread;
res->n_bytes = nread + nread_ckpt;
res->t_ms = t_restore_ms;
queue_results.send(std::move(res));
} break;
@@ -3278,7 +3417,7 @@ private:
if (ctx_dft) {
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
GGML_ASSERT(ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
}
if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) {
@@ -3604,8 +3743,18 @@ private:
if (!do_reset) {
// restore the context checkpoint
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->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) {
GGML_ABORT("failed to restore context checkpoint\n");
}
// restored from a slot file, not guaranteed to load - fall back to full prompt re-processing
SLT_WRN(slot, "%s", "failed to load context checkpoint restored from a slot file\n");
do_reset = true;
}
}
if (!do_reset) {
// restore the draft's speculative state
common_speculative_set_state(spec.get(), slot.id, it->data_spec);
@@ -4300,10 +4449,10 @@ private:
SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size());
ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
GGML_ASSERT(ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
if (slot.ctx_dft) {
ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
GGML_ASSERT(ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY));
}
slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);