server: reuse checkpoint state buffers from a bounded pool - #201
server: reuse checkpoint state buffers from a bounded pool#201danielhanchen wants to merge 2 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. |
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.