[cwriter]

Overview

llama-server currently creates and reuses context-checkpoints through a vector copy.
Since the checkpoints are immutable by design, this patch introduces reference-counted sharing instead.

This PR includes 3 changes to the system:

• Use std::shared_ptr for the context checkpoints, allowing cheap (ref-counted) copies
• Skip zero-initialization - this would be a security risk, but we ensure that the bytes are actually written
• Reject broken requests before inserting into the prompt cache, avoiding leakages and thrashing. This is cheap because it's just an std::move

Additional information

This applies mainly to the host-side copies. The effects become larger as context grows. The effects are most visible when queuing events, as they need to wait on the serialized cache-checkpoint copies.
Tested on 3x Arc Pro B60, Threadripper 1950X x399 platform, 128GB RAM, qwen3.8-flash-next:UD-IQ4XS, RAM prompt cache 8GB, 16 checkpoints, checkpoint spacing 8192 tokens, no MTP, running on the SYCL backend (but should not affect this code).

Workload	Metric	Master	This PR	Reduction
Fresh history, 2-3 retained checkpoints	Cache update	529.213 ms	250.798 ms	52.6%
Fresh history, 2-3 retained checkpoints	First visible token	831.134 ms	542.500 ms	34.7%
Fresh history, 2-3 retained checkpoints	Full 16-token request	1583.031 ms	1295.366 ms	18.2%
Grown history, 8-9 retained checkpoints	Cache update	1039.985 ms	244.963 ms	76.4%
Grown history, 8-9 retained checkpoints	First visible token	1345.784 ms	532.247 ms	60.5%
Grown history, 8-9 retained checkpoints	Full 16-token request	2100.605 ms	1287.922 ms	38.7%

Memory savings:

Workload	Metric (MiB)	Base	PR	Saved
One slot, 2 checkpoints	Steady RSS	1420.5	1421.4	-0.9 (-0.1%)
One slot, 2 checkpoints	Peak RSS	1845.5	1621.1	224.4 (12.2%)
One slot, 8-9 checkpoints	Steady RSS	2828.0	2829.2	-1.2 (-0.0%)
One slot, 8-9 checkpoints	Peak RSS	4209.6	3197.8	1011.8 (24.0%)
Two slots, idle caching	Steady RSS	4757.2	3181.3	1576.0 (33.1%)
Two slots, idle caching	Peak RSS	4757.2	3181.3	1576.0 (33.1%)

Requirements

• I have read and agree with the contributing guidelines - Yes
• AI usage disclosure: Yes, Claude Code Opus 5.5 for issue identification and code draft, Opus Sol 6.1 for PR extraction

[Davidyz]

Would the shared references to the checkpoints also reduce the memory footprint of the prompt cache? If so, it might be useful to include numbers in the table.

[cwriter]

Would the shared references to the checkpoints also reduce the memory footprint of the prompt cache?

No, this patch uses shared_ptr only to have an easy copy mechanism. It does not reduce the memory usage overall, except maybe transiently during copies.

[Davidyz]

No, this patch uses shared_ptr only to have an easy copy mechanism. It does not reduce the memory usage overall, except maybe transiently during copies.

I see. But just for clarification, the scenario I had in mind was when 2 prompts share the same prefix (eg. system prompt+tool schema) but different user messages. With this pr, do they share the same checkpoint object for the matching system prompt? I'm curious because if the original behaviour was a vector copy, it sounds like it's going to store the same checkpoint twice.

[cwriter]

Yes, sorry, you're right: If there are the same prefixes (a growing conversation), this saves (peak) RAM. It does not save RAM unrelated histories, even if they have the same input (there is no deduplication).

Added the Memory table

[Davidyz]

Sweet! I have only a few GB of RAM left once I load the model/prompt cache, so even a few hundred MB of savings would be meaningful to me. Looking forward to seeing this in mainline.

[pwilkin]

/bot review

[ggml-gh-bot]

Automated code review

Review of PR #30137 (llama-server: reference-counted checkpoints). I checked the quick-reject gate, security, and the server/general checklists (paths touched: common/common.h, tools/server/*). No public API or ggml changes; the change is single-purpose and in scope for the server. Perf data is provided, which is good.

