server: on-demand VRAM sharing to time-share one GPU between models
Add release/restore of a model's GPU weight buffers (keeping a host shadow and the KV cache) so several always-loaded llama-server processes can time-share a single GPU without reloading or losing the prompt cache. - llama-model: release_device_weights()/restore_device_weights() capture a compact host shadow (stable iteration order, view-skipping) and free then realloc the device weight buffers; weights_resident() query. - llama-context: release_device()/restore_device() wrappers; decode() auto- restores; public C API llama_context_release_device/restore_device. - server: LLAMA_SLEEP_VRAM_ONLY makes idle-sleep release only the VRAM weights (not a full unload/reload). A cross-process flock token in LLAMA_VRAM_ARENA enforces "resident iff holds token"; an inotify doorbell forces the holder to release on contention. The warden thread only touches the task queue, so releases run on the loop thread and never race a decode. Validated on RX 580 (Vulkan): two models share 8GB, never both resident, correct output under contention, KV cache preserved (no re-prefill). Assisted-by: Claude
This commit is contained in:
@@ -25,6 +25,8 @@
|
||||
#include <filesystem>
|
||||
#include <utility>
|
||||
#include <fstream>
|
||||
#include <thread>
|
||||
#include <atomic>
|
||||
|
||||
// fix problem with std::min and std::max
|
||||
#if defined(_WIN32)
|
||||
@@ -35,6 +37,16 @@
|
||||
#include <windows.h>
|
||||
#endif
|
||||
|
||||
// POSIX file locking + inotify doorbell for the cross-process VRAM arbiter (see vram_share_* below)
|
||||
#if !defined(_WIN32)
|
||||
#include <sys/file.h>
|
||||
#include <sys/stat.h>
|
||||
#include <sys/inotify.h>
|
||||
#include <fcntl.h>
|
||||
#include <unistd.h>
|
||||
#include <poll.h>
|
||||
#endif
|
||||
|
||||
using json = nlohmann::ordered_json;
|
||||
|
||||
constexpr int HTTP_POLLING_SECONDS = 1;
|
||||
@@ -871,6 +883,8 @@ public:
|
||||
}
|
||||
|
||||
~server_context_impl() {
|
||||
// stop the VRAM warden thread (and release the token) before tearing anything down
|
||||
vram_share_shutdown();
|
||||
if (!sleeping) {
|
||||
// destroy() is already called when entering sleeping state
|
||||
// we don't call it again here to avoid double free
|
||||
@@ -894,6 +908,150 @@ private:
|
||||
llama_model * model_dft = nullptr;
|
||||
llama_context * ctx_dft = nullptr;
|
||||
|
||||
// Cross-process VRAM arbiter: when LLAMA_SLEEP_VRAM_ONLY is set, several always-loaded
|
||||
// llama-server processes time-share one GPU. A single flock() on <arena>/token.lock is the
|
||||
// baton: "resident (weights in VRAM) iff I hold the token". A model warms up (locks + restores
|
||||
// weights) right before it decodes, and goes cold (releases weights + unlocks) when it idles,
|
||||
// so only one model holds VRAM at a time and the KV cache is never evicted (no re-prefill).
|
||||
bool vram_only = false; // release-mode sleep active (LLAMA_SLEEP_VRAM_ONLY)
|
||||
bool vram_flock = false; // cross-process flock coordination available
|
||||
std::atomic<bool> vram_cold{true}; // weights currently released (read by warden thread)
|
||||
int vram_lock_fd = -1; // fd for <arena>/token.lock
|
||||
|
||||
// inotify "doorbell": a waiter that wants the VRAM token touches <arena>/doorbell/<pid>, which
|
||||
// wakes the current holder's warden thread so it releases immediately instead of only on idle.
|
||||
std::string vram_doorbell_dir;
|
||||
std::string vram_pid_str;
|
||||
int vram_inotify_fd = -1;
|
||||
std::thread vram_warden;
|
||||
std::atomic<bool> vram_warden_run{false};
|
||||
|
||||
void vram_share_init() {
|
||||
if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) {
|
||||
return;
|
||||
}
|
||||
vram_only = true;
|
||||
vram_cold = false; // weights are resident right after load
|
||||
#if !defined(_WIN32)
|
||||
const char * arena_env = getenv("LLAMA_VRAM_ARENA");
|
||||
const std::string arena = arena_env ? arena_env : "/dev/shm/llama-vram";
|
||||
mkdir(arena.c_str(), 0777);
|
||||
const std::string lock_path = arena + "/token.lock";
|
||||
vram_lock_fd = open(lock_path.c_str(), O_RDWR | O_CREAT, 0666);
|
||||
if (vram_lock_fd >= 0) {
|
||||
vram_flock = true;
|
||||
vram_pid_str = std::to_string(getpid());
|
||||
vram_doorbell_dir = arena + "/doorbell";
|
||||
mkdir(vram_doorbell_dir.c_str(), 0777);
|
||||
vram_inotify_fd = inotify_init1(IN_NONBLOCK);
|
||||
if (vram_inotify_fd >= 0) {
|
||||
inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE);
|
||||
vram_warden_run = true;
|
||||
vram_warden = std::thread([this]{ vram_warden_loop(); });
|
||||
}
|
||||
SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (doorbell %s)\n",
|
||||
lock_path.c_str(), vram_doorbell_dir.c_str());
|
||||
} else {
|
||||
SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str());
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// warden thread: block on the doorbell; when another process rings (wants the token) and we
|
||||
// currently hold it, ask the loop to yield (release) at its next idle point. Only touches the
|
||||
// thread-safe queue - never the GPU - so it cannot race a decode.
|
||||
void vram_warden_loop() {
|
||||
#if !defined(_WIN32)
|
||||
char buf[4096];
|
||||
while (vram_warden_run.load()) {
|
||||
struct pollfd pfd { vram_inotify_fd, POLLIN, 0 };
|
||||
int pr = poll(&pfd, 1, 500); // 500ms so we periodically re-check the run flag
|
||||
if (pr <= 0) {
|
||||
continue;
|
||||
}
|
||||
ssize_t n = read(vram_inotify_fd, buf, sizeof(buf));
|
||||
if (n <= 0) {
|
||||
continue;
|
||||
}
|
||||
bool foreign_ring = false;
|
||||
for (char * p = buf; p < buf + n; ) {
|
||||
struct inotify_event * ev = (struct inotify_event *) p;
|
||||
if (ev->len > 0 && vram_pid_str != ev->name) {
|
||||
foreign_ring = true; // someone else wants the GPU
|
||||
}
|
||||
p += sizeof(struct inotify_event) + ev->len;
|
||||
}
|
||||
if (foreign_ring && !vram_cold) {
|
||||
queue_tasks.request_yield();
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ring the doorbell so the current token holder releases promptly
|
||||
void vram_ring_doorbell() {
|
||||
#if !defined(_WIN32)
|
||||
if (!vram_flock || vram_doorbell_dir.empty()) {
|
||||
return;
|
||||
}
|
||||
const std::string f = vram_doorbell_dir + "/" + vram_pid_str;
|
||||
int fd = open(f.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
|
||||
if (fd >= 0) {
|
||||
ssize_t w = write(fd, "1", 1);
|
||||
(void) w;
|
||||
close(fd);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// acquire the VRAM token (blocking) and bring weights back to the device. Called on the loop
|
||||
// thread right before a decode, so it never races compute.
|
||||
void vram_ensure_warm() {
|
||||
if (!vram_only || !vram_cold) {
|
||||
return;
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
vram_ring_doorbell(); // nudge the current holder to release
|
||||
flock(vram_lock_fd, LOCK_EX); // blocks until the current holder goes cold
|
||||
}
|
||||
#endif
|
||||
llama_context_restore_device(ctx_tgt);
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_restore_device(ctx_dft);
|
||||
}
|
||||
vram_cold = false;
|
||||
}
|
||||
|
||||
void vram_share_shutdown() {
|
||||
#if !defined(_WIN32)
|
||||
vram_warden_run = false;
|
||||
if (vram_warden.joinable()) {
|
||||
vram_warden.join();
|
||||
}
|
||||
if (vram_inotify_fd >= 0) { close(vram_inotify_fd); vram_inotify_fd = -1; }
|
||||
if (vram_lock_fd >= 0) { flock(vram_lock_fd, LOCK_UN); close(vram_lock_fd); vram_lock_fd = -1; }
|
||||
#endif
|
||||
}
|
||||
|
||||
// release the device weights (keeping KV + host shadow) and drop the VRAM token so another
|
||||
// model can warm up. Called on the loop thread when the server goes idle.
|
||||
void vram_go_cold() {
|
||||
if (!vram_only || vram_cold) {
|
||||
return;
|
||||
}
|
||||
llama_context_release_device(ctx_tgt);
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_release_device(ctx_dft);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
flock(vram_lock_fd, LOCK_UN);
|
||||
}
|
||||
#endif
|
||||
vram_cold = true;
|
||||
}
|
||||
|
||||
common_speculative_init_result_ptr spec_init;
|
||||
|
||||
common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
|
||||
@@ -951,6 +1109,25 @@ private:
|
||||
|
||||
void handle_sleeping_state(bool new_state) {
|
||||
GGML_ASSERT(sleeping != new_state);
|
||||
|
||||
// Lightweight VRAM-only sleep: instead of a full unload/reload (which frees RAM and the
|
||||
// KV cache and pays a full reload on wake), just release the model's device (VRAM) weight
|
||||
// buffers to a host shadow (keeping the context, KV cache and host weights) and drop the
|
||||
// shared VRAM token. Waking is handled lazily by vram_ensure_warm() right before the next
|
||||
// decode (which re-acquires the token first), so a single GPU is time-shared between models
|
||||
// with no re-prefill.
|
||||
if (vram_only && ctx_tgt != nullptr) {
|
||||
if (new_state) {
|
||||
SRV_INF("%s", "server entering sleeping state (VRAM-only: releasing device weights, keeping KV cache)\n");
|
||||
vram_go_cold();
|
||||
} else {
|
||||
SRV_INF("%s", "server exiting sleeping state (VRAM-only: weights restored on next decode)\n");
|
||||
// token is acquired and weights restored by vram_ensure_warm() before decode
|
||||
}
|
||||
sleeping = new_state;
|
||||
return;
|
||||
}
|
||||
|
||||
if (new_state) {
|
||||
SRV_INF("%s", "server is entering sleeping state\n");
|
||||
destroy();
|
||||
@@ -1409,6 +1586,12 @@ private:
|
||||
handle_sleeping_state(sleeping);
|
||||
});
|
||||
|
||||
// VRAM arbiter: enable cross-process GPU time-sharing (if LLAMA_SLEEP_VRAM_ONLY is set) and
|
||||
// start cold - release the just-loaded device weights and drop the shared token, so we only
|
||||
// occupy VRAM while actually serving. The first decode re-acquires the token and restores.
|
||||
vram_share_init();
|
||||
vram_go_cold();
|
||||
|
||||
metrics.init();
|
||||
|
||||
if (params_base.cache_idle_slots) {
|
||||
@@ -3589,6 +3772,10 @@ private:
|
||||
n_empty_consecutive = 0;
|
||||
}
|
||||
|
||||
// VRAM arbiter: make sure our weights are resident (acquiring the shared GPU token first)
|
||||
// before we decode. Runs on the loop thread, so it never races an in-flight decode.
|
||||
vram_ensure_warm();
|
||||
|
||||
const int ret = llama_decode(ctx_tgt, batch_view);
|
||||
|
||||
metrics.on_decoded(slots);
|
||||
|
||||
Reference in New Issue
Block a user