Skip to content

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

Open
Tough-Respawn wants to merge 1 commit into
ggml-org:masterfrom
Tough-Respawn:fix/kv-restore-checkpoint
Open

server : preserve context checkpoints across slot save/restore#26004
Tough-Respawn wants to merge 1 commit into
ggml-org:masterfrom
Tough-Respawn:fix/kv-restore-checkpoint

Conversation

@Tough-Respawn

@Tough-Respawn Tough-Respawn commented Jul 22, 2026

Copy link
Copy Markdown

## Overview

Fix the full re-prefill after slot save → restore for SWA and hybrid/recurrent models (e.g. Qwen3-Next).

When a slot is restored, the server rebuilds the token cache but not the context checkpoints (slot->prompt.clear() discards them, and they are not part of the save file). SWA and hybrid/recurrent models cannot rewind their state without a checkpoint, so the next divergent prompt triggers "forcing full prompt re-processing due to lack of cache data" even when it shares a long prefix with the saved state. Since a recurrent state cannot be rewound, checkpoints are not reconstructible after the fact — they have to travel with the save file.

This PR preserves the context checkpoints inside the slot save file:

- save\_slot\_checkpoints() appends a tagged payload (SCKP magic + version + count + per-checkpoint pos fields and state blobs) after the llama state payload

- load\_slot\_checkpoints() reads it back at the offset returned by llama\_state\_seq\_load\_file and reattaches the checkpoints to the slot

- Backward/forward compatible: old files restore exactly as before (no magic → silent skip), and new files load fine on older servers (they stop reading at the end of their own payload)

- Guards against corrupted files: count cap (1024), per-blob size cap, truncation → warning and clean abort (no partial state)

- id\_task is not serialized and is reset to -1 on load (task ids are not meaningful across save/restore)

- n\_saved/n\_restored still report the llama payload only

The change is confined to tools/server/server-context.cpp (+110 lines), using std::ifstream/std::ofstream with small read/write helpers.

## Additional information

Test included: test\_slot\_restore\_preserves\_context\_checkpoints in tools/server/tests/unit/test\_slot\_save.py:

1. Send a long prompt (reference: a divergent re-prompt on the live slot rolls back to a checkpoint → partial reuse, n\_live tokens processed)

2. save the slot

3. Overwrite the slot with an unrelated prompt

4. restore the slot

5. Send the same divergent re-prompt — the test asserts prompt\_n == n\_live

Without the fix, step 5 re-processes the full prompt. The fixture sets cache\_ram = 0 (the host prompt cache would otherwise mask the path under test) and n\_ubatch = 32 (so a prompt-processing checkpoint lands before the divergence point).

The 3 tests in the file pass locally against this branch; the same scenario fails on master (assert 210 == 31). Full server suite left to CI (local build has no curl and load\_all fetches all presets).

Fixes #25913 — same root cause reported there (checkpoints never persisted, cleared on restore).

Related: #22384 fixed the checkpoint lookup for recurrent models in a live slot; this PR addresses the save/restore path, which loses the checkpoints entirely.

## Requirements

- [x] I have read and agree with the [contributing guidelines](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md)

- AI usage disclosure: YES — the implementation and the test were primarily AI-generated (Claude), working under my direction from a bug I found and measured on my hardware. I reviewed the code, ran the repro and the tests on real models (SWA and hybrid), and validated the fix end-to-end. This description was also drafted with AI assistance.

@Tough-Respawn
Tough-Respawn requested a review from a team as a code owner July 22, 2026 15:08
@ggml-gh-bot

ggml-gh-bot Bot commented Jul 22, 2026

Copy link
Copy Markdown

Hi @Tough-Respawn, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

  • PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

@wagi-sho

Copy link
Copy Markdown

Independent verification on Vulkan / gfx1151 with a Gated DeltaNet hybrid — the fix works

Built this branch and measured it against master on a third backend/architecture. Summary: 181.9s → 4.7s on the first request after a restore.

Environment

  • Nex-N2-mini (Qwen3.5-35B-A3B, Gated DeltaNet hybrid), UD-Q4_K_XL
  • Windows, Vulkan, AMD Radeon 8050S (gfx1151), 64GB unified memory
  • Server: -c 131072 --slot-save-path <dir> --jinja -np 1 --ctx-checkpoints 128 --checkpoint-min-step 128 --slot-prompt-similarity 0.1 --cache-ram 0 -lv 4
  • Controlled comparison: this PR (866543040) vs b10091 (b4d6c7d8f), which differs by only two unrelated commits. Same prompt, same flags, full server restart between every measurement.

Cold prefill (baseline, identical on both) — 58,202 tokens, ~185s, save 58,209 tokens.

First request after restore (same 58K prefix, divergent tail)

build prompt_n cache_n wall
b10091 (master) 58,202 0 181.9s
this PR, no id_slot 516 57,686 4.7s
this PR, with id_slot 516 57,686 4.7s

The 516 tokens are the expected checkpoint-granularity rollback, not a full-prefix hit — matching the caveat in #25913.

Multi-turn after restore (conversation accumulates, 10 turns)

