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

@@ -568,9 +568,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

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());
     }
@@ -3701,8 +3709,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) {

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<ggml_backend_ptr> backends;
 

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();
 }

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
     //

src/llama-kv-cache.cpp

@@ -1959,6 +1959,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;
 

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
     //
@@ -272,6 +276,16 @@ private:
     // ggml contexts for the KV cache along with the allocated backend buffers:
     std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> 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<uint8_t> data;    // captured tensor bytes (compact, iteration order)
+    };
+    std::vector<device_buffer_shadow> 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<uint32_t> v_heads;

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();
 }

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
     //

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<llama_memory_i>;

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<bool> vram_cold{true}; // weights currently released (read by warden thread)
     int  vram_lock_fd = -1;      // fd for <arena>/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) {