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:
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user