server: reuse checkpoint state buffers from a bounded pool - #201
server: reuse checkpoint state buffers from a bounded pool#201danielhanchen wants to merge 6 commits into
Conversation
A context checkpoint of a hybrid or recurrent model holds the whole non-rollbackable sequence state, hundreds of MiB, and llama-server allocates and frees one per prompt. An allocation that size always comes from mmap() and always goes back on free, so the first write to a fresh buffer faults in every page. In update_tgt() that is the data_tgt.resize() zero fill, and on a DGX Spark serving a 27B hybrid at 32 slots it is 43.7 ms of the 50.2 ms a checkpoint costs, 288 of the 392 ms of a prefill iteration. Hand the buffer to a bounded pool instead of to the allocator and the pages stay mapped and resident, so the next checkpoint reuses them and pays neither the faults nor the fill. Note the fill is load bearing and must not simply be dropped: with a CUDA target context, llama_state_seq_get_data_ext() copying into pageable host memory that is not yet resident measures ~140x slower than into memory that is (6.5 ms against 917 ms for 149 MiB). On the CPU backend the same change is neutral, so this has to be measured on a GPU. The pool removes the fill by making it unnecessary, not by skipping it. The memory policy is in the comment on common_state_buffer_pool: a byte cap derived from host memory, a count cap, a size floor below which pooling saves nothing, and a trim on idle so the server does not hold the memory when it is not serving. Every cap degrades to today's behaviour rather than to something worse.
Review follow-ups on the pool, all of them about the policy rather than the reuse:
- The byte cap read min(total/16, free/4) of host memory. Every non-Windows host reports
free == total for the CPU device ("free system memory is ill-defined, assume all of it
is free"), so the free term never bound and the stated guarantee did not exist. It is
now a plain fraction of total, which is what it always was, said honestly.
- A failed host memory query arrives as a huge total, not as zero: sysconf(_SC_PHYS_PAGES)
returning -1 is multiplied out with no error check. That made the byte cap useless
exactly where the pool most needed to keep nothing. Bounded.
- The destructor is noexcept and put() can throw. server_prompt_cache::alloc() recovers
from bad_alloc by destroying cached prompts, so a throw there would have turned a cache
shrink into a terminate(). Wrapped.
- trim() was called from update_slots()'s all-idle branch, which is not reached while the
queue is empty: update_slots() only runs after a task. Moved to the task queue's idle
wait, which is the loop that actually runs when nothing is happening, and called with a
zero timeout from the prompt cache's out-of-memory recovery so pooled bytes are always
reclaimable under allocation pressure.
- put() declined when full instead of evicting. For a model whose checkpoints grow through
a prompt that wedges the pool full of buffers no later request can use, holding the
memory at a zero hit rate. It now displaces the smallest pooled buffer, and only one
smaller than itself, so a pool of equal or larger buffers still declines.
- get() dropped the caller's undersized buffer on a hit; it is offered back instead.
- A buffer the caller already owns and that is already large enough is a reuse, and the
cheapest kind, but it was counted as a miss.
No change to what is reused or when. Same md5 from the strict single-slot greedy harness,
and test-recurrent-state-rollback produces output identical to master.
bddcff0 to
468cf0b
Compare
|
Third-party evidence for this PR: the pool was composed onto the pipeline-groups branch and measured on the two-Spark layer split at 128 concurrent rows, which is a much harsher operating point than the single-node 32-slot cells this PR was measured at. It composes, and the honest claim is narrower than the single-node one. Setup, so the arms are comparable: the two commits here were cherry-picked onto
The p99 separation is clean: every ON leg sits below every OFF leg with a 3.3 s gap. The throughput arms overlap (208.80 ON against 208.63 OFF), so nothing should be claimed there. Why the effect is smaller here than the 19 percent p90 you measured at 32 slots on one node, and it is not a contradiction. A trace of the same configuration says where the time goes. Over a 313 s window, of the 8.45 points the bottleneck GPU spends idle, only 1.13 are with neither GPU busy, and inside the prompt-batch iterations Two things worth taking from this into the PR:
Caveat, stated because the bracket is uneven: five ON legs against two OFF legs. A third OFF leg would make the throughput question answerable rather than merely unclaimed. |
|
@codex review |
|
@codex security review |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 375237a8ce
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // of total, not free: every non-Windows host reports free == total for the CPU device | ||
| return mem_total / 16; |
There was a problem hiding this comment.
In a Linux container whose cgroup limit is substantially smaller than host RAM, this cap is based on host-wide physical memory: the CPU backend implementation in ggml/src/ggml-cpu/ggml-cpu.cpp obtains mem_total from sysconf(_SC_PHYS_PAGES) without consulting cgroups. For example, a 4 GiB pod on a 256 GiB host can therefore retain up to 16 GiB of resident checkpoint buffers, allowing the cgroup OOM killer to terminate the server before the advertised bound or idle trimming helps. The cap needs to incorporate the process/container memory limit when one is present.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 041849375d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // limit_size does not count pooled bytes, so both trims are needed: the first releases | ||
| // what was already pooled, the second the buffers update() just evicted, which the | ||
| // checkpoint destructors hand to the pool rather than to the allocator. | ||
| common_state_buffer_pool::instance().trim(0); |
There was a problem hiding this comment.
When memory is tight and the cache-entry removals above have just returned large checkpoint buffers to the new pool, these buffers can be the reason resize() fails even though the same allocation would have succeeded once they were freed. This trim releases them, but the function then reduces the cache limit and returns without retrying, so an otherwise recoverable prompt-cache save is discarded and may permanently shrink the cache (potentially to one byte when size() is now zero). Retry the allocation after trimming before taking the existing failure path.
Useful? React with 👍 / 👎.
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
What this changes, in four questions1. Before the PR, what happened? Every sequence-state checkpoint allocated a fresh multi-hundred-MiB buffer and freed it again, one per prompt. At that size glibc goes straight to 2. After the PR, what happens? Freed checkpoint buffers go to a bounded process-wide pool instead of the allocator, so the pages stay mapped, resident and dirty and the next checkpoint of the same size reuses them. Retention is capped by bytes (1/16 of host memory) and by count, a full pool only accepts a buffer by displacing a smaller one, and 3. Is this a real issue or a fake one? Real, and describable as a sequence a user hits. Serve any hybrid or recurrent model with prompt caching on. Each prompt writes a checkpoint of hundreds of MiB; each one is a fresh Two related defects found during review and fixed here as well:
4. If merged, does it break anything? This code is not RPC-gated and has no two-node path, so it is live on every
Verified: |
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Simulation against the merge base,
|
| tree | before | after | reclaimed | pool held after |
|---|---|---|---|---|
fbf9abcc7 master, no pool |
4592792 kB | 3909692 kB | 683100 kB | n/a |
head 1db6cb358 |
4592216 kB | 3909084 kB | 683132 kB | 0.000 MiB |
| head minus only those two trims | 4592704 kB | 4218340 kB | 374364 kB | 301.503 MiB |
Without them the recovery loses 308768 kB, exactly the 6 pooled checkpoints of 50.251 MiB, and
the server's own log confirms 301.503 MiB still held against 0.000 MiB with them. With them it
matches master to within 32 kB out of 683 MB.
The isolated version of the same case, in the unit harness: destroying 12 cached checkpoints of
150 MiB reclaims 0 kB, and the subsequent trim(0) reclaims 1843248 kB.
The idle/sleep trim is load bearing too
First measured in steady state with --cache-ram 8192, where the pool's high water mark is 2
buffers because the same buffer is recycled (12 of 14 gets are hits). The trim has nothing to
release and head differs from head-without-the-trim by 428 kB. Recorded rather than discarded: that
is a real configuration in which the flip cannot show.
With --cache-ram 300 so cached prompts are evicted and their checkpoints reach the pool, both arms
reach an identical 502.506 MiB held, hwm 12, from each server's own log:
| tree | RSS at idle_start | RSS after sleep entered |
|---|---|---|
| head | 4195524 kB | 460132 kB |
head minus both trims in server_queue::start_loop |
4195992 kB | 975200 kB |
Difference 515068 kB = 503.0 MiB, matching the pool's own accounting to 0.5 MiB. Without it the
server sleeps still holding half a gigabyte of pooled buffers at the moment it advertises the memory
as released.
The pre-hardening pool wedges, and head fixes it
ee6555ac7 against 1db6cb358, same 43-assertion harness:
| case | ee6555ac7 |
1db6cb358 |
|---|---|---|
| full pool of 64 small buffers, then a 200 MiB buffer | refused | accepted, one smaller evicted |
| 40 growing checkpoints | hwm 64, 0 evictions | hwm 40, 1 eviction |
n_hwm after trim() |
preserved, 64 -> 64 | reset to 0, 4 -> 0 |
Latent use-after-destruction, confirmed by AddressSanitizer
instance() is a function-local static, so it is destroyed in reverse order of completion of
construction. A common_prompt_checkpoint with static storage duration constructed before it would
be destroyed after it and call put() on a destroyed object. [basic.start.term]/5 makes that
undefined outright, and the destructor's try { } catch (...) { } cannot help, because it is not
an exception.
Reproduced under ASan, as a pair, because whether it fires depends on one thing:
| pool at exit | result |
|---|---|
| empty | exit 0, 0 ASan reports |
| holding 96 MiB | exit 1, heap-use-after-free |
ERROR: AddressSanitizer: heap-use-after-free
READ of size 8 at 0x503000002c20 thread T0
#13 common_state_buffer_pool::put(std::vector<unsigned char>&&) common/common.cpp:2364
#14 common_prompt_checkpoint::~common_prompt_checkpoint() common/common.cpp:2406
#15 __run_exit_handlers stdlib/exit.c:108
common.cpp:2364 is free_bufs.push_back(std::move(src)). ~vector deallocates and leaves its
pointers dangling; an empty pool is the case that hides it, because trim()'s shrink_to_fit()
nulls them and the same push_back merely allocates afresh.
This is not reachable today. Every checkpoint holder in the tree has automatic storage
duration, including server_context ctx_server in main() at tools/server/server.cpp:170, so it
is destroyed before static destruction begins. It is a constraint the PR silently creates for
anyone who later gives a checkpoint static storage duration. A leaky singleton
(static auto * p = new common_state_buffer_pool; return *p;) removes it for one deliberate leak.
Everything else checks out
43 of 43 assertions pass on head, under ASan and under TSan. Under TSan, 0 data races and 0
lock-order inversions across 8 concurrent workers doing 1666 get/put cycles with a ninth
thread calling trim(0) every millisecond; each worker tags the first, middle and last byte of its
buffer and reads them back, which would catch a buffer handed to two owners. The only TSan findings
are three instances of the exit-time destruction-order issue above.
- Cap is machine-proportionate, not Spark-tuned.
cap_bytes7.606 GiB against/proc/meminfo
MemTotal 121.689 GiB, a ratio of 0.0625, read independently of the code under test. Derived at
runtime fromsysconf(_SC_PHYS_PAGES); the> (1ull<<50)guard turns a container's-1into
cap_bytes = 0, andput()'scap > cap_bytesthen declines everything, which is master. - The 32 MiB floor is exactly right. glibc's
DEFAULT_MMAP_THRESHOLD_MAXis
4*1024*1024*sizeof(long), 32 MiB on LP64, and the dynamic threshold is only grown under a
<= DEFAULT_MMAP_THRESHOLD_MAXtest, so it can never exceed it. - Eviction terminates in every case, honours both caps, declines an over-cap buffer without
disturbing the pool, and best fit never returns an undersized buffer. - Moves survived the user-declared destructor: move constructible, move assignable, nothrow
move constructible, the move steals the pointer rather than copying, and astd::list::splice
does not move data. lock.unlock()inserver_queue::start_loopis safe. Theunique_lockis declared inside
the innerwhile, so it is destroyed at the end of the iteration and reconstructed at the top of
the next.- Trim really unmaps. 8 x 150 MiB held simultaneously and touched: RSS 75024 -> 1303856 kB,
trim(60 s)correctly a no-op,trim(0)back to 75024 kB.
KV cache and prefix cache
The correctness table in the body uses cache_prompt: false. These runs use cache_prompt: true, temperature 0, top_k 1, seed 42, one slot and one request in flight, seven distinct
prompts, on a recurrent model that both writes and reloads checkpoints:
| reply | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|---|---|---|---|---|---|---|---|
| master | 101effe7 | 64f25d29 | 6c0ac344 | d55410f8 | 7dab2b66 | a8440b7d | 3caf7343 |
| head | 101effe7 | 64f25d29 | 6c0ac344 | d55410f8 | 7dab2b66 | a8440b7d | 3caf7343 |
| head minus the OOM trims | 101effe7 | 64f25d29 | 6c0ac344 | d55410f8 | 7dab2b66 | a8440b7d | 3caf7343 |
All identical. Reply 7 is the one issued while the bad_alloc fires, so the recovery path is inside
the compared output. Handing a checkpoint a buffer that still contains the previous checkpoint's
bytes is safe because update_tgt/update_dft overwrite the whole buffer and abort if
llama_state_seq_get_data_ext does not fill it.
Two smaller notes
trim()setsn_hwm = 0, so thehwmthe server logs understates the peak after any idle period
or OOM recovery.ee6555ac7did not do this.- A monotonically growing checkpoint gets 0 pool hits out of 40 and leaves the pool holding
4440 MiB in 40 buffers, because it never fits anything already pooled and every previous buffer
is retained. Both caps hold and a trim reclaims it all, but the worst case is nearer the byte cap
than the "high water mark 17 buffers, 2.5 GiB" in the body suggests; that number is one workload,
and the cap is what actually bounds it.
…tructor
common_state_buffer_pool::instance() returned a function-local static, which is destroyed in
reverse order of completion of construction. A common_prompt_checkpoint with static storage
duration constructed before it would therefore be destroyed after it, and ~common_prompt_checkpoint
would call put() on a destroyed object. [basic.start.term]/5 makes that undefined outright, and the
try/catch in the destructor does not help, because it is not an exception.
Reproduced under AddressSanitizer with a namespace-scope checkpoint, as a pair, because whether it
fires depends only on whether the pool still holds a buffer when it is destroyed:
pool at exit result
empty exit 0, 0 ASan reports
holding 96 MiB exit 1, heap-use-after-free
ERROR: AddressSanitizer: heap-use-after-free
READ of size 8 thread T0
#13 common_state_buffer_pool::put(std::vector<unsigned char>&&) common/common.cpp:2364
#14 common_prompt_checkpoint::~common_prompt_checkpoint() common/common.cpp:2406
#15 __run_exit_handlers stdlib/exit.c:108
common.cpp:2364 is free_bufs.push_back(). ~vector deallocates and leaves its pointers dangling; an
empty pool is the case that hides it, because trim()'s shrink_to_fit() nulls them and the same
push_back merely allocates afresh.
This is latent today, not live: every checkpoint holder in the tree has automatic storage duration,
including server_context ctx_server in main(). It is a constraint this feature silently creates for
anyone who later gives a checkpoint static storage duration, and a leaky singleton removes it for
the cost of one deliberate leak of a process-wide object at exit.
After this commit the same ASan run is clean, 0 reports, with 43 of 43 pool assertions still
passing under both ASan and ThreadSanitizer (0 data races, 0 lock-order inversions).
|
Pushed A function-local static is destroyed in reverse order of completion of construction, so a
43 of 43 pool assertions still pass, under ASan and under TSan, with 0 data races and 0 lock-order This was latent rather than live, since every checkpoint holder in the tree has automatic storage Nothing else in the PR is touched, and the trims measured in the comment above are unaffected. |
1 participant
A 19% cut in p90 time to first token on a hybrid model at 32 slots, with throughput unchanged
to slightly up. The cost removed is in prefill, on the critical path for TTFT; the decode path
never touches it. Host work per prefill iteration falls from 556.7 ms to 129.7 ms.
A context checkpoint of a hybrid or recurrent model holds the whole non-rollbackable sequence
state. For a 27B
qwen35at 16k context that is 149 MiB, and llama-server allocates andfrees one per prompt. Any allocation that size comes straight from
mmap()and goes straightback on
free(), so the first write to a fresh buffer faults in every one of its pages.That first write is
data_tgt.resize()incommon_prompt_checkpoint::update_tgt(), and on aDGX Spark serving 32 slots it is 43.2 ms of the 54.9 ms a checkpoint costs, which is
454.8 of the 556.7 ms of a prefill iteration.
create_checkpointis 82% of the prefill batchbuild, and the prefill batch build is the single longest host stall in the server. It is a time
to first token cost: the decode path never touches it.
This hands the buffer to a bounded pool instead of to the allocator. The pages stay mapped and
resident, so the next checkpoint reuses them.
The fill is load bearing, and is not removed here
The obvious change is to stop zero filling a buffer that
llama_state_seq_get_data_ext()overwrites immediately. It was tried, and it is wrong. With a CUDA target context the copy into
pageable host memory that is not yet resident runs about 140x slower, 6.5 ms against 917 ms
for 149 MiB, for a 24% throughput regression and triple the TTFT. On the CPU backend the same
change is neutral, so a CPU-only measurement would have shipped it.
This PR does not remove the fill. It makes it disappear by making it unnecessary: on a pool hit
the buffer already has the right size, so
resize()is a no-op over memory that is alreadymapped, resident and dirty, which is the state the fill existed to produce.
Per-phase table
Qwen3.8-27B-UD-Q4_K_XL, one DGX Spark, no RPC,--parallel 32, 32 concurrent, npp 128 ntg 256,--cache-ram 0. Four traced cells, every one of them inside a single thermal-cap window at1690 MHz, cross-checked against the thermal guard's own log. Prefill iterations:
54.85 ms to 20.35 ms per checkpoint, and 556.7 ms to 129.7 ms of host work per prefill
iteration, minus 77%.
create_checkpointalone goes from 454.8 to 129.4 ms per iteration.The residual 14 ms of resize is the cold first wave: a miss still pays the full 43 ms, a hit
pays nothing, and about a third of the checkpoints in a cell are misses.
Decode iterations are untouched:
post_decode6.86 and 6.96 ms/iter on master against6.89 with the pool, decode batch build about 10 us/call either way.
Time to first token
This is where the change is visible to a user. Two independent brackets, base / new / base with
one server load per arm, every arm inside one capped window at 1690 MHz:
TTFT p90 and p99 both fall 19%. Throughput at 32 slots is +3%, and at 8 and 1 slots it is
inside the bracket, which is what one checkpoint per prompt predicts: those cells are not prefill
bound.
A third bracket on the same branch: 99.01 / 100.80 / -- tok/s, TTFT p90 9730.7 to 7966.6 ms,
minus 18.1%. Three brackets, base p90 9538 to 9856 ms and new p90 7687 to 7967 ms.
A second hybrid,
Qwen3.5-4B, A/B/A with the GPU clocks pinned by the harness for the whole run:TTFT p90 falls 34% at 32 slots and 13% at 8, with throughput inside the bracket. The smaller
model has a smaller checkpoint, so the whole of the win shows up as latency.
Neutrality where checkpoints are not created
Two bracketed controls, all arms in one capped state.
-ctxcp 0, same model, checkpoints off: 101.10 / 100.78 / 101.21 tok/s, TTFT p907157 / 7163 / 7176 ms.
qwen2arch, no recurrent state and no SWA, so checkpoints are never created: 1197.6 /1274.5 / 1217.7 tok/s at 32 slots, 555.9 / 593.6 / 561.7 at 8, 115.6 / 113.8 / 114.2 at 1.
Nothing below
MIN_BUFFER_BYTES(32 MiB, the allocator's mmap threshold) is ever pooled, so fora model whose checkpoints are small the pool is never entered at all: an SWA
tinygemma3makes0.356 MiB checkpoints and the pool declines every one of them.
Memory policy
Documented on
common_state_buffer_pool. A byte cap of 1/16 of total host memory, a count cap of64 buffers, a 32 MiB size floor, an eviction rule that displaces the smallest pooled buffer and
only one smaller than the buffer coming in, and
trim()from the task queue's idle wait and fromthe prompt cache's out-of-memory recovery. Every cap declines the buffer and lets it be freed,
which is exactly master's behaviour, so a machine that cannot afford the pool degrades to today
rather than to something worse.
Worst case extra resident memory is the byte cap, 7.59 GiB on a 121 GiB machine. The measured
case is far below it, because the pool only ever receives buffers the process had just freed and
hands them straight back out: high water mark 17 buffers, 2.5 GiB, and the server's VmHWM is
6.4738 GB with the pool against 6.4677 GB without it, a rise of about 6 MB.
A negative result, measured and dropped
A fourth commit wrote the recycled buffer once before the state copy, on the theory that the
device-to-host copy is fastest into host pages the CPU wrote last. It was dropped. Measured
A/B/A with all three cells inside one capped window, microseconds per
create_checkpoint:The pass costs 2.2 ms and saves 0.7. The effect is real but does not pay for itself, and whole
cell throughput at 32 slots was 100.32 / 102.27 / 102.96 tok/s, inside the bracket.
The numbers that originally justified it came from two cells that were not in the same clock
state. The useful lesson is not "check the clock": it is that a two-cell A against B has nothing
in it that can disagree with itself, so any difference it shows is indistinguishable from drift.
Every comparison above is A/B/A, and every cell in it is cross-referenced against the machine's
thermal guard log to confirm all three arms were in one clock state.
Correctness
Greedy, temperature 0, top_k 1, seed 42,
cache_promptfalse, one slot and one request inflight so the batch composition is fixed, md5 of the concatenated output:
tests/test-recurrent-state-rollbackon the 4B hybrid produces output byte identical to master,including the same pre-existing dirty-ctx mismatch.
The single-slot harness is used deliberately. With several requests decoded in one batch these
models are not run to run reproducible: base against base produced three different md5s from one
binary, so a concurrent harness is not usable as a correctness control.