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