fix: checkpoint restore for hybrid/recurrent models#1976
Open
razlani wants to merge 1 commit into
Open
Conversation
- skip prompt cache load when slot has good match
firecoperana
requested changes
Jun 16, 2026
Comment on lines
+3581
to
+3585
| // Recurrent/hybrid models need checkpoints much sooner to avoid gaps that force | ||
| // full reprocessing - 4 tokens vs 64 for transformer-only. | ||
| const bool has_recurrent = llama_model_has_recurrent(llama_get_model(slot.ctx)); | ||
| const int checkpoint_min_tokens = has_recurrent ? 4 : 64; | ||
| do_checkpoint = do_checkpoint && pos_min >= 0 && slot.cache_tokens.n_tokens() >= checkpoint_min_tokens; |
Collaborator
There was a problem hiding this comment.
Why do you need to save a checkpoint when you only need to reprocess 4 tokens compared to 64?
Comment on lines
+3517
to
+3519
| if (has_recurrent) { | ||
| return cur.pos_max <= pos_next; | ||
| } |
Collaborator
There was a problem hiding this comment.
There was a recent PR that we have to revert because this breaks the model when you regenerate the same conversation, so it does not work.
Comment on lines
+3836
to
+3840
| if (has_recurrent_ckpt) { | ||
| // KV is cleared but keep token tracking for checkpoint restore. | ||
| // apply_checkpoint() will find a checkpoint at or before n_past | ||
| // and restore the recurrent state from it. | ||
| } else { |
Collaborator
There was a problem hiding this comment.
Don't think this is correct. We always need to set n_past to 0. No need to change this part.
Comment on lines
+3491
to
+3500
| int32_t pos_min = llama_kv_cache_seq_pos_min(slot.ctx, slot.id); | ||
| const bool kv_empty = pos_min < 0; | ||
|
|
||
| // For recurrent/hybrid models: also enter checkpoint search when KV is empty | ||
| // (slot went idle and KV was cleared) but cache_tokens still reflects a prior | ||
| // decode and the new prompt fully covers that cached prefix. | ||
| const bool full_match_kv_lost = has_recurrent && kv_empty && | ||
| slot.n_past > 0 && | ||
| slot.n_past == slot.cache_tokens.n_tokens() && | ||
| !slot.server_cached_prompt.checkpoints.empty(); |
Collaborator
There was a problem hiding this comment.
When you clear the kv cache, checkpoints should be cleared too, so there is nothing to search.
SamuelOliveirads
left a comment
SamuelOliveirads
left a comment
Collaborator
There was a problem hiding this comment.
As far from fire's point I didn't find any other problems.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1762. Overlaps with #1888 (same root issue) - my fix addresses an additional root cause and includes a prompt cache load guard.
Three bugs compound to make checkpoint restore non-functional on hybrid/recurrent models (Qwen3.6, Mamba). Every turn after idle forces a full re-prefill from token 0.
apply_checkpoint()runsapply_checkpointgate condition excludes the recovery case (n_past zeroed by Offload Bitnet token embeddings to the GPU #1)pos_minwhich has wrong semantics for hybrid modelscache_tokens/n_past/checkpoints when partial KV delete fails for recurrent models, instead of nuking everythingfull_match_kv_lostpath intoapply_checkpoint()for when KV is empty but checkpoints existpos_max-based checkpoint matching for recurrent modelsf_keep >= cache_ram_similarity)Tested via (Qwen3.6-35B-A3B, RTX 5070 Ti,
-np 1 -c 131072):preserve_thinking: true: 0 full drops, >130 perfect cache hits in a sessionSingle-slot only. Multi-slot recurrent has separate issues (#1932/#1933).