Summary of what I verified before the findings: every path that fills a common_state_data buffer also validates the fill - common_prompt_checkpoint::update_tgt/update_dft (common/common.cpp:2355-2389) abort on a size mismatch, ckpt_read_buf (tools/server/server-context.cpp:2606-2613) checks ifs.read before returning success, and prompt_save (tools/server/server-context.cpp:326-332) now checks the return of llama_state_seq_get_data_ext and pops the entry on failure. So the no-init allocator does not currently leak uninitialized bytes. I also verified that checkpoints are only mutated before publication into slot.prompt.checkpoints (create_checkpoint, server-context.cpp:2569-2585) and in load_slot_checkpoints before assign (server-context.cpp:2716-2722), and the shared_ptr<const common_prompt_checkpoint> type enforces immutability afterwards - the sharing itself looks sound.

Will slow the review

(point 1) Cache accounting now over-counts shared checkpoint memory, and eviction cannot free it. server_prompt_cache_state::size() (tools/server/server-task.h:651-658) and the admission/eviction logic in server_prompt_cache::alloc (tools/server/server-task.cpp:1734-1775) still count every checkpoint once per holder. Three consequences: (a) a cache entry whose checkpoints are already resident in the slot (marginal real cost near zero) can be rejected with "exceeds cache size limit", silently disabling prompt caching for that slot; (b) update()/alloc may evict real data to make accounting room for bytes that are already in RAM and would not be freed by the eviction; (c) --cache-ram now bounds an inflated number, so the cache holds less useful content for the same real memory. Since the PR's stated goal is memory savings for constrained users, please address this - e.g. count each unique checkpoint once (dedupe by pointer) or exclude slot-resident checkpoints from the per-state size.

(point 2) PR description/diff mismatch: the description's third bullet ("Reject broken requests before inserting into the prompt cache ... cheap because it's just an std::move") does not correspond to anything in this diff. The actual change validates the state after insertion and pops it on failure (server-context.cpp:326-332). Please correct the description so reviewers evaluate what the code actually does.

(point 3) prompt_save reaches into prompt_cache.states.pop_back() (tools/server/server-context.cpp:330) - server_slot now needs to know that alloc returned &states.back(). Consider having the cache handle the rollback itself (e.g. alloc returns a small guard, or a dealloc method) so the invariant "the returned pointer is the back of the list" stays inside server_prompt_cache. Minor, but it removes a hidden coupling.

(point 4) Test coverage: the new sharing between slot, cache entries, and copy_prompt_to/copy_state_to is not covered by tests. test_kv_keep_only_active.py exercises cache-ram without checkpoints, and test_slot_save.py exercises the file appendix. Please at least run the cache-ram test with --ctx-checkpoints > 0 (save -> idle eviction -> restore-from-cache -> slot save) to cover the shared-checkpoint paths end to end.

Nits

(point 5) common/common.h:1265 names std::allocator explicitly but the header does not include <memory> (it only works transitively today, same as the pre-existing std::unique_ptr usage). Adding the include makes the new code self-contained.

(point 6) In create_checkpoint (tools/server/server-context.cpp:2585-2590), the trace log after push_back(std::move(ckpt)) still reads through the cur reference into the moved-from shared_ptr. It is currently safe (the list keeps the object alive), but capturing the logged values before the move would be more robust against future edits. Similarly, the two consecutive const auto & cur = **it; blocks in the restore path (server-context.cpp:3757-3761 and 3764-3765) could be one binding.

Security note (no finding, but worth keeping in mind)

The no-init vector is correct today only because all three fill sites validate completeness; the type itself enforces nothing. Any future common_state_data that is resized and only partially filled before being passed to llama_state_seq_set_data_ext or written to a slot file would push uninitialized heap bytes into model state (potentially observable in generated output) or to disk. The header comment on the allocator covers this, which is good - just make sure future users of common_state_data keep that contract.

Overall: the approach (immutable, reference-counted checkpoints) is the right one and the fill-path validation is done properly. Please address the accounting issue (point 1) and the description mismatch (point 2) before requesting maintainer review.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[sarahwu185]

Hi @Davidyz, I think #30683 may help address the reuse issue.

[Nthalk]

Here for my support of this, I run this pr (augmented with a few patches of my own). It's stable, and saves anywhere between 6% and 30% of time when it hits. I have several downstream patches for better caching that gets my workload up to 94% cached for limited pinned streams - allowing a fused (cache load + processed) TPS of 26k tokens processed a second on a 5090 with qwen 3.8 27b q6ms @ 70k ctx.

