[Fast Start] Charge daemon-held weights against `gpu_memory_utilization` (#57298)
Signed-off-by: liusy58 <mg21330037@smail.nju.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Xun Sun <UNIDY2002@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Xun Sun <UNIDY2002@outlook.com>
siyu committed
Sep 30, 2026 at 01:58 UTC
f5ed3d9c03ec634fbfe670baaf234569bbdecf96
7 files changed
+187
-14
tests/model_executor/model_loader/test_weight_cache.py
+15
@@ -354,3 +354,18 @@ def test_weight_cache_key_distinguishes_dp_ranks():
354
assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"]
355
assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"]
356
assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"]
357
+
358
+
359
+def test_ipc_loader_copy_mode_reports_no_external_weight_memory():
360
+ """Copy mode clones the weights into the engine, so nothing is external
361
+ (vllm_config is never touched)."""
362
+ from vllm.config import LoadConfig
363
+ from vllm.model_executor.model_loader.weight_cache.ipc_loader import (
364
+ IpcModelLoader,
365
+ )
366
+
367
+ copy_mode = LoadConfig(
368
+ load_format="ipc_cache", model_loader_extra_config={"mode": "copy"}
369
+ )
370
+ loader = IpcModelLoader(copy_mode)
371
+ assert loader.get_external_weight_memory(None) == 0 # type: ignore[arg-type]
tests/v1/worker/test_utils.py
+37
-1
@@ -13,7 +13,7 @@ import torch
13
14
import vllm.v1.attention.backends.mla.index_group as index_group_module
15
import vllm.v1.hisparse.runtime as hisparse_runtime_module
16
-from vllm.config import CUDAGraphMode
16
+from vllm.config import CacheConfig, CUDAGraphMode
17
from vllm.config.mamba import MambaBackendEnum, MambaConfig
18
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (
19
worker as hisparse_worker_module,
@@ -25,12 +25,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
25
_SlotMappingStaging,
26
)
27
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
28
+from vllm.utils.mem_constants import GiB_bytes
29
+from vllm.utils.mem_utils import MemorySnapshot
30
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
31
from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
32
from vllm.v1.worker.utils import (
33
bind_kv_cache,
34
bind_kv_cache_to_layers,
35
copy_kv_cache_blocks_inplace,
36
+ request_memory,
37
)
38
39
@@ -1390,3 +1393,36 @@ def test_bind_kv_cache_draft_model(default_vllm_config):
1393
assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"]
1394
assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"]
1395
assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"]
1396
+
1397
+
1398
+def _memory_snapshot(total_gib: int, free_gib: int) -> MemorySnapshot:
1399
+ return MemorySnapshot(
1400
+ free_memory=free_gib * GiB_bytes,
1401
+ total_memory=total_gib * GiB_bytes,
1402
+ device="cpu",
1403
+ auto_measure=False,
1404
+ )
1405
+
1406
+
1407
+def test_request_memory_charges_external_weights():
1408
+ """Externally held weights are charged against the utilization budget;
1409
+ the engine is granted only the remainder."""
1410
+ cache_config = CacheConfig(gpu_memory_utilization=0.9)
1411
+ # 70 GiB of a 100 GiB device is held externally before the worker starts.
1412
+ snapshot = _memory_snapshot(total_gib=100, free_gib=30)
1413
+
1414
+ with pytest.raises(ValueError, match="less than desired"):
1415
+ request_memory(snapshot, cache_config)
1416
+
1417
+ # 90 GiB budget - 70 GiB external = 20 GiB for the engine.
1418
+ assert request_memory(snapshot, cache_config, 70 * GiB_bytes) == 20 * GiB_bytes
1419
+
1420
+ # The external weights alone exceed the utilization budget.
1421
+ with pytest.raises(ValueError, match="exceed the desired"):
1422
+ request_memory(snapshot, cache_config, 95 * GiB_bytes)
1423
+
1424
+ # Other tenants squeeze free memory below the engine's remainder.
1425
+ with pytest.raises(ValueError, match="less than the engine's budget"):
1426
+ request_memory(
1427
+ _memory_snapshot(total_gib=100, free_gib=10), cache_config, 70 * GiB_bytes
1428
+ )
vllm/model_executor/model_loader/base_loader.py
+5
@@ -28,6 +28,11 @@ class BaseModelLoader(ABC):
28
def __init__(self, load_config: LoadConfig):
29
self.load_config = load_config
30
31
+ def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
32
+ """Get weights memory from external process;
33
+ 0 when the weights are not external."""
34
+ return 0
35
+
36
@abstractmethod
37
def download_model(self, model_config: ModelConfig) -> None:
38
"""Download a model so that it can be immediately loaded."""
vllm/model_executor/model_loader/weight_cache/daemon.py
+18
-1
@@ -78,6 +78,7 @@ cached and keep loading from disk in the engine.
78
79
import contextlib
80
import fcntl
81
+import gc
82
import multiprocessing
83
import os
84
import socket
@@ -116,6 +117,7 @@ from vllm.model_executor.model_loader.weight_cache.utils import (
117
is_draft_model_cacheable,
118
)
119
from vllm.platforms import current_platform
120
+from vllm.utils.mem_utils import format_gib
121
from vllm.utils.torch_utils import set_default_torch_dtype
122
from vllm.v1.worker.workspace import init_workspace_manager
123
@@ -249,10 +251,17 @@ class WeightCacheDaemon:
251
with set_current_vllm_config(self.vllm_config):
252
ensure_model_parallel_initialized(self.tp_size, self.pp_size)
253
self.model = self.get_model()
254
+ # Loading and post-processing leave freed transients in the caching
255
+ # allocator; return them so engines sharing the GPU can use them.
256
+ gc.collect()
257
+ torch.accelerator.empty_cache()
258
logger.info(
253
- "Weight cache %s daemon rank %d loaded model",
259
+ "Weight cache %s daemon rank %d loaded model (%s GiB allocated, "
260
+ "%s GiB reserved)",
261
self.role,
262
self.global_rank,
263
+ format_gib(torch.accelerator.memory_allocated()),
264
+ format_gib(torch.accelerator.memory_reserved()),
265
)
266
267
def get_model(self) -> torch.nn.Module:
@@ -369,6 +378,14 @@ class WeightCacheDaemon:
378
cmd = request.get("cmd")
379
if cmd == "get_state":
380
self._handle_get_state(conn, request)
381
+ elif cmd == "get_memory":
382
+ send_msg(
383
+ conn,
384
+ {
385
+ "status": "ok",
386
+ "memory_bytes": torch.accelerator.memory_allocated(),
387
+ },
388
+ )
389
elif cmd == "release":
390
self._handle_release(conn)
391
else:
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+66
-10
@@ -6,7 +6,9 @@ daemon via CUDA IPC instead of loading from disk."""
6
import dataclasses
7
import socket
8
import time
9
+from collections.abc import Callable
10
from copy import copy
11
+from typing import TypeVar
12
13
import torch
14
import torch.nn as nn
@@ -23,6 +25,7 @@ from vllm.logger import init_logger
25
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
26
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
27
from vllm.model_executor.model_loader.utils import (
28
+ get_draft_load_config,
29
initialize_model,
30
process_weights_after_loading,
31
)
@@ -40,6 +43,9 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
43
send_msg,
44
verify_socket_owner,
45
)
46
+from vllm.model_executor.model_loader.weight_cache.utils import (
47
+ is_draft_model_cacheable,
48
+)
49
from vllm.model_executor.utils import weights_already_processed
50
from vllm.tracing import instrument
51
from vllm.utils.torch_utils import set_default_torch_dtype
@@ -50,6 +56,8 @@ _CONNECT_TIMEOUT_S = 5.0
56
_STATE_TIMEOUT_S = 300.0
57
_STARTUP_RETRY_INTERVAL_S = 0.5
58
59
+_T = TypeVar("_T")
60
+
61
62
class IpcModelLoader(BaseModelLoader):
63
"""Loads a model by mapping the weight cache daemon's tensors via CUDA IPC.
@@ -110,6 +118,52 @@ class IpcModelLoader(BaseModelLoader):
118
f"{load_config.load_format}: {sorted(extra_config)}"
119
)
120
121
+ def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
122
+ # Copy mode clones the weights into this process; nothing external.
123
+ if self.mode != "zero_copy":
124
+ return 0
125
+ total = self._daemon_memory()
126
+ if is_draft_model_cacheable(vllm_config.speculative_config):
127
+ # The draft group is queried with the same config the engine
128
+ # would load the draft with.
129
+ draft_config = get_draft_load_config(vllm_config)
130
+ if draft_config.load_format == "ipc_cache":
131
+ draft_loader = IpcModelLoader(draft_config)
132
+ # Only a zero-copy draft-group daemon holds external weights;
133
+ # an explicit draft_load_config without is_draft would resolve
134
+ # back to the target socket and double-count it.
135
+ if draft_loader.is_draft and draft_loader.mode == "zero_copy":
136
+ total += draft_loader._daemon_memory()
137
+ return total
138
+
139
+ def _daemon_memory(self) -> int:
140
+ if not self.fallback:
141
+ # The loader waits for the daemon at load time, so the weights
142
+ # will be zero-copy mapped for sure; mirror that wait here since
143
+ # returning 0 would over-grant the memory budget.
144
+ return self._with_startup_wait(self._query_daemon_memory)
145
+ # An unreachable daemon means a disk load, i.e. nothing external.
146
+ try:
147
+ return self._query_daemon_memory()
148
+ except (WeightCacheUnavailableError, ConnectionError, OSError) as e:
149
+ logger.warning(
150
+ "Cannot query weight cache daemon memory (%s); "
151
+ "assuming the weights are not externally held",
152
+ e,
153
+ )
154
+ return 0
155
+
156
+ def _query_daemon_memory(self) -> int:
157
+ with self._connect(self.connect_timeout_s) as conn:
158
+ send_msg(conn, {"cmd": "get_memory"})
159
+ response = recv_msg(conn)
160
+ if response.get("status") != "ok":
161
+ raise WeightCacheUnavailableError(
162
+ "Weight cache daemon rejected the memory query: "
163
+ f"{response.get('message')}"
164
+ )
165
+ return int(response.get("memory_bytes", 0))
166
+
167
def download_model(self, model_config: ModelConfig) -> None:
168
DefaultModelLoader(self._fallback_load_config()).download_model(model_config)
169
@@ -308,23 +362,25 @@ class IpcModelLoader(BaseModelLoader):
362
def _request_state_with_startup_wait(
363
self, cache_config: WeightCacheKey
364
) -> WeightCacheState:
365
+ return self._with_startup_wait(lambda: self._request_state(cache_config))
366
+
367
+ def _with_startup_wait(self, op: Callable[[], _T]) -> _T:
368
+ """Retry op until the daemon answers or the state timeout elapses;
369
+ the daemon may still be loading the model when the engine starts."""
370
deadline = time.monotonic() + self.state_timeout_s
312
- waiting_logged = False
371
while True:
372
try:
315
- return self._request_state(cache_config)
316
- except WeightCacheUnavailableError as e:
373
+ return op()
374
+ except (WeightCacheUnavailableError, ConnectionError, OSError) as e:
375
if time.monotonic() >= deadline:
376
raise WeightCacheUnavailableError(
377
"Weight cache daemon did not become ready within "
320
- f"{self.state_timeout_s:.1f}s"
378
+ f"{self.state_timeout_s:.1f}s: {e}"
379
) from e
322
- if not waiting_logged:
323
- logger.info(
324
- "Waiting up to %.1fs for the weight cache daemon to start",
325
- self.state_timeout_s,
326
- )
327
- waiting_logged = True
380
+ logger.info_once(
381
+ "Waiting up to %.1fs for the weight cache daemon to start",
382
+ self.state_timeout_s,
383
+ )
384
time.sleep(
385
max(
386
0.0,
vllm/v1/worker/gpu_worker.py
+8
-1
@@ -57,6 +57,7 @@ from vllm.distributed.weight_transfer import (
57
)
58
from vllm.logger import init_logger
59
from vllm.lora.request import LoRARequest
60
+from vllm.model_executor.model_loader import get_model_loader
61
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
62
from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
63
from vllm.platforms import current_platform
@@ -487,7 +488,13 @@ class Worker(WorkerBase):
488
489
# take current memory snapshot
490
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
490
- self.requested_memory = request_memory(init_snapshot, self.cache_config)
491
+ # Weights from external model loader process
492
+ external_weight_memory = get_model_loader(
493
+ self.load_config
494
+ ).get_external_weight_memory(self.vllm_config)
495
+ self.requested_memory = request_memory(
496
+ init_snapshot, self.cache_config, external_weight_memory
497
+ )
498
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
499
logger.debug(
500
"worker requested memory: %sGiB", format_gib(self.requested_memory)
vllm/v1/worker/utils.py
+38
-1
@@ -536,7 +536,11 @@ def sanity_check_mm_encoder_outputs(
536
)
537
538
539
-def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int:
539
+def request_memory(
540
+ init_snapshot: MemorySnapshot,
541
+ cache_config: CacheConfig,
542
+ external_weight_memory: int = 0,
543
+) -> int:
544
"""Calculate the amount of memory required by vLLM, then validate
545
that the current amount of free memory is sufficient for that.
546
"""
@@ -544,6 +548,39 @@ def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) ->
548
init_snapshot.total_memory * cache_config.gpu_memory_utilization
549
)
550
551
+ if external_weight_memory > 0:
552
+ engine_memory = requested_memory - external_weight_memory
553
+ if engine_memory <= 0:
554
+ raise ValueError(
555
+ f"Externally held weights on device {init_snapshot.device_} "
556
+ f"({format_gib(external_weight_memory)}/"
557
+ f"{format_gib(init_snapshot.total_memory)} GiB) exceed the "
558
+ "desired GPU memory utilization "
559
+ f"({cache_config.gpu_memory_utilization}, "
560
+ f"{format_gib(requested_memory)} GiB). Increase GPU memory "
561
+ "utilization or reduce GPU memory used by other processes."
562
+ )
563
+ if init_snapshot.free_memory < engine_memory:
564
+ raise ValueError(
565
+ f"Free memory on device {init_snapshot.device_} "
566
+ f"({format_gib(init_snapshot.free_memory)}/"
567
+ f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
568
+ "is less than the engine's budget after excluding "
569
+ "external process's weights "
570
+ f"({cache_config.gpu_memory_utilization}, "
571
+ f"{format_gib(engine_memory)} GiB). Decrease GPU memory "
572
+ "utilization or reduce GPU memory used by other processes."
573
+ )
574
+ logger.info_once(
575
+ "Weights are held outside this process: of the %s GiB "
576
+ "utilization budget, %s GiB is externally held and "
577
+ "%s GiB remains for the engine's own allocations.",
578
+ format_gib(requested_memory),
579
+ format_gib(external_weight_memory),
580
+ format_gib(engine_memory),
581
+ )
582
+ return engine_memory
583
+
584
if init_snapshot.free_memory < requested_memory:
585
raise ValueError(
586
f"Free memory on device {init_snapshot.device_} "