server: on-demand mmproj - free encoder VRAM on the text path
The vision encoder sits idle in VRAM on every text request. With LLAMA_MMPROJ_ONDEMAND=1 the server keeps the encoder in RAM (released after load) and brings it into VRAM just-in-time before mtmd_batch_encode, releasing it again afterwards - so the encoder's VRAM is free for KV / expert cache on the common text-only path. When the encoder does not fit (that VRAM has been claimed), evict the LLM backbone weights first: they are not needed while the encoder runs (the decode that uses them happens afterwards), their host shadow is read-only (cheap free, no D2H), and the KV / prompt cache is left untouched. Order on the way back matters: free the encoder before re-uploading the weights so the peak stays within VRAM. Also wire the encoder release/restore into the VRAM arbiter: vram_go_cold releases it; the just-in-time encode-path restore replaces the previous restore in vram_ensure_warm. Assisted-by: Claude
This commit is contained in:
@@ -780,7 +780,36 @@ struct server_slot {
|
||||
// TODO @ngxson : move this log line to debug when it become more stable
|
||||
SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
|
||||
|
||||
// Bring the vision/audio encoder into VRAM just for this encode. In on-demand mode the
|
||||
// encoder is normally kept in RAM (its VRAM freed for KV / expert cache). If it does not
|
||||
// fit, evict the LLM backbone weights first: they are NOT needed while the encoder runs (the
|
||||
// decode that uses them happens afterwards) and their host shadow is read-only, so this is a
|
||||
// cheap free (no D2H) and leaves the KV cache / prompt cache untouched.
|
||||
static const bool mmproj_ondemand = getenv("LLAMA_MMPROJ_ONDEMAND") != nullptr;
|
||||
bool weights_evicted = false;
|
||||
if (mctx && !mtmd_restore_device(mctx)) {
|
||||
if (mmproj_ondemand) {
|
||||
llama_context_release_device(ctx_tgt, /* evict_kv = */ false); // free backbone, keep KV
|
||||
weights_evicted = true;
|
||||
}
|
||||
if (!mtmd_restore_device(mctx)) {
|
||||
if (weights_evicted) { llama_context_restore_device(ctx_tgt); }
|
||||
SLT_ERR(*this, "%s", "failed to bring the multimodal encoder into VRAM for encoding\n");
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
res = mtmd_batch_encode(mbatch.get());
|
||||
|
||||
// release the encoder VRAM again (on-demand), then restore the backbone weights for decode.
|
||||
// Order matters: free the encoder BEFORE re-uploading the weights so the peak stays within VRAM.
|
||||
if (mctx && mmproj_ondemand) {
|
||||
mtmd_release_device(mctx);
|
||||
}
|
||||
if (weights_evicted) {
|
||||
llama_context_restore_device(ctx_tgt);
|
||||
}
|
||||
|
||||
if (res != 0) {
|
||||
SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
|
||||
return -1;
|
||||
@@ -1058,6 +1087,11 @@ private:
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_restore_device(ctx_dft);
|
||||
}
|
||||
// NOTE: the multimodal (vision/audio) encoder is intentionally NOT restored here. It is
|
||||
// brought into VRAM just-in-time before an image/audio encode (process_mtmd_chunk) and, in
|
||||
// on-demand mode, released again right after - so a warm text-only model holds no encoder
|
||||
// VRAM (that space is free for KV / expert cache). A cold->warm wake for a *text* request
|
||||
// therefore leaves the encoder in RAM; an image request restores it at encode time.
|
||||
vram_cold = false;
|
||||
}
|
||||
|
||||
@@ -1082,6 +1116,11 @@ private:
|
||||
if (ctx_dft != nullptr) {
|
||||
llama_context_release_device(ctx_dft, vram_evict_kv);
|
||||
}
|
||||
// also release the multimodal (vision/audio) encoder weights (dead weight while cold);
|
||||
// its read-only host shadow is captured once and rebuilt on wake by vram_ensure_warm()
|
||||
if (mctx != nullptr) {
|
||||
mtmd_release_device(mctx);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
if (vram_flock) {
|
||||
flock(vram_lock_fd, LOCK_UN);
|
||||
@@ -1436,6 +1475,13 @@ private:
|
||||
}
|
||||
SRV_INF("loaded multimodal model, '%s'\n", mmproj_path.c_str());
|
||||
|
||||
// on-demand encoder: keep the vision/audio encoder weights in RAM (out of VRAM) until an
|
||||
// image/audio actually needs encoding, freeing that VRAM for KV / expert cache on the
|
||||
// common text-only path. process_mtmd_chunk() brings the encoder in just-in-time.
|
||||
if (getenv("LLAMA_MMPROJ_ONDEMAND") != nullptr) {
|
||||
mtmd_release_device(mctx);
|
||||
}
|
||||
|
||||
if (params_base.ctx_shift) {
|
||||
params_base.ctx_shift = false;
|
||||
SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled");
|
||||
|
||||
Reference in New Issue
Block a user