[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.