[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>

tests/model_executor/model_loader/test_weight_cache.py

@@ -354,3 +354,18 @@ def test_weight_cache_key_distinguishes_dp_ranks():
     assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"]
     assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"]
     assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"]
+
+
+def test_ipc_loader_copy_mode_reports_no_external_weight_memory():
+    """Copy mode clones the weights into the engine, so nothing is external
+    (vllm_config is never touched)."""
+    from vllm.config import LoadConfig
+    from vllm.model_executor.model_loader.weight_cache.ipc_loader import (
+        IpcModelLoader,
+    )
+
+    copy_mode = LoadConfig(
+        load_format="ipc_cache", model_loader_extra_config={"mode": "copy"}
+    )
+    loader = IpcModelLoader(copy_mode)
+    assert loader.get_external_weight_memory(None) == 0  # type: ignore[arg-type]

tests/v1/worker/test_utils.py

@@ -13,7 +13,7 @@ import torch
 
 import vllm.v1.attention.backends.mla.index_group as index_group_module
 import vllm.v1.hisparse.runtime as hisparse_runtime_module
-from vllm.config import CUDAGraphMode
+from vllm.config import CacheConfig, CUDAGraphMode
 from vllm.config.mamba import MambaBackendEnum, MambaConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (
     worker as hisparse_worker_module,
@@ -25,12 +25,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
     _SlotMappingStaging,
 )
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.mem_utils import MemorySnapshot
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
 from vllm.v1.worker.utils import (
     bind_kv_cache,
     bind_kv_cache_to_layers,
     copy_kv_cache_blocks_inplace,
+    request_memory,
 )
 
 
@@ -1390,3 +1393,36 @@ def test_bind_kv_cache_draft_model(default_vllm_config):
     assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"]
     assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"]
     assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"]
+
+
+def _memory_snapshot(total_gib: int, free_gib: int) -> MemorySnapshot:
+    return MemorySnapshot(
+        free_memory=free_gib * GiB_bytes,
+        total_memory=total_gib * GiB_bytes,
+        device="cpu",
+        auto_measure=False,
+    )
+
+
+def test_request_memory_charges_external_weights():
+    """Externally held weights are charged against the utilization budget;
+    the engine is granted only the remainder."""
+    cache_config = CacheConfig(gpu_memory_utilization=0.9)
+    # 70 GiB of a 100 GiB device is held externally before the worker starts.
+    snapshot = _memory_snapshot(total_gib=100, free_gib=30)
+
+    with pytest.raises(ValueError, match="less than desired"):
+        request_memory(snapshot, cache_config)
+
+    # 90 GiB budget - 70 GiB external = 20 GiB for the engine.
+    assert request_memory(snapshot, cache_config, 70 * GiB_bytes) == 20 * GiB_bytes
+
+    # The external weights alone exceed the utilization budget.
+    with pytest.raises(ValueError, match="exceed the desired"):
+        request_memory(snapshot, cache_config, 95 * GiB_bytes)
+
+    # Other tenants squeeze free memory below the engine's remainder.
+    with pytest.raises(ValueError, match="less than the engine's budget"):
+        request_memory(
+            _memory_snapshot(total_gib=100, free_gib=10), cache_config, 70 * GiB_bytes
+        )

vllm/model_executor/model_loader/base_loader.py

@@ -28,6 +28,11 @@ class BaseModelLoader(ABC):
     def __init__(self, load_config: LoadConfig):
         self.load_config = load_config
 
+    def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
+        """Get weights memory from external process;
+        0 when the weights are not external."""
+        return 0
+
     @abstractmethod
     def download_model(self, model_config: ModelConfig) -> None:
         """Download a model so that it can be immediately loaded."""

vllm/model_executor/model_loader/weight_cache/daemon.py

@@ -78,6 +78,7 @@ cached and keep loading from disk in the engine.
 
 import contextlib
 import fcntl
+import gc
 import multiprocessing
 import os
 import socket
@@ -116,6 +117,7 @@ from vllm.model_executor.model_loader.weight_cache.utils import (
     is_draft_model_cacheable,
 )
 from vllm.platforms import current_platform
+from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import set_default_torch_dtype
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -249,10 +251,17 @@ class WeightCacheDaemon:
         with set_current_vllm_config(self.vllm_config):
             ensure_model_parallel_initialized(self.tp_size, self.pp_size)
             self.model = self.get_model()
+        # Loading and post-processing leave freed transients in the caching
+        # allocator; return them so engines sharing the GPU can use them.
+        gc.collect()
+        torch.accelerator.empty_cache()
         logger.info(
-            "Weight cache %s daemon rank %d loaded model",
+            "Weight cache %s daemon rank %d loaded model (%s GiB allocated, "
+            "%s GiB reserved)",
             self.role,
             self.global_rank,
+            format_gib(torch.accelerator.memory_allocated()),
+            format_gib(torch.accelerator.memory_reserved()),
         )
 
     def get_model(self) -> torch.nn.Module:
@@ -369,6 +378,14 @@ class WeightCacheDaemon:
         cmd = request.get("cmd")
         if cmd == "get_state":
             self._handle_get_state(conn, request)
+        elif cmd == "get_memory":
+            send_msg(
+                conn,
+                {
+                    "status": "ok",
+                    "memory_bytes": torch.accelerator.memory_allocated(),
+                },
+            )
         elif cmd == "release":
             self._handle_release(conn)
         else:

vllm/model_executor/model_loader/weight_cache/ipc_loader.py

@@ -6,7 +6,9 @@ daemon via CUDA IPC instead of loading from disk."""
 import dataclasses
 import socket
 import time
+from collections.abc import Callable
 from copy import copy
+from typing import TypeVar
 
 import torch
 import torch.nn as nn
@@ -23,6 +25,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
 from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
 from vllm.model_executor.model_loader.utils import (
+    get_draft_load_config,
     initialize_model,
     process_weights_after_loading,
 )
@@ -40,6 +43,9 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
     send_msg,
     verify_socket_owner,
 )
+from vllm.model_executor.model_loader.weight_cache.utils import (
+    is_draft_model_cacheable,
+)
 from vllm.model_executor.utils import weights_already_processed
 from vllm.tracing import instrument
 from vllm.utils.torch_utils import set_default_torch_dtype
@@ -50,6 +56,8 @@ _CONNECT_TIMEOUT_S = 5.0
 _STATE_TIMEOUT_S = 300.0
 _STARTUP_RETRY_INTERVAL_S = 0.5
 
+_T = TypeVar("_T")
+
 
 class IpcModelLoader(BaseModelLoader):
     """Loads a model by mapping the weight cache daemon's tensors via CUDA IPC.
@@ -110,6 +118,52 @@ class IpcModelLoader(BaseModelLoader):
                 f"{load_config.load_format}: {sorted(extra_config)}"
             )
 
