[myml]

Overview

Context checkpoints are written with LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY, which asks for only the parts of the memory that cannot be rolled back.

llama_memory_hybrid::state_write() implements that by skipping mem_attn and writing only mem_recr, and state_read() does the same on the way back, so the target's attention cache is deliberately left out of every checkpoint and stays resident:

// src/llama-memory-hybrid.cpp
if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
    mem_attn->state_write(io, seq_id, flags);
}
mem_recr->state_write(io, seq_id, flags);

The draft model's context is a plain llama_kv_cache — a prefix-addressed attention cache, the same kind of memory as mem_attn. But llama_kv_cache::state_write() discards the flags and serializes every cell, so the PARTIAL_ONLY request has no effect on it:

// src/llama-kv-cache.cpp:2180
void llama_kv_cache::state_write(llama_io_wstr_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
    GGML_UNUSED(flags);
    for (uint32_t i = 0; i < cells.size(); ++i) {

So the whole draft KV ends up in every checkpoint. Since the draft is sized to the target's n_ctx, checkpoint size grows linearly with the context instead of being the constant cost of the recurrent state.

Skip the draft when creating a long-term checkpoint.

Measured

Qwen3.8-27B (16 attention + 48 recurrent layers), MTP draft, q8_0 KV, -c 262144, -ctxcp 144, checkpoints 1..8 over a 13789-token prompt:

checkpoint (n_tokens)	before	after
1 (2048)	157.673 MiB	149.626 MiB
2 (4097)	165.724 MiB	149.626 MiB
3 (6146)	173.775 MiB	149.626 MiB
4 (8195)	181.825 MiB	149.626 MiB
5 (10244)	189.876 MiB	149.626 MiB
6 (12293)	197.927 MiB	149.626 MiB
7 (13273)	201.778 MiB	149.626 MiB
8 (13771)	203.734 MiB	149.626 MiB

The slope is 46.061 MiB over 11723 tokens, i.e. 4.02 KiB per token, and extrapolating it to zero tokens lands exactly on the 149.626 MiB that the recurrent state alone costs (which is RS buffer / (1 + n_rs_seq), as expected).

The draft term therefore overtakes the recurrent state around 38k tokens, and at 256k it is 1.0 GiB per checkpoint — over the 144 checkpoints -ctxcp allows that is ~150 GiB instead of ~21 GiB.

Greedy output and draft acceptance are unchanged.

Why no matching change on the way back

The draft's cells are position-addressed and survive like mem_attn's, and the restore path already tolerates a checkpoint without draft data:

• common_prompt_checkpoint::load_dft() returns early when data_dft.empty() (common/common.cpp), so the draft memory is left untouched rather than cleared.
• The clear() in llama_kv_cache::state_read_sinfo() is guarded by seq_id == -1, and an empty per-stream cell count is skipped outright, so loading a partial state does not wipe the cache.

The per-draft-step rollback checkpoint in slot.spec_ckpt is untouched: it goes through update_dft() directly with ON_DEVICE and does not depend on the long-term checkpoints.

Requirements

• I have read and agree with the contributing guidelines: YES
• AI usage disclosure: YES — an AI coding agent assisted with reading the checkpoint/memory code, drafting the change and collecting the measurements above. The analysis, the decision to submit, and every claim in this description were reviewed and verified against the source by me (I run this in production on the machine the numbers come from), and I take responsibility for the patch.

[ggml-gh-bot]

Hi @myml, thanks for your contribution!

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

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

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

[DanoPTT]

Tested on gfx1201 (R9700), Qwen3.8-27B Q5_K_M + draft-mtp, f16 KV cache (K and V), ~100k context:

	checkpoint size
master	542.9 MiB
this PR	149.6 MiB

Output after restoring a checkpoint is identical to master (greedy, 4 requests), draft acceptance unchanged. Running it in production since today.

[myml]

merged #30514