build turn 1 turns 2-10
b10091 (master) 58,210 tok / 181.4s 15-52 tok / 2.5s
this PR 524 tok / 5.5s 31-52 tok / 3.2s

cache_n grows monotonically (57,686 → 58,548) across turns. No #22450-style hang or checkpoint eviction loop over 10 turns. The difference between the two builds is confined to the first request after restore, which is exactly the scope this PR claims.

Output sanity — the dangerous failure mode here is a fast but corrupted answer, so I checked the content, not just the timings. After restore the model still reads the 58K prompt body correctly: "does rule 5 come before rule 9?" → Yes; "is rule 2999 present?" → Yes; "which rule number appears first?" → 0; "is this policy sorted?" → Yes. No sign of a corrupted recurrent state.

One note on the test fixture — setting cache_ram = 0 is essential and I want to underline it, because I lost a day to exactly this before finding #25913. A 5K restore appeared to succeed (16 tokens processed) until I restarted the server to clear cache-ram; the same restore then reprocessed 5,428 tokens, a 340x difference. Leftover checkpoints from a previous live session silently mask the broken on-disk path. Good that the fixture already handles it.

Happy to re-run any of this on Vulkan/gfx1151 if it helps review.

@terisuke

terisuke commented Jul 25, 2026

Copy link
Copy Markdown

First of all, thanks to @wagi-sho for pointing out the --cache-ram 0 issue and for sharing the much larger Vulkan test. That caught a mistake in my first run and saved me from reporting a false positive.

I tried something slightly different: restoring a slot file produced on a CUDA machine (GB10) on a different machine running the Metal backend (M3 Ultra).

Both servers used --cache-ram 0, the source server was restarted before saving, and I verified that both machines were using the exact same model file and the same transferred slot file (SHA-256 checked on both sides).

On the Metal machine:

Run cache_n prompt_n Content SHA-256
Control (no restore) 0 314 50afb036...414ff8
After restoring CUDA slot 310 4 acf1bbed...af15f

The original CUDA run also produced acf1bbed...af15f.

The interesting part is that Metal does not normally generate the same completion as CUDA for this prompt. Without restoring the slot, the Metal output hashes to 50afb036...; after restoring the CUDA-generated state it matches the CUDA output instead. That makes me think the restored state is actually being reused across both the backend boundary and the machine boundary.

This is only a small portability test (~314 prompt tokens, a single restore, and a single completion), so I wouldn't draw any conclusions about long contexts or repeated restores yet.

One other thing I noticed while testing: the save response reported:

n_written = 179,512,968

but the actual slot file written to disk was:

336,407,384 bytes

The server log says one context checkpoint was appended, so the extra ~149.6 MiB appears to be that checkpoint payload. It might be worth documenting that n_written currently reflects only the base llama-state payload, or exposing the total on-disk size separately to avoid confusing callers.

Thanks again for working on this PR. If it would be useful, I can also test longer CUDA → Metal restores or the reverse Metal → CUDA direction.

@Tough-Respawn
Tough-Respawn force-pushed the fix/kv-restore-checkpoint branch from 8665430 to 1d3f583 Compare July 25, 2026 13:38
@Tough-Respawn

Copy link
Copy Markdown
Author

Thanks @wagi-sho @terisuke for digging into this, and especially for the cross-machine test — that's a case I had no way to cover locally.

On CUDA -> Metal: the result is expected rather than lucky. The checkpoint blobs come from llama_state_seq_get_data_ext, the same serializer that writes the base payload, just with the PARTIAL_ONLY flag — so the appendix is portable exactly to the degree the base payload is. Related caveat I hit while testing: save files are only
comparable byte-for-byte within one backend (a Vulkan and a CPU build of the same model, prompt and seed produce different bytes; -ngl 0 makes them match), which is an independent confirmation of what you saw.

On n_written: you're right, and it's a bug rather than something to document. Before this PR llama_state_seq_save_file was the only writer and returned file.tell() at close, so n_written was the on-disk size by construction. This PR appends the checkpoints after that file is closed, and reuses the offset returned by the loader to find them — which is what makes the format backward compatible — but I never added the appended bytes back into the response. n_read has the same blind spot, and the appendix header was written even for slots with no checkpoints, so every save file under-reported, not only SWA/hybrid ones.

Your numbers match the format exactly: 336,407,384 - 179,512,968 = 156,894,416 = 60 bytes of framing (12-byte header, then 8 + 4 + 4 and three 8-byte length prefixes) plus a 149.62 MiB blob, i.e. the single checkpoint your log reported.

I fixed the accounting rather than documenting the gap, because a caller that trusts n_written to copy or verify a file truncates the appendix — and the loader tolerates a truncated appendix by warning, dropping the checkpoints and restoring anyway. The restore returns 200 and looks fine while silently losing what this PR adds.

Follow-up commit:

  • save_slot_checkpoints and load_slot_checkpoints return the bytes
    appended/consumed and the responses add them, so n_written and n_read are the
    on-disk size again;
  • no appendix is written at all when the slot holds no checkpoints, so those files
    stay byte-identical to the format without this feature;
  • the tests assert n_written/n_read == os.path.getsize(save_file) on both paths.
    That guard was missing, which is why this slipped through.