+    def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
+        # Copy mode clones the weights into this process; nothing external.
+        if self.mode != "zero_copy":
+            return 0
+        total = self._daemon_memory()
+        if is_draft_model_cacheable(vllm_config.speculative_config):
+            # The draft group is queried with the same config the engine
+            # would load the draft with.
+            draft_config = get_draft_load_config(vllm_config)
+            if draft_config.load_format == "ipc_cache":
+                draft_loader = IpcModelLoader(draft_config)
+                # Only a zero-copy draft-group daemon holds external weights;
+                # an explicit draft_load_config without is_draft would resolve
+                # back to the target socket and double-count it.
+                if draft_loader.is_draft and draft_loader.mode == "zero_copy":
+                    total += draft_loader._daemon_memory()
+        return total
+
+    def _daemon_memory(self) -> int:
+        if not self.fallback:
+            # The loader waits for the daemon at load time, so the weights
+            # will be zero-copy mapped for sure; mirror that wait here since
+            # returning 0 would over-grant the memory budget.
+            return self._with_startup_wait(self._query_daemon_memory)
+        # An unreachable daemon means a disk load, i.e. nothing external.
+        try:
+            return self._query_daemon_memory()
+        except (WeightCacheUnavailableError, ConnectionError, OSError) as e:
+            logger.warning(
+                "Cannot query weight cache daemon memory (%s); "
+                "assuming the weights are not externally held",
+                e,
+            )
+            return 0
+
+    def _query_daemon_memory(self) -> int:
+        with self._connect(self.connect_timeout_s) as conn:
+            send_msg(conn, {"cmd": "get_memory"})
+            response = recv_msg(conn)
+        if response.get("status") != "ok":
+            raise WeightCacheUnavailableError(
+                "Weight cache daemon rejected the memory query: "
+                f"{response.get('message')}"
+            )
+        return int(response.get("memory_bytes", 0))
+
     def download_model(self, model_config: ModelConfig) -> None:
         DefaultModelLoader(self._fallback_load_config()).download_model(model_config)
 