[cwriter]

Is there something that is still required to get this merged?

[Nthalk]

Yeah, I've been running this since day one. It's great for throughput.

[ggerganov]

Will take a look - please rebase again on latest master when you get the chance.

[cwriter]

Thanks, rebased

[ServeurpersoCom]

Tested on a rootless container with an RTX PRO 6000, Qwen3.6-35B-A3B UD-Q8_K_XL on one slot, two conversations alternating for 10 turns each so every request saves and reloads the prompt cache, 15 checkpoints of 63 MiB per conversation.

Without this PR the time spent outside prompt processing grows to 182 ms per request at 25k tokens and the RAM peak includes a full copy of the checkpoints.

With this PR it drops to 104 ms, the RAM peak is 940 MiB lower and all 20 outputs are identical to master.

LGTM!

@cwriter few small nits, nothing blocking:

• the third bullet of the description does not match the diff anymore, the state is validated after insertion and popped on failure
• src/llama-impl.h already has no_init for the same purpose, you could mention why common needs its own allocator

[bhargav191098]

@ServeurpersoCom @cwriter Following up on #31391, there was still some CPU stall when saving Qwen3.8-27B prompt state.

The CUDA gather improved the transfer, but each save in this workload still deep-copied two retained recurrent-state checkpoints, totalling 299.25 MiB. These were earlier checkpoints already on the CPU, copied into the RAM prompt-cache entry—not the new GPU-to-CPU snapshot transfer.

In instrumented baseline runs, cache insertion/checkpoint copying took 42.66 ms, and allocating/zeroing the new snapshot buffer took another 28.98 ms. This PR addresses those costs by sharing immutable checkpoints and skipping buffer zero-initialization before serialization fills them.

I applied the diff at c73ed8dbd to master 149f5641b, which already includes #31391, and ran three matched fresh-server pairs:

Metric	Master	Master + this PR	Change
Mean snapshot-save interval	95.23 ms	44.30 ms	−53.5%
Maximum snapshot-save interval	128.15 ms	54.88 ms	−57.2%
Output throughput	176.74 tok/s	182.86 tok/s	+3.5%
Mean TTFT	483.09 ms	407.54 ms	−15.6%

RTX 5090, CUDA 13.4, Release build, Qwen3.8-27B Q4_K_M; np4/KVU/concurrency 4, context 20,480, FA on, Q8_0 K/V. AIPerf: 64 streaming requests, ISL512/OSL512, output standard deviation 128. No profiler or speculative decoding.

These are medians of per-run statistics. Every run completed 64 requests and recorded 60 saves, with matching request payloads, output-token counts and snapshot lengths. Save intervals include CPU work and backend completion; the saves exercise the shared LRU/idle-slot prompt_save() path.

Here to provide some numbers in support of this PR! :D

[Nthalk]

Adding data on top of this PR (running it since day one, now at c73ed8d on master 5851a8e incl. #31391; RTX 5090, Qwen3.8-27B Q6_K_S, 128k ctx, q8_0 KV, --cache-ram 40960, --ctx-checkpoints 8).

A further duplicate copy remains after this PR: when a task restores a checkpoint at the last-user-message boundary, create_checkpoint then supersedes it with a byte-identical fresh save (~330 MiB, ~0.3 s here). We saw this on 45% of restores (591/1,321). A small local patch that keeps the checkpoint the current task just restored (other same-position checkpoints are still superseded) cut <500-new-token prefill 0.67 → 0.49 s; stacked with this PR, per-request prefill went 1.15 → 1.01 s (500–2k), 2.29 → 1.94 s (2k–8k). Happy to open it as a follow-up once this lands.

Also confirming review point 1: size() counts shared checkpoints once per holder, so --cache-ram evicts early.

[ggerganov]

@ServeurpersoCom Please merge if all is good on your end.

[ServeurpersoCom]

Thanks @cwriter, and thanks @bhargav191098 and @Nthalk for the extra numbers.

Two follow-ups, none of them a regression vs master: the prompt cache still counts a shared checkpoint once per holder, so --cache-ram evicts earlier than it needs to now that holders share memory (point 1 of the bot review, confirmed by @Nthalk); and the remaining nits from my review (description bullet, reuse of no_init from llama-impl.h).

@Nthalk the restored checkpoint that gets superseded by an identical save is what #30683 addresses, it will be rebased on top of this one.