Skip to content

server: fix checkpoint handling for hybrid/recurrent models (#24055) #25592

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

Open
krim404 wants to merge 2 commits into
base: master
Choose a base branch
from
+59 −9
Open

server: fix checkpoint handling for hybrid/recurrent models (#24055) #25592

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
468457d
server : fix checkpoint handling for hybrid/recurrent models (#24055)
Jul 12, 2026
e75a77a
server : drop the unsafe rs rollback fast-path and adopt restored che…
Jul 14, 2026
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
4 changes: 3 additions & 1 deletion src/llama-memory-hybrid.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,9 @@ void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p
}

llama_pos llama_memory_hybrid::seq_pos_min(llama_seq_id seq_id) const {
// the min of the total cache is the max of the two caches' min values
// the min of the total cache is the max of the two caches' min values.
// the recurrent state is valid only at its latest position, so the combined min must
// not report positions that the recurrent state cannot serve
return std::max(mem_attn->seq_pos_min(seq_id), mem_recr->seq_pos_min(seq_id));
}

Expand Down
4 changes: 4 additions & 0 deletions src/llama-memory-recurrent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,10 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
cell.pos = p0 - 1;
return true;
}
// cannot roll back beyond the available snapshots - the caller has to
// restore a checkpoint or reprocess the sequence
LLAMA_LOG_DEBUG("%s: cannot roll back recurrent state of seq %d by %d tokens (n_rs_seq = %u)\n",
__func__, (int) seq_id, (int) rollback, n_rs_seq);
return false;
}
// invalidate tails which will be cleared
Expand Down
60 changes: 52 additions & 8 deletions tools/server/server-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2282,10 +2282,26 @@ struct server_context_impl {
return true;
}

// whether the memory state is valid only at its exact final position (hybrid/recurrent),
// as opposed to a range of positions (SWA)
bool ctx_tgt_state_exact() const {
return ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL ||
ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS;
}

// n_tokens_cur: the number of tokens added to the batch for the current slot
void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) {
const int id_task = slot.task->id;

// an equivalent checkpoint already exists (e.g. it was just restored)
if (!slot.prompt.checkpoints.empty() &&
slot.prompt.checkpoints.back().n_tokens == slot.prompt.n_tokens() - n_tokens_cur &&
slot.prompt.checkpoints.back().pos_max == pos_max) {
// adopt the checkpoint so the min-step eviction below does not erase it
slot.prompt.checkpoints.back().id_task = id_task;
return;
}

// evict checkpoints within min-step of a previous checkpoint, unless they were
// created by the current task
int64_t last = -1;
Expand Down Expand Up @@ -2317,16 +2333,21 @@ struct server_context_impl {
cur.id_task = id_task;

// [TAG_CHECKPOINTS_FIX_POS_MIN]
// TODO: here we incorrectly deterimne that the saved checkpoint data covers the [pos_min, pos_max] range
// this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
// the state of hybrid/recurrent memory is valid only at its exact final position
// TODO: for SWA models the saved range can still claim more than it actually covers:
// https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
if (ctx_tgt_state_exact()) {
pos_min = pos_max;
}

cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max);

cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// stash the draft's speculative state with the checkpoint
common_speculative_get_state(spec.get(), slot.id, cur.data_spec);

SLT_TRC(slot,
SLT_INF(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,
cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
Expand Down Expand Up @@ -3223,6 +3244,10 @@ struct server_context_impl {

llama_pos pos_next = slot.prompt.tokens.pos_next(n_past);

// pos_next can be reduced below by a checkpoint restore - remember the
// divergence point for the checkpoint invalidation
const llama_pos pos_next_lcp = pos_next;

// ref: https://github.com/ggml-org/llama.cpp/pull/24110
const bool has_new_tokens = (n_past < slot.task->n_tokens());

Expand Down Expand Up @@ -3280,13 +3305,24 @@ struct server_context_impl {
}

if (pos_min >= pos_min_thold) {
// whether the checkpoints hold a state that is valid only at its exact
// final position (hybrid/recurrent memory), as opposed to a range (SWA)
const bool ckpt_exact = ctx_tgt_state_exact();

// search for a context checkpoint
const auto it = std::find_if(
slot.prompt.checkpoints.rbegin(),
slot.prompt.checkpoints.rend(),
[&](const auto & cur) {
// 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);
if (ckpt_exact) {
// usable only if the tokens up to and including its position are
// a prefix of the new prompt, with at least one token left to
// process [TAG_PROMPT_LOGITS]. the state is self-contained, so
// the SWA slack in pos_min_thold does not apply
return cur.pos_max < pos_next - (has_new_tokens ? 0 : 1);
}
// workaround for [TAG_CHECKPOINTS_FIX_POS_MIN]
if (cur.pos_max > pos_next) {
return false;
Expand All @@ -3306,11 +3342,11 @@ struct server_context_impl {

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);
SLT_INF(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);
}

if (do_reset) {
SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n",
SLT_INF(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n",
"https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055");
pos_next = 0;
n_past = 0;
Expand All @@ -3319,11 +3355,19 @@ struct server_context_impl {
}

{
// erase any checkpoints with pos_max > pos_next
// erase any checkpoints that cover diverged content - once the new
// tokens are decoded, their staleness would become undetectable
const llama_pos pos_stale = std::min(pos_next_lcp, slot.task->tokens.pos_next());

// an exact checkpoint at the divergence position irreversibly contains
// the diverged token, while a range (SWA) checkpoint gets that entry
// overwritten when decoding resumes from it
const llama_pos pos_stale_min = ctx_tgt_state_exact() ? pos_stale : pos_stale + 1;

for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) {
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);
if (cur.pos_max > pos_next || cur.pos_max >= pos_stale_min) {
SLT_INF(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);
} else {
++it;
Expand Down