@@ -308,23 +362,25 @@ class IpcModelLoader(BaseModelLoader):
     def _request_state_with_startup_wait(
         self, cache_config: WeightCacheKey
     ) -> WeightCacheState:
+        return self._with_startup_wait(lambda: self._request_state(cache_config))
+
+    def _with_startup_wait(self, op: Callable[[], _T]) -> _T:
+        """Retry op until the daemon answers or the state timeout elapses;
+        the daemon may still be loading the model when the engine starts."""
         deadline = time.monotonic() + self.state_timeout_s
-        waiting_logged = False
         while True:
             try:
-                return self._request_state(cache_config)
-            except WeightCacheUnavailableError as e:
+                return op()
+            except (WeightCacheUnavailableError, ConnectionError, OSError) as e:
                 if time.monotonic() >= deadline:
                     raise WeightCacheUnavailableError(
                         "Weight cache daemon did not become ready within "
-                        f"{self.state_timeout_s:.1f}s"
+                        f"{self.state_timeout_s:.1f}s: {e}"
                     ) from e
-                if not waiting_logged:
-                    logger.info(
-                        "Waiting up to %.1fs for the weight cache daemon to start",
-                        self.state_timeout_s,
-                    )
-                    waiting_logged = True
+                logger.info_once(
+                    "Waiting up to %.1fs for the weight cache daemon to start",
+                    self.state_timeout_s,
+                )
                 time.sleep(
                     max(
                         0.0,

vllm/v1/worker/gpu_worker.py

@@ -57,6 +57,7 @@ from vllm.distributed.weight_transfer import (
 )
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.model_loader import get_model_loader
 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
 from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
 from vllm.platforms import current_platform
@@ -487,7 +488,13 @@ class Worker(WorkerBase):
 
             # take current memory snapshot
             self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
-            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            # Weights from external model loader process
+            external_weight_memory = get_model_loader(
+                self.load_config
+            ).get_external_weight_memory(self.vllm_config)
+            self.requested_memory = request_memory(
+                init_snapshot, self.cache_config, external_weight_memory
+            )
             logger.debug("worker init memory snapshot: %r", self.init_snapshot)
             logger.debug(
                 "worker requested memory: %sGiB", format_gib(self.requested_memory)

vllm/v1/worker/utils.py

@@ -536,7 +536,11 @@ def sanity_check_mm_encoder_outputs(
     )
 
 
-def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int:
+def request_memory(
+    init_snapshot: MemorySnapshot,
+    cache_config: CacheConfig,
+    external_weight_memory: int = 0,
+) -> int:
     """Calculate the amount of memory required by vLLM, then validate
     that the current amount of free memory is sufficient for that.
     """
@@ -544,6 +548,39 @@ def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) ->
         init_snapshot.total_memory * cache_config.gpu_memory_utilization
     )
 
+    if external_weight_memory > 0:
+        engine_memory = requested_memory - external_weight_memory
+        if engine_memory <= 0:
+            raise ValueError(
+                f"Externally held weights on device {init_snapshot.device_} "
+                f"({format_gib(external_weight_memory)}/"
+                f"{format_gib(init_snapshot.total_memory)} GiB) exceed the "
+                "desired GPU memory utilization "
+                f"({cache_config.gpu_memory_utilization}, "
+                f"{format_gib(requested_memory)} GiB). Increase GPU memory "
+                "utilization or reduce GPU memory used by other processes."
+            )
+        if init_snapshot.free_memory < engine_memory:
+            raise ValueError(
+                f"Free memory on device {init_snapshot.device_} "
+                f"({format_gib(init_snapshot.free_memory)}/"
+                f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
+                "is less than the engine's budget after excluding "
+                "external process's weights "
+                f"({cache_config.gpu_memory_utilization}, "
+                f"{format_gib(engine_memory)} GiB). Decrease GPU memory "
+                "utilization or reduce GPU memory used by other processes."
+            )
+        logger.info_once(
+            "Weights are held outside this process: of the %s GiB "
+            "utilization budget, %s GiB is externally held and "
+            "%s GiB remains for the engine's own allocations.",
+            format_gib(requested_memory),
+            format_gib(external_weight_memory),
+            format_gib(engine_memory),
+        )
+        return engine_memory
+
     if init_snapshot.free_memory < requested_memory:
         raise ValueError(
             f"Free memory on device {init_snapshot.device_} "