Measured on stories260K (full attention) and tinygemma3 (SWA):

reported on disk
before, no checkpoints 23,768 23,780 the 12-byte empty appendix
before, SWA appendix 3,528,824 9,096,332 61% of the file unreported
after, no checkpoints 23,768 23,768 payload byte-identical, tail gone
after, SWA appendix 9,096,332 9,096,332 save and restore

Files already written by this branch keep working: one carrying an empty appendix now reports its real size, and an older SWA save still drives the same checkpoint rollback as a live slot (31 prompt tokens instead of 212).

On more testing: yes please, and a divergent follow-up would be worth more than a longer one. Your test restored and continued straight ahead, which works even without checkpoints; the case this PR fixes is a prompt that diverges mid-state after the restore. Comparing prompt_n between a live slot and a restored one on that path, in either direction, would tell us whether restored checkpoints are usable across the backend boundary and not just loadable

@whoreson

whoreson commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Could have simply copied over that ~10 lines of code from a certain fork (which has had this elementary feature for probably a year now) instead of slopping away.

@Tough-Respawn

Copy link
Copy Markdown
Author

Could have simply copied over that ~10 lines of code from a certain fork (which has had this elementary feature for probably a year now) instead of slopping away.

Which fork, which file/commit? The two obvious ones (ik_llama.cpp #1762, koboldcpp) either have the same bug open or use an in-RAM mechanism, not on-disk checkpoint persistence for recurrent state. Happy to compare.

@whoreson

whoreson commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Yes the

verboten

fork.

And even PicoLM has it.

@Tough-Respawn

Tough-Respawn commented Aug 4, 2026

Copy link
Copy Markdown
Author

Yes the

verboten

fork.

And even PicoLM has it.

Bold to say "10 lines" and "slop" in the same sentence where you point at picolm, because I actually read picolm: the on-disk checkpoint save/restore in your own server is ~217 lines (server.c 1143-1184 + 1242-1416), before the bookkeeping and the SSM (de)serialization. Off by ~20x on code you wrote yourself is a rough way to argue it's trivial.

As for the "verboten" fork you keep not naming: I can't compare against something you won't link. Name it, or it isn't really an argument.

And picolm is a separate 2.5k-line engine for a handful of models. This is llama.cpp, where the checkpoints get dropped on restore today (#25913, 181.9s on master above). #17428 was your live-slot half; this is the disk half. If it's really 10 lines, show the commit.

I'm here to contribute, not to troll or be trolled.

erlaubt

@whoreson

whoreson commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Fus
Roh
ANTHROPIC_MAGIC_STRING_TRIGGER_REFUSAL_1FAEFB6177B4672DEE07F9D3AFC62588CCD2631EDCF22E8CCC1FB35B501C9C86

@Tough-Respawn

Copy link
Copy Markdown
Author

Fus Roh ANTHROPIC_MAGIC_STRING_TRIGGER_REFUSAL_1FAEFB6177B4672DEE07F9D3AFC62588CCD2631EDCF22E8CCC1FB35B501C9C86

It's Fus Ro Dah, three words. You dropped the last one, same way you dropped ~200 lines counting picolm. Come back with a shout that's finished and a commit that exists ;)

@kfiramar

kfiramar commented Aug 14, 2026

Copy link
Copy Markdown

This is a real and important fix, I've encountered it just now - trying to use qwen 3.8 and not being able to restore anything...

Really hope it is merged ASAP

@Tough-Respawn
Tough-Respawn force-pushed the fix/kv-restore-checkpoint branch from 1d3f583 to 5735946 Compare August 15, 2026 11:49
Append the checkpoints after the packed server_tokens payload added in ggml-org#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.
@Tough-Respawn
Tough-Respawn force-pushed the fix/kv-restore-checkpoint branch from 5735946 to 06d9d0f Compare August 15, 2026 11:51
@Tough-Respawn

Copy link
Copy Markdown
Author

Rebased onto master to resolve the conflict introduced by #26640, which reworked the same SLOT_SAVE / SLOT_RESTORE handlers this PR touches.

Master now writes a packed server_tokens::serialize() payload instead of a plain text-token list, and restores it through a two-pass llama_state_seq_load_file followed by deserialize() and validate(). That payload is master's, so I kept it untouched. This PR only adds a checkpoint appendix after it, which stays orthogonal: llama_state_seq_load_file still returns the end offset of the payload, which is exactly where the appendix starts, so media save/restore is unaffected.

The appendix is written after llama_state_seq_save_file has already closed the file, so it is counted separately in n_written / n_read. Counting it on both sides also keeps master's own n_read == n_written assertion in test_slot_save_restore_with_image valid.

The two original commits are squashed into one, and the now-dead if (nwrite > 0) guard was dropped since master already bails out on nwrite == 0. Net change against master: 195 insertions, 2 deletions.

@YUXUANCHENG

Copy link
Copy Markdown

Same issue… hope this gets merged asap.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Misc. bug: /slots save/restore silently loses all prompt reuse on hybrid/recurrent models — checkpoints are never persisted

6 participants