server: reuse checkpoint state buffers from a bounded pool

common/common.cpp

@@ -2251,6 +2251,171 @@ bool common_prompt_batch_decode(
     return true;
 }
 
+common_state_buffer_pool & common_state_buffer_pool::instance() {
+    // 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() {
+    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);
+    }
+
+    GGML_UNUSED(mem_free);
+
+    // 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;
+    }
+
+    // of total, not free: every non-Windows host reports free == total for the CPU device
+    return mem_total / 16;
+}
+
+void common_state_buffer_pool::get(std::vector<uint8_t> & dst, size_t size) {
+    std::vector<uint8_t> prev; // the caller's previous storage, offered back below
+
+    {
+        std::lock_guard<std::mutex> lock(mtx);
+
+        st.n_get++;
+
+        t_last_us = ggml_time_us();
+
+        if (dst.capacity() >= size) {
+            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) {
+                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();
+
+                prev = std::move(free_bufs[best]);
+                free_bufs.erase(free_bufs.begin() + best);
+
+                std::swap(dst, prev);
+
+                st.n_hit++;
+            }
+        }
+    }
+
+    // too small here, but may still serve a smaller checkpoint later
+    put(std::move(prev));
+
+    // usually a no-op leaving stale bytes: callers overwrite the whole buffer and abort if the
+    // state does not fill it
+    dst.resize(size);
+}
+
+void common_state_buffer_pool::put(std::vector<uint8_t> && src) {
+    const size_t cap = src.capacity();
+
+    if (cap < MIN_BUFFER_BYTES) {
+        return; // below the allocator's mmap threshold: holding it would save nothing
+    }
+
+    std::lock_guard<std::mutex> 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 (cap > cap_bytes) {
+        return; // one buffer of this size would spend the whole budget
+    }
+
+    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++;
+    }
+
+    // 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;
+
+    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<std::mutex> 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;
+    n_hwm      = 0;
+}
+
+common_state_buffer_pool::stats common_state_buffer_pool::get_stats() const {
+    std::lock_guard<std::mutex> 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() {
+    // 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 (...) {
+    }
+}
+
 size_t common_prompt_checkpoint::size() const {
     return data_tgt.size() + data_dft.size() + data_spec.size();
 }
@@ -2289,7 +2454,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 +2472,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) {

common/common.h

@@ -14,6 +14,7 @@
 #include <string_view>
 #include <vector>
 #include <map>
+#include <mutex>
 #include <algorithm>
 #include <fstream>
 
@@ -1134,6 +1135,58 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *);
 // prompt utils
 //
 
+// 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.
+//
+// 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.
+//
+// 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 {
+    static constexpr size_t MAX_BUFFERS = 64;
+
+    // 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 {
+        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_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
+    };
+
+    // contents of `dst` are NOT preserved; every caller overwrites the whole buffer
+    void get(std::vector<uint8_t> & dst, size_t size);
+
+    void put(std::vector<uint8_t> && src);
+
+    void trim(int64_t idle_us);
+
+    stats get_stats() const;
+
+    static common_state_buffer_pool & instance();
+
+private:
+    mutable std::mutex mtx;
+
+    std::vector<std::vector<uint8_t>> 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 +1203,18 @@ struct common_prompt_checkpoint {
     // (e.g. eagle3's deferred-boundary g_embd row)
     std::vector<uint8_t> data_spec;
 
+    common_prompt_checkpoint() = default;
+
+    // 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;
+    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;

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, %" 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);
+                cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024,
+                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

tools/server/server-queue.cpp

@@ -1,3 +1,4 @@
+#include "common.h"
 #include "server-task.h"
 #include "server-queue.h"
 
@@ -332,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);
@@ -357,6 +363,12 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
                 if (res) {
                     break; // new task arrived or terminate
                 }
+
+                // 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);
+
                 // otherwise, loop again to check sleeping condition
             }
         }

tools/server/server-task.cpp

@@ -1767,12 +1767,19 @@ 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());
 
+        // 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<size_t>(1, 0.4*size());
 
         SRV_WRN(" - cache size limit reduced to %.3f MiB\n", limit_size / (1024.0 * 1024.0));
 
         update();
 
+        common_state_buffer_pool::instance().trim(0);
+
         return nullptr;
     }