From 468457d8116e431bc3643b4d2591cd2361a475dc Mon Sep 17 00:00:00 2001 From: krim Date: Sun, 12 Jul 2026 17:01:35 +0200 Subject: [PATCH 1/2] server : fix checkpoint handling for hybrid/recurrent models (#24055) The recurrent state is only valid at the exact position it was saved at, but checkpoints claimed the whole [pos_min, pos_max] range reported by the memory module, and restoring relied on that by accident. Record the actual position instead and only restore a checkpoint when that position is still inside the common prefix of the new prompt. Also erase checkpoints that overlap edited history - once the new tokens are decoded they would look valid again while holding state from content that no longer exists. While at it, make the cache actually useful for agentic clients that strip reasoning from previous turns: always checkpoint near the end of the prompt (that is where the next request diverges), evict the checkpoint closest to its neighbor instead of the oldest one so early anchors survive compaction, try the bounded n_rs_seq rollback before searching checkpoints, and log checkpoint activity at INFO so you can see what the cache is doing without -lv. --- src/llama-memory-hybrid.cpp | 4 ++- src/llama-memory-recurrent.cpp | 4 +++ tools/server/server-context.cpp | 55 +++++++++++++++++++++++++++------ 3 files changed, 53 insertions(+), 10 deletions(-) diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..27d9c5cf76c 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -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)); } diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 3d6c6db876b..b1d9c208bb3 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -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 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7564ad4e9cf..58cbb59b57b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2286,6 +2286,13 @@ struct server_context_impl { 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) { + return; + } + // evict checkpoints within min-step of a previous checkpoint, unless they were // created by the current task int64_t last = -1; @@ -2317,8 +2324,14 @@ 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_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS) { + 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); @@ -2326,7 +2339,7 @@ struct server_context_impl { // 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); @@ -3223,6 +3236,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()); @@ -3279,7 +3296,19 @@ struct server_context_impl { SLT_WRN(slot, "%s\n", st1.str().c_str()); } - if (pos_min >= pos_min_thold) { + // with enough per-token snapshots, the recurrent state can be rolled + // back to pos_next directly and no checkpoint is needed + const bool can_rollback = + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && + pos_min < pos_min_thold + (llama_pos) llama_n_rs_seq(ctx_tgt); + + if (pos_min >= pos_min_thold && !can_rollback) { + // 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_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + // search for a context checkpoint const auto it = std::find_if( slot.prompt.checkpoints.rbegin(), @@ -3287,6 +3316,11 @@ struct server_context_impl { [&](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 + return cur.pos_max < pos_min_thold; + } // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] if (cur.pos_max > pos_next) { return false; @@ -3306,11 +3340,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; @@ -3319,11 +3353,14 @@ 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, (llama_pos) slot.task->n_tokens()); + 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) { + 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; From e75a77ab6ce8b7c0e82268cac28a1f1a98d175a9 Mon Sep 17 00:00:00 2001 From: krim Date: Tue, 14 Jul 2026 18:55:30 +0200 Subject: [PATCH 2/2] server : drop the unsafe rs rollback fast-path and adopt restored checkpoints The n_rs_seq rollback snapshots are only guaranteed valid for tokens decoded in the last ubatch, so rolling back across decode boundaries could silently restore a stale state - always go through the checkpoint path instead. Adopt a restored checkpoint into the current task so the min-step eviction does not erase it, compare the stale bound in position space (mtmd positions differ from token counts), and keep range (SWA) checkpoints at the exact divergence position since resuming overwrites that entry anyway. --- tools/server/server-context.cpp | 39 +++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 58cbb59b57b..ed7e6c1c93e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2282,6 +2282,13 @@ 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; @@ -2290,6 +2297,8 @@ struct server_context_impl { 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; } @@ -2327,8 +2336,7 @@ struct server_context_impl { // 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_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS) { + if (ctx_tgt_state_exact()) { pos_min = pos_max; } @@ -3296,18 +3304,10 @@ struct server_context_impl { SLT_WRN(slot, "%s\n", st1.str().c_str()); } - // with enough per-token snapshots, the recurrent state can be rolled - // back to pos_next directly and no checkpoint is needed - const bool can_rollback = - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && - pos_min < pos_min_thold + (llama_pos) llama_n_rs_seq(ctx_tgt); - - if (pos_min >= pos_min_thold && !can_rollback) { + 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_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + const bool ckpt_exact = ctx_tgt_state_exact(); // search for a context checkpoint const auto it = std::find_if( @@ -3318,8 +3318,10 @@ struct server_context_impl { 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 - return cur.pos_max < pos_min_thold; + // 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) { @@ -3355,11 +3357,16 @@ struct server_context_impl { { // 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, (llama_pos) slot.task->n_tokens()); + 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 || cur.pos_max >= pos_stale) { + 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 {