From ee6555ac7e6a0fe5d3ddfaf2b3b1f173f2086d01 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 6 Sep 2026 10:19:24 -0700 Subject: [PATCH 1/6] server: reuse checkpoint state buffers from a bounded pool 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. --- common/common.cpp | 131 +++++++++++++++++++++++++++++++- common/common.h | 102 +++++++++++++++++++++++++ tools/server/server-context.cpp | 12 ++- 3 files changed, 241 insertions(+), 4 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d3..6b4352cc57d9 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2251,6 +2251,133 @@ bool common_prompt_batch_decode( return true; } +common_state_buffer_pool & common_state_buffer_pool::instance() { + static common_state_buffer_pool pool; + return pool; +} + +// derived once, from host memory as the CPU backend device reports it. an unknown host memory +// size means no pooling at all, which is the safe direction. +static size_t common_state_buffer_pool_cap() { + size_t mem_free = 0; + size_t mem_total = 0; + + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (cpu_dev != nullptr) { + ggml_backend_dev_memory(cpu_dev, &mem_free, &mem_total); + } + + if (mem_total == 0) { + return 0; + } + + return std::min(mem_total / 16, mem_free / 4); +} + +void common_state_buffer_pool::get(std::vector & dst, size_t size) { + { + std::lock_guard lock(mtx); + + st.n_get++; + + t_last_us = ggml_time_us(); + + // an allocation that is already big enough is already resident - nothing to do + if (dst.capacity() < size) { + // best fit, so a large pooled buffer is not spent on a small request + size_t best = free_bufs.size(); + for (size_t i = 0; i < free_bufs.size(); ++i) { + if (free_bufs[i].capacity() >= size && + (best == free_bufs.size() || free_bufs[i].capacity() < free_bufs[best].capacity())) { + best = i; + } + } + + if (best < free_bufs.size()) { + held_bytes -= free_bufs[best].capacity(); + + // the pooled buffer goes to the caller and the caller's undersized one is + // dropped: it belongs to a different size class and is of no use here + std::swap(dst, free_bufs[best]); + free_bufs.erase(free_bufs.begin() + best); + + st.n_hit++; + } + } + } + + // on a hit the pooled buffer normally already has exactly this size, so this is a no-op and + // the stale bytes are left in place. that is safe because every caller overwrites the whole + // buffer and aborts if the state does not fill it. when the sizes differ, the fill only + // touches pages that are already resident. + dst.resize(size); +} + +void common_state_buffer_pool::put(std::vector && src) { + const size_t cap = src.capacity(); + + if (cap < MIN_BUFFER_BYTES) { + return; + } + + std::lock_guard lock(mtx); + + st.n_put++; + + t_last_us = ggml_time_us(); + + if (!cap_known) { + cap_bytes = common_state_buffer_pool_cap(); + cap_known = true; + } + + if (free_bufs.size() >= MAX_BUFFERS || held_bytes + cap > cap_bytes) { + return; // declined: src is freed by its own destructor, as it would be without the pool + } + + // the logical size is kept, not cleared: a later get() of the same size is then a no-op + // resize rather than a full zero fill of the buffer + free_bufs.push_back(std::move(src)); + + held_bytes += cap; + + n_hwm = std::max(n_hwm, free_bufs.size()); + + st.n_keep++; +} + +void common_state_buffer_pool::trim(int64_t idle_us) { + std::lock_guard lock(mtx); + + if (free_bufs.empty() || ggml_time_us() - t_last_us < idle_us) { + return; + } + + free_bufs.clear(); + free_bufs.shrink_to_fit(); + + held_bytes = 0; +} + +common_state_buffer_pool::stats common_state_buffer_pool::get_stats() const { + std::lock_guard lock(mtx); + + stats res = st; + + res.held_bytes = held_bytes; + res.cap_bytes = cap_bytes; + res.n_hwm = n_hwm; + + return res; +} + +common_prompt_checkpoint::~common_prompt_checkpoint() { + auto & pool = common_state_buffer_pool::instance(); + + pool.put(std::move(data_tgt)); + pool.put(std::move(data_dft)); +} + size_t common_prompt_checkpoint::size() const { return data_tgt.size() + data_dft.size() + data_spec.size(); } @@ -2289,7 +2416,7 @@ void common_prompt_checkpoint::update_tgt( const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags); - data_tgt.resize(ckpt_size); + common_state_buffer_pool::instance().get(data_tgt, ckpt_size); const size_t n = llama_state_seq_get_data_ext(ctx, data_tgt.data(), ckpt_size, seq_id, flags); if (n != ckpt_size) { @@ -2307,7 +2434,7 @@ void common_prompt_checkpoint::update_dft( const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags); - data_dft.resize(ckpt_size); + common_state_buffer_pool::instance().get(data_dft, ckpt_size); const size_t n = llama_state_seq_get_data_ext(ctx, data_dft.data(), ckpt_size, seq_id, flags); if (n != ckpt_size) { diff --git a/common/common.h b/common/common.h index de49dac9f63a..4aeff0ae5121 100644 --- a/common/common.h +++ b/common/common.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -1134,6 +1135,93 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); // prompt utils // +// A bounded, process-wide pool of reusable byte buffers for sequence-state checkpoints. +// +// Checkpoint buffers are large - for a hybrid/recurrent model the non-rollbackable state of a +// single sequence runs to hundreds of MiB - and one is allocated and freed per prompt. Any +// allocation that size comes straight from mmap() and goes straight back on free, so the first +// write to a fresh buffer faults in every one of its pages. In llama-server that shows up as a +// ~44 ms zero fill inside `update_tgt()` before the state copy can even start, repeated for +// every checkpoint of every prompt. +// +// Handing the buffer to this pool instead of to the allocator keeps the pages mapped, resident +// and dirty, so the next checkpoint of the same size reuses them and pays neither the faults +// nor the fill. +// +// Note that the zero fill itself is load bearing and must not simply be removed: with a CUDA +// target context, `llama_state_seq_get_data_ext()` copying into pageable host memory that is +// not yet resident runs about 140x slower than into memory that is. The pool removes the fill +// by making it unnecessary, not by skipping it. +// +// Memory policy. +// +// - Byte cap. At most `cap_bytes` is retained, derived once from host memory as +// min(1/16 of total, 1/4 of free). A machine with little memory therefore keeps little or +// nothing, and when host memory cannot be determined at all the pool keeps nothing. +// - Count cap. At most `MAX_BUFFERS` buffers, a sanity bound; the byte cap is the real one. +// - Size floor. Buffers below `MIN_BUFFER_BYTES` are never retained. Below the allocator's +// mmap threshold there is no fault storm to avoid, so pooling them would only hold memory. +// - Release. `trim()` drops buffers the pool has not touched for a while, so an idle server +// gives the memory back instead of holding it for the life of the process. +// +// When any cap is reached the buffer is simply freed, which is exactly the behaviour without +// the pool: a machine that cannot afford the pool degrades to today's behaviour, never to +// something worse. +// +// On the peak. The pool only ever receives buffers that the process had just freed, and hands +// them straight back out, so for a steady workload the sum of live plus pooled buffers is the +// count the process already peaked at without the pool. The byte cap and `trim()` bound the +// one case where that is not true: a server whose live checkpoint count structurally shrinks. +struct common_state_buffer_pool { + // sanity bound on the buffer count; the byte cap is the limit that actually binds + static constexpr size_t MAX_BUFFERS = 64; + + // buffers below this size are left to the allocator: they do not come from mmap(), so + // reusing them saves nothing + static constexpr size_t MIN_BUFFER_BYTES = 32ull*1024*1024; + + struct stats { + uint64_t n_get = 0; // buffers requested + uint64_t n_hit = 0; // ... served from the pool + uint64_t n_put = 0; // buffers offered back + uint64_t n_keep = 0; // ... retained + size_t held_bytes = 0; + size_t cap_bytes = 0; + size_t n_hwm = 0; // high water mark of retained buffers + }; + + // resize `dst` to `size`, reusing a pooled allocation when one fits. + // the previous contents of `dst` are not preserved; every caller overwrites the whole + // buffer immediately, and preserving them would defeat the point of the reuse. + void get(std::vector & dst, size_t size); + + // offer `src` to the pool. if the policy declines it, `src` is left alone and freed by its + // own destructor, exactly as it would be without the pool. + void put(std::vector && src); + + // free every pooled buffer if the pool has not been used for `idle_us`. cheap and safe to + // call often; the server calls it whenever all of its slots are idle. + void trim(int64_t idle_us); + + stats get_stats() const; + + static common_state_buffer_pool & instance(); + +private: + mutable std::mutex mtx; + + std::vector> free_bufs; + + size_t held_bytes = 0; + size_t cap_bytes = 0; + size_t n_hwm = 0; + bool cap_known = false; + + int64_t t_last_us = 0; + + stats st; +}; + struct common_prompt_checkpoint { int64_t n_tokens; @@ -1150,6 +1238,20 @@ struct common_prompt_checkpoint { // (e.g. eagle3's deferred-boundary g_embd row) std::vector data_spec; + common_prompt_checkpoint() = default; + + // returns the state buffers to common_state_buffer_pool instead of to the allocator. + // the copy and move members are defaulted explicitly: declaring a destructor would + // otherwise suppress the implicit moves and turn every list splice into a deep copy of a + // multi-hundred-MiB buffer. + ~common_prompt_checkpoint(); + + common_prompt_checkpoint(const common_prompt_checkpoint &) = default; + common_prompt_checkpoint(common_prompt_checkpoint &&) noexcept = default; + + common_prompt_checkpoint & operator=(const common_prompt_checkpoint &) = default; + common_prompt_checkpoint & operator=(common_prompt_checkpoint &&) noexcept = default; + size_t size() const; bool empty() const; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..6dc62cc70bcf 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2249,10 +2249,13 @@ struct server_context_impl { // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); + const auto pool = common_state_buffer_pool::instance().get_stats(); + SLT_TRC(slot, - "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", + "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB, buffer pool: %" PRIu64 "/%" PRIu64 " reused, %.3f of %.3f MiB held, hwm %zu)\n", (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, - cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); + cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024, + pool.n_hit, pool.n_get, pool.held_bytes / (1024.0 * 1024.0), pool.cap_bytes / (1024.0 * 1024.0), pool.n_hwm); } // returns false to decline the task, it is offered again after the decode is done @@ -2702,6 +2705,11 @@ struct server_context_impl { if (all_idle) { SRV_TRC("%s", "all slots are idle\n"); + // give the checkpoint buffers back once the server has been quiet for a while. + // the delay matters: back-to-back request waves go briefly all-idle between + // waves, and dropping the pool there would throw away every reuse. + common_state_buffer_pool::instance().trim(30ll*1000*1000); + metrics_flush_idle(); return; // skip further processing From 468cf0b4d92db6d066a070452d1dbe375018406d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 6 Sep 2026 12:35:48 -0700 Subject: [PATCH 2/6] server: harden the checkpoint buffer pool's memory policy 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. --- common/common.cpp | 83 ++++++++++++++++++++++++++------- common/common.h | 26 +++++++---- tools/server/server-context.cpp | 9 +--- tools/server/server-queue.cpp | 10 ++++ tools/server/server-task.cpp | 6 +++ 5 files changed, 103 insertions(+), 31 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 6b4352cc57d9..8da6e9eb5769 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2256,8 +2256,8 @@ common_state_buffer_pool & common_state_buffer_pool::instance() { return pool; } -// derived once, from host memory as the CPU backend device reports it. an unknown host memory -// size means no pooling at all, which is the safe direction. +// derived once, from host memory as the CPU backend device reports it. a host memory size that +// cannot be trusted means no pooling at all, which is the safe direction. static size_t common_state_buffer_pool_cap() { size_t mem_free = 0; size_t mem_total = 0; @@ -2267,14 +2267,25 @@ static size_t common_state_buffer_pool_cap() { ggml_backend_dev_memory(cpu_dev, &mem_free, &mem_total); } - if (mem_total == 0) { + GGML_UNUSED(mem_free); + + // zero means the CPU device could not report a size. an absurd size means the query itself + // failed without saying so: ggml_backend_cpu_device_get_memory() multiplies out + // sysconf(_SC_PHYS_PAGES) with no error check, so a -1 from a restricted container arrives + // here as ~1.8e19. treat both as unknown and keep nothing. + if (mem_total == 0 || mem_total > (1ull << 50)) { return 0; } - return std::min(mem_total / 16, mem_free / 4); + // a fraction of TOTAL host memory, deliberately not of free: every non-Windows host reports + // free == total for the CPU device ("free system memory is ill-defined, assume all of it is + // free"), so a free-based cap would be the same number with a false claim attached to it. + return mem_total / 16; } void common_state_buffer_pool::get(std::vector & dst, size_t size) { + std::vector prev; // the caller's previous storage, offered back below + { std::lock_guard lock(mtx); @@ -2282,8 +2293,11 @@ void common_state_buffer_pool::get(std::vector & dst, size_t size) { t_last_us = ggml_time_us(); - // an allocation that is already big enough is already resident - nothing to do - if (dst.capacity() < size) { + if (dst.capacity() >= size) { + // the caller's own allocation is already big enough, so it is already resident. + // that is a reuse too, and the cheapest kind. + st.n_hit++; + } else { // best fit, so a large pooled buffer is not spent on a small request size_t best = free_bufs.size(); for (size_t i = 0; i < free_bufs.size(); ++i) { @@ -2296,16 +2310,21 @@ void common_state_buffer_pool::get(std::vector & dst, size_t size) { if (best < free_bufs.size()) { held_bytes -= free_bufs[best].capacity(); - // the pooled buffer goes to the caller and the caller's undersized one is - // dropped: it belongs to a different size class and is of no use here - std::swap(dst, free_bufs[best]); + prev = std::move(free_bufs[best]); free_bufs.erase(free_bufs.begin() + best); + // dst takes the pooled buffer and prev takes the caller's old one + std::swap(dst, prev); + st.n_hit++; } } } + // the caller's old buffer was too small for this request but may still serve a smaller + // checkpoint later, so offer it back rather than dropping it on the floor + put(std::move(prev)); + // on a hit the pooled buffer normally already has exactly this size, so this is a no-op and // the stale bytes are left in place. that is safe because every caller overwrites the whole // buffer and aborts if the state does not fill it. when the sizes differ, the fill only @@ -2317,7 +2336,7 @@ void common_state_buffer_pool::put(std::vector && src) { const size_t cap = src.capacity(); if (cap < MIN_BUFFER_BYTES) { - return; + return; // below the allocator's mmap threshold: holding it would save nothing } std::lock_guard lock(mtx); @@ -2331,8 +2350,31 @@ void common_state_buffer_pool::put(std::vector && src) { cap_known = true; } - if (free_bufs.size() >= MAX_BUFFERS || held_bytes + cap > cap_bytes) { - return; // declined: src is freed by its own destructor, as it would be without the pool + if (cap > cap_bytes) { + return; // one buffer of this size would spend the whole budget + } + + // eviction. make room by dropping the smallest pooled buffers, but only ones smaller than + // the buffer coming in. a pool already full of buffers at least this large has nothing to + // gain from the swap, so decline instead. that rule is what keeps a workload whose + // checkpoints grow through a prompt from wedging the pool full of small buffers that no + // later and larger request can use. + while (free_bufs.size() >= MAX_BUFFERS || held_bytes + cap > cap_bytes) { + size_t worst = 0; + for (size_t i = 1; i < free_bufs.size(); ++i) { + if (free_bufs[i].capacity() < free_bufs[worst].capacity()) { + worst = i; + } + } + + if (free_bufs.empty() || free_bufs[worst].capacity() >= cap) { + return; // declined: src is freed by its own destructor, as it would be without the pool + } + + held_bytes -= free_bufs[worst].capacity(); + free_bufs.erase(free_bufs.begin() + worst); + + st.n_evict++; } // the logical size is kept, not cleared: a later get() of the same size is then a no-op @@ -2357,6 +2399,7 @@ void common_state_buffer_pool::trim(int64_t idle_us) { free_bufs.shrink_to_fit(); held_bytes = 0; + n_hwm = 0; } common_state_buffer_pool::stats common_state_buffer_pool::get_stats() const { @@ -2372,10 +2415,18 @@ common_state_buffer_pool::stats common_state_buffer_pool::get_stats() const { } common_prompt_checkpoint::~common_prompt_checkpoint() { - auto & pool = common_state_buffer_pool::instance(); - - pool.put(std::move(data_tgt)); - pool.put(std::move(data_dft)); + // a destructor is noexcept and put() can throw: it locks a mutex and pushes to a vector. + // that matters on exactly the path this has to survive, because server_prompt_cache::alloc() + // recovers from a std::bad_alloc by destroying cached prompts, and a throw from here would + // turn a graceful cache shrink into a terminate(). + try { + auto & pool = common_state_buffer_pool::instance(); + + pool.put(std::move(data_tgt)); + pool.put(std::move(data_dft)); + } catch (...) { + // the buffers are freed by their own destructors instead + } } size_t common_prompt_checkpoint::size() const { diff --git a/common/common.h b/common/common.h index 4aeff0ae5121..f74784dafbbc 100644 --- a/common/common.h +++ b/common/common.h @@ -1155,14 +1155,23 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); // // Memory policy. // -// - Byte cap. At most `cap_bytes` is retained, derived once from host memory as -// min(1/16 of total, 1/4 of free). A machine with little memory therefore keeps little or -// nothing, and when host memory cannot be determined at all the pool keeps nothing. -// - Count cap. At most `MAX_BUFFERS` buffers, a sanity bound; the byte cap is the real one. +// - Byte cap. At most `cap_bytes` is retained, derived once as 1/16 of total host memory, so +// the pool is proportionate to the machine it runs on. When host memory cannot be +// determined the pool keeps nothing. +// - Count cap. At most `MAX_BUFFERS` buffers. With `n_ctx_checkpoints` defaulting to 32 per +// slot a busy server can release far more than that at once, so this cap does bind, and +// the eviction rule below decides which buffers survive it. +// - Eviction. A buffer offered to a full pool displaces the smallest pooled buffer, but only +// one smaller than itself; otherwise it is declined. Without that rule a workload whose +// checkpoints grow through a prompt would fill the pool with small buffers that no later +// request can use, and the hit rate would fall to zero while the memory stayed held. // - Size floor. Buffers below `MIN_BUFFER_BYTES` are never retained. Below the allocator's // mmap threshold there is no fault storm to avoid, so pooling them would only hold memory. -// - Release. `trim()` drops buffers the pool has not touched for a while, so an idle server -// gives the memory back instead of holding it for the life of the process. +// - Release. `trim()` drops buffers the pool has not touched for a while. The server calls it +// from the task queue's idle wait, so a server that goes quiet gives the memory back, and +// with a zero timeout from the prompt cache's out-of-memory recovery, so pooled bytes are +// always reclaimable under allocation pressure even though they are not counted against +// `--cache-ram`. // // When any cap is reached the buffer is simply freed, which is exactly the behaviour without // the pool: a machine that cannot afford the pool degrades to today's behaviour, never to @@ -1173,7 +1182,7 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); // count the process already peaked at without the pool. The byte cap and `trim()` bound the // one case where that is not true: a server whose live checkpoint count structurally shrinks. struct common_state_buffer_pool { - // sanity bound on the buffer count; the byte cap is the limit that actually binds + // bound on the buffer count. this does bind in practice, see the eviction rule above static constexpr size_t MAX_BUFFERS = 64; // buffers below this size are left to the allocator: they do not come from mmap(), so @@ -1184,7 +1193,8 @@ struct common_state_buffer_pool { uint64_t n_get = 0; // buffers requested uint64_t n_hit = 0; // ... served from the pool uint64_t n_put = 0; // buffers offered back - uint64_t n_keep = 0; // ... retained + uint64_t n_keep = 0; // ... retained + uint64_t n_evict = 0; // pooled buffers displaced to make room size_t held_bytes = 0; size_t cap_bytes = 0; size_t n_hwm = 0; // high water mark of retained buffers diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6dc62cc70bcf..7f3870ac5b82 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2252,10 +2252,10 @@ struct server_context_impl { const auto pool = common_state_buffer_pool::instance().get_stats(); SLT_TRC(slot, - "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB, buffer pool: %" PRIu64 "/%" PRIu64 " reused, %.3f of %.3f MiB held, hwm %zu)\n", + "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB, buffer pool: %" PRIu64 "/%" PRIu64 " reused, %" PRIu64 " evicted, %.3f of %.3f MiB held, hwm %zu)\n", (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024, - pool.n_hit, pool.n_get, pool.held_bytes / (1024.0 * 1024.0), pool.cap_bytes / (1024.0 * 1024.0), pool.n_hwm); + pool.n_hit, pool.n_get, pool.n_evict, pool.held_bytes / (1024.0 * 1024.0), pool.cap_bytes / (1024.0 * 1024.0), pool.n_hwm); } // returns false to decline the task, it is offered again after the decode is done @@ -2705,11 +2705,6 @@ struct server_context_impl { if (all_idle) { SRV_TRC("%s", "all slots are idle\n"); - // give the checkpoint buffers back once the server has been quiet for a while. - // the delay matters: back-to-back request waves go briefly all-idle between - // waves, and dropping the pool there would throw away every reuse. - common_state_buffer_pool::instance().trim(30ll*1000*1000); - metrics_flush_idle(); return; // skip further processing diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..b42362c70aa0 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -1,3 +1,4 @@ +#include "common.h" #include "server-task.h" #include "server-queue.h" @@ -357,6 +358,15 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { if (res) { break; // new task arrived or terminate } + + // nothing to do. hand back any checkpoint state buffers the pool is still + // holding, so a server that has gone quiet does not keep them for the life of + // the process. this is the only loop that runs while the queue is empty, so it + // is the only place the release can happen. done without the lock: trim() + // unmaps hundreds of MiB and must not delay an arriving task. + lock.unlock(); + common_state_buffer_pool::instance().trim(30ll*1000*1000); + // otherwise, loop again to check sleeping condition } } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..ab562b744812 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1767,6 +1767,12 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro } catch (const std::bad_alloc & e) { SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what()); + // the checkpoint buffer pool holds host memory that is reusable but not currently in + // use, and it is not counted against limit_size. release all of it before shrinking the + // cache: without this the retry reclaims nothing, because the cached prompts destroyed + // by update() below hand their buffers to the pool instead of to the allocator. + common_state_buffer_pool::instance().trim(0); + limit_size = std::max(1, 0.4*size()); SRV_WRN(" - cache size limit reduced to %.3f MiB\n", limit_size / (1024.0 * 1024.0)); From 375237a8ce8bb523ea6e79ba361be5b01a80789b Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 18:41:34 -0700 Subject: [PATCH 3/6] server: trim comments in the checkpoint buffer pool --- common/common.cpp | 40 +++++--------------- common/common.h | 69 ++++++----------------------------- tools/server/server-queue.cpp | 7 +--- tools/server/server-task.cpp | 6 +-- 4 files changed, 25 insertions(+), 97 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 8da6e9eb5769..210e88b616ca 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2256,8 +2256,6 @@ common_state_buffer_pool & common_state_buffer_pool::instance() { return pool; } -// derived once, from host memory as the CPU backend device reports it. a host memory size that -// cannot be trusted means no pooling at all, which is the safe direction. static size_t common_state_buffer_pool_cap() { size_t mem_free = 0; size_t mem_total = 0; @@ -2269,17 +2267,13 @@ static size_t common_state_buffer_pool_cap() { GGML_UNUSED(mem_free); - // zero means the CPU device could not report a size. an absurd size means the query itself - // failed without saying so: ggml_backend_cpu_device_get_memory() multiplies out - // sysconf(_SC_PHYS_PAGES) with no error check, so a -1 from a restricted container arrives - // here as ~1.8e19. treat both as unknown and keep nothing. + // ggml_backend_cpu_device_get_memory() multiplies out sysconf(_SC_PHYS_PAGES) unchecked, so a + // -1 from a restricted container arrives as ~1.8e19: absurd means unknown, not huge. if (mem_total == 0 || mem_total > (1ull << 50)) { return 0; } - // a fraction of TOTAL host memory, deliberately not of free: every non-Windows host reports - // free == total for the CPU device ("free system memory is ill-defined, assume all of it is - // free"), so a free-based cap would be the same number with a false claim attached to it. + // of total, not free: every non-Windows host reports free == total for the CPU device return mem_total / 16; } @@ -2294,8 +2288,6 @@ void common_state_buffer_pool::get(std::vector & dst, size_t size) { t_last_us = ggml_time_us(); if (dst.capacity() >= size) { - // the caller's own allocation is already big enough, so it is already resident. - // that is a reuse too, and the cheapest kind. st.n_hit++; } else { // best fit, so a large pooled buffer is not spent on a small request @@ -2313,7 +2305,6 @@ void common_state_buffer_pool::get(std::vector & dst, size_t size) { prev = std::move(free_bufs[best]); free_bufs.erase(free_bufs.begin() + best); - // dst takes the pooled buffer and prev takes the caller's old one std::swap(dst, prev); st.n_hit++; @@ -2321,14 +2312,11 @@ void common_state_buffer_pool::get(std::vector & dst, size_t size) { } } - // the caller's old buffer was too small for this request but may still serve a smaller - // checkpoint later, so offer it back rather than dropping it on the floor + // too small here, but may still serve a smaller checkpoint later put(std::move(prev)); - // on a hit the pooled buffer normally already has exactly this size, so this is a no-op and - // the stale bytes are left in place. that is safe because every caller overwrites the whole - // buffer and aborts if the state does not fill it. when the sizes differ, the fill only - // touches pages that are already resident. + // usually a no-op leaving stale bytes: callers overwrite the whole buffer and abort if the + // state does not fill it dst.resize(size); } @@ -2354,11 +2342,6 @@ void common_state_buffer_pool::put(std::vector && src) { return; // one buffer of this size would spend the whole budget } - // eviction. make room by dropping the smallest pooled buffers, but only ones smaller than - // the buffer coming in. a pool already full of buffers at least this large has nothing to - // gain from the swap, so decline instead. that rule is what keeps a workload whose - // checkpoints grow through a prompt from wedging the pool full of small buffers that no - // later and larger request can use. while (free_bufs.size() >= MAX_BUFFERS || held_bytes + cap > cap_bytes) { size_t worst = 0; for (size_t i = 1; i < free_bufs.size(); ++i) { @@ -2377,8 +2360,7 @@ void common_state_buffer_pool::put(std::vector && src) { st.n_evict++; } - // the logical size is kept, not cleared: a later get() of the same size is then a no-op - // resize rather than a full zero fill of the buffer + // size kept, not cleared: a later get() of the same size is then a no-op resize free_bufs.push_back(std::move(src)); held_bytes += cap; @@ -2415,17 +2397,15 @@ common_state_buffer_pool::stats common_state_buffer_pool::get_stats() const { } common_prompt_checkpoint::~common_prompt_checkpoint() { - // a destructor is noexcept and put() can throw: it locks a mutex and pushes to a vector. - // that matters on exactly the path this has to survive, because server_prompt_cache::alloc() - // recovers from a std::bad_alloc by destroying cached prompts, and a throw from here would - // turn a graceful cache shrink into a terminate(). + // put() locks and pushes, so it can throw, and a destructor is noexcept. that is fatal on the + // one path this must survive: server_prompt_cache::alloc() destroys cached prompts to recover + // from std::bad_alloc, and a throw here would make that graceful shrink a terminate(). try { auto & pool = common_state_buffer_pool::instance(); pool.put(std::move(data_tgt)); pool.put(std::move(data_dft)); } catch (...) { - // the buffers are freed by their own destructors instead } } diff --git a/common/common.h b/common/common.h index f74784dafbbc..0561eaeb5bcb 100644 --- a/common/common.h +++ b/common/common.h @@ -1135,58 +1135,19 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); // prompt utils // -// A bounded, process-wide pool of reusable byte buffers for sequence-state checkpoints. +// Bounded process-wide pool of reusable checkpoint state buffers. At hundreds of MiB each these +// come from mmap() and go back on free, so every fresh buffer faults in all of its pages. // -// Checkpoint buffers are large - for a hybrid/recurrent model the non-rollbackable state of a -// single sequence runs to hundreds of MiB - and one is allocated and freed per prompt. Any -// allocation that size comes straight from mmap() and goes straight back on free, so the first -// write to a fresh buffer faults in every one of its pages. In llama-server that shows up as a -// ~44 ms zero fill inside `update_tgt()` before the state copy can even start, repeated for -// every checkpoint of every prompt. +// Do NOT instead delete the zero fill this avoids: with a CUDA target context, +// llama_state_seq_get_data_ext() into non-resident pageable host memory runs about 140x slower. // -// Handing the buffer to this pool instead of to the allocator keeps the pages mapped, resident -// and dirty, so the next checkpoint of the same size reuses them and pays neither the faults -// nor the fill. -// -// Note that the zero fill itself is load bearing and must not simply be removed: with a CUDA -// target context, `llama_state_seq_get_data_ext()` copying into pageable host memory that is -// not yet resident runs about 140x slower than into memory that is. The pool removes the fill -// by making it unnecessary, not by skipping it. -// -// Memory policy. -// -// - Byte cap. At most `cap_bytes` is retained, derived once as 1/16 of total host memory, so -// the pool is proportionate to the machine it runs on. When host memory cannot be -// determined the pool keeps nothing. -// - Count cap. At most `MAX_BUFFERS` buffers. With `n_ctx_checkpoints` defaulting to 32 per -// slot a busy server can release far more than that at once, so this cap does bind, and -// the eviction rule below decides which buffers survive it. -// - Eviction. A buffer offered to a full pool displaces the smallest pooled buffer, but only -// one smaller than itself; otherwise it is declined. Without that rule a workload whose -// checkpoints grow through a prompt would fill the pool with small buffers that no later -// request can use, and the hit rate would fall to zero while the memory stayed held. -// - Size floor. Buffers below `MIN_BUFFER_BYTES` are never retained. Below the allocator's -// mmap threshold there is no fault storm to avoid, so pooling them would only hold memory. -// - Release. `trim()` drops buffers the pool has not touched for a while. The server calls it -// from the task queue's idle wait, so a server that goes quiet gives the memory back, and -// with a zero timeout from the prompt cache's out-of-memory recovery, so pooled bytes are -// always reclaimable under allocation pressure even though they are not counted against -// `--cache-ram`. -// -// When any cap is reached the buffer is simply freed, which is exactly the behaviour without -// the pool: a machine that cannot afford the pool degrades to today's behaviour, never to -// something worse. -// -// On the peak. The pool only ever receives buffers that the process had just freed, and hands -// them straight back out, so for a steady workload the sum of live plus pooled buffers is the -// count the process already peaked at without the pool. The byte cap and `trim()` bound the -// one case where that is not true: a server whose live checkpoint count structurally shrinks. +// A full pool only accepts a buffer by displacing a SMALLER one, else checkpoints growing through +// a prompt wedge it full of small buffers no later request can use. trim() keeps pooled bytes +// reclaimable, since they do not count against `--cache-ram`. struct common_state_buffer_pool { - // bound on the buffer count. this does bind in practice, see the eviction rule above static constexpr size_t MAX_BUFFERS = 64; - // buffers below this size are left to the allocator: they do not come from mmap(), so - // reusing them saves nothing + // below the allocator's mmap threshold there is no fault storm to avoid static constexpr size_t MIN_BUFFER_BYTES = 32ull*1024*1024; struct stats { @@ -1200,17 +1161,11 @@ struct common_state_buffer_pool { size_t n_hwm = 0; // high water mark of retained buffers }; - // resize `dst` to `size`, reusing a pooled allocation when one fits. - // the previous contents of `dst` are not preserved; every caller overwrites the whole - // buffer immediately, and preserving them would defeat the point of the reuse. + // contents of `dst` are NOT preserved; every caller overwrites the whole buffer void get(std::vector & dst, size_t size); - // offer `src` to the pool. if the policy declines it, `src` is left alone and freed by its - // own destructor, exactly as it would be without the pool. void put(std::vector && src); - // free every pooled buffer if the pool has not been used for `idle_us`. cheap and safe to - // call often; the server calls it whenever all of its slots are idle. void trim(int64_t idle_us); stats get_stats() const; @@ -1250,10 +1205,8 @@ struct common_prompt_checkpoint { common_prompt_checkpoint() = default; - // returns the state buffers to common_state_buffer_pool instead of to the allocator. - // the copy and move members are defaulted explicitly: declaring a destructor would - // otherwise suppress the implicit moves and turn every list splice into a deep copy of a - // multi-hundred-MiB buffer. + // the copy/move members below are defaulted explicitly: declaring this destructor suppresses + // the implicit moves, turning every list splice into a deep copy of a multi-hundred-MiB buffer. ~common_prompt_checkpoint(); common_prompt_checkpoint(const common_prompt_checkpoint &) = default; diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index b42362c70aa0..16f11f7b8f49 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -359,11 +359,8 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { break; // new task arrived or terminate } - // nothing to do. hand back any checkpoint state buffers the pool is still - // holding, so a server that has gone quiet does not keep them for the life of - // the process. this is the only loop that runs while the queue is empty, so it - // is the only place the release can happen. done without the lock: trim() - // unmaps hundreds of MiB and must not delay an arriving task. + // the only loop running while the queue is empty, so the only place a quiet server + // gives pooled buffers back. unlocked: trim() unmaps hundreds of MiB. lock.unlock(); common_state_buffer_pool::instance().trim(30ll*1000*1000); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index ab562b744812..be36e12c533c 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1767,10 +1767,8 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro } catch (const std::bad_alloc & e) { SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what()); - // the checkpoint buffer pool holds host memory that is reusable but not currently in - // use, and it is not counted against limit_size. release all of it before shrinking the - // cache: without this the retry reclaims nothing, because the cached prompts destroyed - // by update() below hand their buffers to the pool instead of to the allocator. + // release first or the retry reclaims nothing: update() below hands the destroyed prompts' + // buffers to the pool, which limit_size does not count, rather than to the allocator common_state_buffer_pool::instance().trim(0); limit_size = std::max(1, 0.4*size()); From 041849375daf2051946b991c61722d4e75b46034 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 18:51:47 -0700 Subject: [PATCH 4/6] server: trim the checkpoint pool after evicting cache entries too --- tools/server/server-task.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index be36e12c533c..85a437cc6039 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1767,8 +1767,9 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro } catch (const std::bad_alloc & e) { SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what()); - // release first or the retry reclaims nothing: update() below hands the destroyed prompts' - // buffers to the pool, which limit_size does not count, rather than to the allocator + // 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); limit_size = std::max(1, 0.4*size()); @@ -1777,6 +1778,8 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro update(); + common_state_buffer_pool::instance().trim(0); + return nullptr; } From 1db6cb3588734769c5e083c85bd32e02237c5b92 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 19:10:48 -0700 Subject: [PATCH 5/6] server: release the checkpoint pool when entering the idle sleep state --- tools/server/server-queue.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 16f11f7b8f49..a4d80dde335e 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -333,6 +333,11 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { cb(true); } req_stop_sleeping = false; + // Sleep is when the server claims the memory is gone, and it preempts the timed + // trim below whenever the sleep threshold is inside that idle window, so release + // unconditionally here. Held under the lock deliberately: dropping it around the + // unmap would open a wakeup window between the callbacks and the wait. + common_state_buffer_pool::instance().trim(0); // wait until we are requested to exit sleeping state condition_tasks.wait(lock, [&]{ return (!running || req_stop_sleeping); From 00ce29bd731937d2f5ef1c648dd9b042d15c17dd Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 9 Sep 2026 06:35:15 -0700 Subject: [PATCH 6/6] server: do not let the checkpoint pool outlive-order a checkpoint destructor 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&&) 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). --- common/common.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 210e88b616ca..69e358d914a0 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2252,8 +2252,15 @@ bool common_prompt_batch_decode( } common_state_buffer_pool & common_state_buffer_pool::instance() { - static common_state_buffer_pool pool; - return pool; + // Deliberately leaked. A function-local static is destroyed in reverse order of completion of + // construction, so a common_prompt_checkpoint with static storage duration constructed before + // this one would be destroyed after it and its destructor would call put() on a destroyed + // object. [basic.start.term]/5 makes that undefined, and the try/catch in the destructor + // cannot help, because it is not an exception. No such holder exists today; this keeps it from + // becoming a silent use-after-free the day one does. The pool is process-wide and one + // allocation, so leaking it at exit costs nothing. + static common_state_buffer_pool * pool = new common_state_buffer_pool(); + return *pool; } static size_t common_state_buffer_pool_cap() {