From 5ddcf40fc1d3bfddffeaf65cc526307c2e944b9b Mon Sep 17 00:00:00 2001 From: Lumpiasty Date: Sun, 26 Jul 2026 14:55:39 +0200 Subject: [PATCH] llama: evict recurrent/SSM state on device release The recurrent (SSM/conv) state of hybrid models (e.g. Qwen3.5) was left resident when a model's device buffers were released for on-demand VRAM sharing - llama_memory_recurrent::release_device_buffers() was a no-op default. Implement it (and restore_device_buffers) with the same capture-host-shadow / free / reallocate pattern as llama_kv_cache, so llama_memory_hybrid now evicts both its attention KV and its recurrent state. The state is read-write, so its shadow is recaptured on every release. Assisted-by: Claude --- src/llama-memory-hybrid.cpp | 2 +- src/llama-memory-recurrent.cpp | 82 +++++++++++++++++++++++++++++++++- src/llama-memory-recurrent.h | 14 ++++++ 3 files changed, 95 insertions(+), 3 deletions(-) diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index ebd36b1a1..5c405dbd3 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -202,7 +202,7 @@ void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, } void llama_memory_hybrid::release_device_buffers() { - // evict the attention KV (grows with context); the recurrent state uses the no-op default + // evict both the attention KV (grows with context) and the recurrent/SSM state mem_attn->release_device_buffers(); mem_recr->release_device_buffers(); } diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index ef82eb976..9afba4c3f 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -140,7 +140,9 @@ void llama_memory_recurrent::clear(bool data) { if (data) { for (auto & [_, buf] : ctxs_bufs) { - ggml_backend_buffer_clear(buf.get(), 0); + if (buf) { // may be null if evicted for on-demand VRAM sharing + ggml_backend_buffer_clear(buf.get(), 0); + } } } @@ -399,6 +401,7 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) { std::map llama_memory_recurrent::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { + if (!buf) { continue; } // may be null if evicted for on-demand VRAM sharing ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get()); } return ret; @@ -700,12 +703,87 @@ bool llama_memory_recurrent::get_can_shift() const { size_t llama_memory_recurrent::total_size() const { size_t size = 0; for (const auto & [_, buf] : ctxs_bufs) { - size += ggml_backend_buffer_get_size(buf.get()); + if (buf) { // may be null if evicted for on-demand VRAM sharing + size += ggml_backend_buffer_get_size(buf.get()); + } } return size; } +void llama_memory_recurrent::release_device_buffers() { + // Same mechanism as llama_kv_cache: the recurrent (SSM/conv) state is read-write, so its host + // shadow is (re)captured on every release. The caller must have synchronized the backend. + 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; + } + auto & sh = dev_shadows[i]; + sh.releasable = true; + sh.buft = ggml_backend_buffer_get_type(buf); + + 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(); + 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 recurrent state from device\n", __func__, freed / 1024.0 / 1024.0); + } +} + +bool llama_memory_recurrent::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 recurrent 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(); + return true; +} + size_t llama_memory_recurrent::size_r_bytes() const { size_t size_r_bytes = 0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748..86af82ac1 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -66,6 +66,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; + uint32_t head = 0; // the location where the batch will be placed in the cache (see find_slot()) uint32_t size = 0; // total number of cells, shared across all sequences uint32_t used = 0; // used cells (i.e. at least one seq_id) @@ -121,6 +125,16 @@ private: // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; + // on-demand device eviction (see release_device_buffers): host shadow of each device buffer's + // live contents (recaptured on every release since the recurrent state is read-write) + struct device_buffer_shadow { + ggml_backend_buffer_type_t buft = nullptr; + bool releasable = false; + std::vector data; + }; + std::vector dev_shadows; // parallel to ctxs_bufs + bool dev_released = false; + size_t total_size() const; size_t size_r_bytes() const;