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
This commit is contained in:
@@ -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;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user