From 7398c40eda5e004c13fa0c37af7337279bb7fd07 Mon Sep 17 00:00:00 2001 From: Lumpiasty Date: Sat, 25 Jul 2026 02:02:18 +0200 Subject: [PATCH] server: optionally evict the KV cache too (Phase 2 of VRAM sharing) Extend on-demand device residency to the KV cache so that when a model's KV plus another model would not fit in VRAM, the KV can also be evicted to a host shadow (D2H on release, H2D on restore) instead of only the weights. - llama_memory_i: add release_device_buffers()/restore_device_buffers() (default no-op). Implemented in llama_kv_cache (D2H shadow of the live ctxs_bufs, freed and reallocated like the weights); llama_memory_hybrid and llama_kv_cache_iswa delegate to their child caches. - llama_context::release_device(evict_kv): also evict the memory's device buffers when requested; restore_device() rebuilds them. Public API llama_context_release_device gains an evict_kv flag. - server: LLAMA_SLEEP_EVICT_KV=1 enables it. Off by default (weights-only), since the KV shadow adds a D2H/H2D copy of the live cache each cycle. Validated on RX 580 (Vulkan), 4B @ 32k ctx: weights-only cold VRAM 1750 MB (KV stays); weights+KV cold VRAM 726 MB (KV freed, ~1 GB reclaimed). KV survives the round-trip: prompt cache reused after the cycle (prompt_n 4 vs 42), correct output. Assisted-by: Claude --- include/llama.h | 6 ++- src/llama-context.cpp | 20 ++++++--- src/llama-context.h | 9 +++- src/llama-kv-cache-iswa.cpp | 11 +++++ src/llama-kv-cache-iswa.h | 3 ++ src/llama-kv-cache.cpp | 74 +++++++++++++++++++++++++++++++++ src/llama-kv-cache.h | 14 +++++++ src/llama-memory-hybrid.cpp | 12 ++++++ src/llama-memory-hybrid.h | 3 ++ src/llama-memory.h | 10 +++++ tools/server/server-context.cpp | 6 ++- 11 files changed, 156 insertions(+), 12 deletions(-) diff --git a/include/llama.h b/include/llama.h index dd933a232..75abc9f5a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -560,9 +560,11 @@ extern "C" { // On-demand device (VRAM) residency: free the model's GPU weight buffers (keeping a host // shadow so no reload is needed) to hand VRAM to another model, and rebuild them on demand. - // The KV cache is left resident. Any subsequent llama_decode auto-restores the weights. + // Any subsequent llama_decode auto-restores. If evict_kv is true, the KV cache device buffers + // are also evicted to a host shadow (for when the KV would not fit alongside the other model), + // at the cost of a D2H/H2D copy of the live KV each cycle; otherwise the KV stays resident. // Intended for time-sharing a single GPU between several always-loaded models. - LLAMA_API void llama_context_release_device(struct llama_context * ctx); + LLAMA_API void llama_context_release_device(struct llama_context * ctx, bool evict_kv); LLAMA_API void llama_context_restore_device(struct llama_context * ctx); LLAMA_API bool llama_context_weights_resident(const struct llama_context * ctx); LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4abb55650..950ce2847 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -728,21 +728,29 @@ void llama_context::synchronize() { t_compute_start_us = 0; } -void llama_context::release_device() { - if (!model.weights_resident()) { +void llama_context::release_device(bool evict_kv) { + if (!model.weights_resident() && !(evict_kv && !kv_device_evicted)) { return; } // ensure no compute is in flight before freeing the device buffers synchronize(); model.release_device_weights(); + if (evict_kv && memory && !kv_device_evicted) { + memory->release_device_buffers(); + kv_device_evicted = true; + } } void llama_context::restore_device() { - if (model.weights_resident()) { + if (model.weights_resident() && !kv_device_evicted) { return; } model.restore_device_weights(); - // make sure all weight uploads have completed before any compute reads them + if (kv_device_evicted && memory) { + memory->restore_device_buffers(); + kv_device_evicted = false; + } + // make sure all weight/KV uploads have completed before any compute reads them if (sched) { ggml_backend_sched_synchronize(sched.get()); } @@ -3700,8 +3708,8 @@ void llama_synchronize(llama_context * ctx) { ctx->synchronize(); } -void llama_context_release_device(llama_context * ctx) { - ctx->release_device(); +void llama_context_release_device(llama_context * ctx, bool evict_kv) { + ctx->release_device(evict_kv); } void llama_context_restore_device(llama_context * ctx) { diff --git a/src/llama-context.h b/src/llama-context.h index e50fbd515..b3893c9d8 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -59,8 +59,10 @@ struct llama_context { // on-demand device (VRAM) residency: free / rebuild the model's GPU weight buffers so that // VRAM can be time-shared with another model. release_device() synchronizes first; decode() - // auto-restores when it finds the weights have been released. - void release_device(); + // auto-restores when it finds the weights have been released. If evict_kv is set, the KV cache + // device buffers are also evicted to a host shadow (for when the KV would not fit alongside the + // other model) - this adds a D2H/H2D copy of the live KV on each cycle. + void release_device(bool evict_kv = false); void restore_device(); const llama_model & get_model() const; @@ -351,6 +353,9 @@ private: bool sched_need_reserve = true; + // on-demand device residency: true while the KV cache device buffers have been evicted to host + bool kv_device_evicted = false; + ggml_backend_t backend_cpu = nullptr; std::vector backends; diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index e91866469..56fdfe5c1 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -272,6 +272,17 @@ void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, kv_swa->state_read(io, seq_id, flags); } +void llama_kv_cache_iswa::release_device_buffers() { + kv_base->release_device_buffers(); + kv_swa->release_device_buffers(); +} + +bool llama_kv_cache_iswa::restore_device_buffers() { + bool ok = kv_base->restore_device_buffers(); + ok = kv_swa->restore_device_buffers() && ok; + return ok; +} + llama_kv_cache * llama_kv_cache_iswa::get_base() const { return kv_base.get(); } diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 7dab6eaa8..653bcc4e5 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -83,6 +83,9 @@ public: void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + void release_device_buffers() override; + bool restore_device_buffers() override; + // // llama_kv_cache_iswa specific API // diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index e70583e64..8ca118366 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1807,6 +1807,80 @@ size_t llama_kv_cache::total_size() const { return size; } +void llama_kv_cache::release_device_buffers() { + // NOTE: the caller must have synchronized the backend so nothing references these buffers. + // The KV cache is read-write, so its host shadow is (re)captured fresh on every release. + if (dev_released) { + return; + } + dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{}); + size_t freed = 0; + for (size_t i = 0; i < ctxs_bufs.size(); ++i) { + ggml_context * ctx = ctxs_bufs[i].first.get(); + ggml_backend_buffer_t buf = ctxs_bufs[i].second.get(); + if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) { + continue; // only real device (VRAM) buffers are evictable + } + auto & sh = dev_shadows[i]; + sh.releasable = true; + sh.buft = ggml_backend_buffer_get_type(buf); + + // capture live contents compactly in stable iteration order (skip views, which alias a base) + size_t total = 0; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src == nullptr) { total += ggml_nbytes(t); } + } + sh.data.resize(total); + size_t off = 0; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src != nullptr) { continue; } + const size_t n = ggml_nbytes(t); + ggml_backend_tensor_get(t, sh.data.data() + off, 0, n); + off += n; + } + + freed += ggml_backend_buffer_get_size(buf); + ctxs_bufs[i].second.reset(); // free the device buffer + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = nullptr; + t->data = nullptr; + } + } + dev_released = true; + if (freed > 0) { + LLAMA_LOG_INFO("%s: released %.2f MiB of KV cache from device\n", __func__, freed / 1024.0 / 1024.0); + } +} + +bool llama_kv_cache::restore_device_buffers() { + if (!dev_released) { + return true; + } + for (size_t i = 0; i < ctxs_bufs.size(); ++i) { + auto & sh = dev_shadows[i]; + if (!sh.releasable) { + continue; + } + ggml_context * ctx = ctxs_bufs[i].first.get(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft); + if (buf == nullptr) { + LLAMA_LOG_ERROR("%s: failed to reallocate KV cache device buffer (out of VRAM?)\n", __func__); + return false; + } + size_t off = 0; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src != nullptr) { continue; } + const size_t n = ggml_nbytes(t); + ggml_backend_tensor_set(t, sh.data.data() + off, 0, n); + off += n; + } + ctxs_bufs[i].second.reset(buf); + } + dev_released = false; + dev_shadows.clear(); // recaptured on next release + return true; +} + size_t llama_kv_cache::size_k_bytes() const { size_t size_k_bytes = 0; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 531d99dbd..abacbe532 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -149,6 +149,10 @@ public: void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + // on-demand device (VRAM) residency (see llama_memory_i) + void release_device_buffers() override; + bool restore_device_buffers() override; + // // llama_kv_cache specific API // @@ -265,6 +269,16 @@ private: // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; + // on-demand device eviction: host shadow of each device buffer's live contents (re-captured on + // every release since the KV is read-write), plus the buffer type needed to reallocate it. + struct device_buffer_shadow { + ggml_backend_buffer_type_t buft = nullptr; + bool releasable = false; // device (VRAM) buffer that was freed + std::vector data; // captured tensor bytes (compact, iteration order) + }; + std::vector dev_shadows; // parallel to ctxs_bufs + bool dev_released = false; + // the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot()) // note: this is not part of the KV state and it's only used to speed-up the find_slot() method std::vector v_heads; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9..ebd36b1a1 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -201,6 +201,18 @@ void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, mem_recr->state_read(io, seq_id, flags); } +void llama_memory_hybrid::release_device_buffers() { + // evict the attention KV (grows with context); the recurrent state uses the no-op default + mem_attn->release_device_buffers(); + mem_recr->release_device_buffers(); +} + +bool llama_memory_hybrid::restore_device_buffers() { + bool ok = mem_attn->restore_device_buffers(); + ok = mem_recr->restore_device_buffers() && ok; + return ok; +} + llama_kv_cache * llama_memory_hybrid::get_mem_attn() const { return mem_attn.get(); } diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74..318076def 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -76,6 +76,9 @@ public: void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + void release_device_buffers() override; + bool restore_device_buffers() override; + // // llama_memory_hybrid specific API // diff --git a/src/llama-memory.h b/src/llama-memory.h index db8253966..e6f3cc0ba 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -124,6 +124,16 @@ struct llama_memory_i { virtual void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const = 0; virtual void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) = 0; + + // + // on-demand device (VRAM) residency + // + + // Free this memory's device (VRAM) buffers to a host shadow and rebuild them on demand, to + // hand VRAM to another model while keeping the cached data (no re-prefill). The caller must + // have synchronized the backend first. Default: no-op (the memory stays resident). + virtual void release_device_buffers() {} + virtual bool restore_device_buffers() { return true; } }; using llama_memory_ptr = std::unique_ptr; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 468f34bd3..deebbfa6a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -914,6 +914,7 @@ private: // weights) right before it decodes, and goes cold (releases weights + unlocks) when it idles, // so only one model holds VRAM at a time and the KV cache is never evicted (no re-prefill). bool vram_only = false; // release-mode sleep active (LLAMA_SLEEP_VRAM_ONLY) + bool vram_evict_kv = false; // also evict the KV cache to host on release (LLAMA_SLEEP_EVICT_KV) bool vram_flock = false; // cross-process flock coordination available std::atomic vram_cold{true}; // weights currently released (read by warden thread) int vram_lock_fd = -1; // fd for /token.lock @@ -931,6 +932,7 @@ private: return; } vram_only = true; + vram_evict_kv = getenv("LLAMA_SLEEP_EVICT_KV") != nullptr; vram_cold = false; // weights are resident right after load #if !defined(_WIN32) const char * arena_env = getenv("LLAMA_VRAM_ARENA"); @@ -1040,9 +1042,9 @@ private: if (!vram_only || vram_cold) { return; } - llama_context_release_device(ctx_tgt); + llama_context_release_device(ctx_tgt, vram_evict_kv); if (ctx_dft != nullptr) { - llama_context_release_device(ctx_dft); + llama_context_release_device(ctx_dft, vram_evict_kv); } #if !defined(_WIN32) if (vram_flock) {