diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 950ce2847..16e72cda8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -738,11 +738,15 @@ void llama_context::release_device(bool evict_kv) { if (evict_kv && memory && !kv_device_evicted) { memory->release_device_buffers(); kv_device_evicted = true; + // also free the scheduler and its worst-case compute buffer (hundreds of MiB) so a cold + // model holds essentially no VRAM. Rebuilt lazily by sched_reserve() on restore. + sched.reset(); + sched_need_reserve = true; } } void llama_context::restore_device() { - if (model.weights_resident() && !kv_device_evicted) { + if (model.weights_resident() && !kv_device_evicted && sched) { return; } model.restore_device_weights(); @@ -750,6 +754,10 @@ void llama_context::restore_device() { memory->restore_device_buffers(); kv_device_evicted = false; } + // rebuild the scheduler + compute buffer if they were freed on release (evict_kv mode) + if (!sched) { + sched_reserve(); + } // make sure all weight/KV uploads have completed before any compute reads them if (sched) { ggml_backend_sched_synchronize(sched.get());