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
This commit is contained in:
2026-07-26 21:58:43 +02:00
parent fd7cbd5f4f
commit 5ddcf40fc1
3 changed files with 95 additions and 3 deletions
+1 -1
View File
@@ -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();
}
+80 -2
View File
@@ -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<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> 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;
+14
View File
@@ -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<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> 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<uint8_t> data;
};
std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
bool dev_released = false;
size_t total_size() const;
size_t size_r_bytes() const;