diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d3..69e358d914a0 100644 --- a/common/common.cpp +++ b/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 & dst, size_t size) { + std::vector prev; // the caller's previous storage, offered back below + + { + std::lock_guard 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 && 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 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 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 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) { diff --git a/common/common.h b/common/common.h index de49dac9f63a..0561eaeb5bcb 100644 --- a/common/common.h +++ b/common/common.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -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 & dst, size_t size); + + void put(std::vector && src); + + 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 +1203,18 @@ struct common_prompt_checkpoint { // (e.g. eagle3's deferred-boundary g_embd row) std::vector 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; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..7f3870ac5b82 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, %" 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 diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..a4d80dde335e 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" @@ -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 } } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..85a437cc6039 100644 --- a/tools/server/server-task.cpp +++ b/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(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; }