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:
@@ -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() {
|
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_attn->release_device_buffers();
|
||||||
mem_recr->release_device_buffers();
|
mem_recr->release_device_buffers();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -140,7 +140,9 @@ void llama_memory_recurrent::clear(bool data) {
|
|||||||
|
|
||||||
if (data) {
|
if (data) {
|
||||||
for (auto & [_, buf] : ctxs_bufs) {
|
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> llama_memory_recurrent::memory_breakdown() const {
|
||||||
std::map<ggml_backend_buffer_type_t, size_t> ret;
|
std::map<ggml_backend_buffer_type_t, size_t> ret;
|
||||||
for (const auto & [_, buf] : ctxs_bufs) {
|
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());
|
ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
|
||||||
}
|
}
|
||||||
return ret;
|
return ret;
|
||||||
@@ -700,12 +703,87 @@ bool llama_memory_recurrent::get_can_shift() const {
|
|||||||
size_t llama_memory_recurrent::total_size() const {
|
size_t llama_memory_recurrent::total_size() const {
|
||||||
size_t size = 0;
|
size_t size = 0;
|
||||||
for (const auto & [_, buf] : ctxs_bufs) {
|
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;
|
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 llama_memory_recurrent::size_r_bytes() const {
|
||||||
size_t size_r_bytes = 0;
|
size_t size_r_bytes = 0;
|
||||||
|
|
||||||
|
|||||||
@@ -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_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;
|
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 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 size = 0; // total number of cells, shared across all sequences
|
||||||
uint32_t used = 0; // used cells (i.e. at least one seq_id)
|
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:
|
// 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;
|
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 total_size() const;
|
||||||
|
|
||||||
size_t size_r_bytes() const;
|
size_t size_r_bytes() const;
|
||||||
|
|||||||
Reference in New Issue
Block a user