@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dbf49dad11

[Fast Start] Use GPU uuid as socket folder identifier (#56669)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

Isotr0py committed Sep 14, 2026 at 14:28 UTC dbf49dad11e30ec5d464be74b715a9a9e31c2c5a

4 files changed +30 -69

tests/model_executor/model_loader/test_weight_cache.py

+3 -5

@@ -7,6 +7,7 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve

7 identical outputs.

8 """

10 +import glob

11 import os

12 import shutil

13 import subprocess

@@ -21,7 +22,6 @@ import pytest

23 from vllm import SamplingParams

24 from vllm.assets.image import ImageAsset

24 -from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path

25 from vllm.platforms import current_platform

27 DAEMON_TIMEOUT_S = 600

@@ -79,9 +79,7 @@ class WeightCacheDaemon:

79 # Poll for the socket files rather than a log line: model loading can

80 # pull in JIT compilers that swap the process's stderr and swallow

81 # everything logged afterwards, making log-based readiness flaky.

82 - expected = [

83 - get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size)

84 - ]

82 + pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock")

83 deadline = time.monotonic() + timeout_s

84 while time.monotonic() < deadline:

85 if self._proc.poll() is not None:

@@ -89,7 +87,7 @@ class WeightCacheDaemon:

87 f"Weight cache daemon exited with {self._proc.returncode}:\n"

88 f"{self._logs()}"

89 )

92 - if all(os.path.exists(path) for path in expected):

90 + if len(glob.glob(pattern)) >= self.tp_size:

91 return

92 time.sleep(1.0)

93 raise TimeoutError(

vllm/model_executor/model_loader/weight_cache/daemon.py

+8 -18

@@ -60,6 +60,7 @@ from vllm.distributed import (

60 ensure_model_parallel_initialized,

61 init_distributed_environment,

62 )

63 +from vllm.engine.arg_utils import EngineArgs

64 from vllm.logger import init_logger

65 from vllm.model_executor.model_loader import get_model_loader

66 from vllm.model_executor.model_loader.utils import process_weights_after_loading

@@ -70,13 +71,15 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (

71 check_ipc_platform_support,

72 check_ipc_quant_support,

73 ensure_private_socket_dir,

73 - get_physical_device_id,

74 + get_current_device_uuid,

75 get_socket_path,

76 recv_msg,

77 send_msg,

78 verify_peer_is_owner,

79 )

80 from vllm.platforms import current_platform

81 +from vllm.utils.argparse_utils import FlexibleArgumentParser

82 +from vllm.utils.network_utils import get_distributed_init_method, get_open_port

83 from vllm.utils.torch_utils import set_default_torch_dtype

85 logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")

@@ -209,7 +212,7 @@ class WeightCacheDaemon:

212 ready_callback: Invoked once the socket is bound and listening,

213 so the launcher can report overall readiness.

214 """

212 - socket_path = self._socket_path()

215 + socket_path = self._socket_path

216 ensure_private_socket_dir(

217 os.path.dirname(socket_path), strict_perms=self.socket_dir is None

218 )

@@ -268,12 +271,9 @@ class WeightCacheDaemon:

271 ) from e

272 return lock_fd

273

274 + @property

275 def _socket_path(self) -> str:

272 - device_index = torch.accelerator.current_device_index()

273 - gpu_id = get_physical_device_id(device_index)

274 - if gpu_id is None:

275 - gpu_id = device_index

276 - return get_socket_path(gpu_id, self.socket_dir)

276 + return get_socket_path(get_current_device_uuid(), self.socket_dir)

277

278 def _handle_connection(self, conn: socket.socket) -> None:

279 request = recv_msg(conn)

@@ -304,7 +304,7 @@ class WeightCacheDaemon:

304 "status": "ok",

305 "entries": self.entries,

306 "aliases": self.aliases,

307 - "gpu_uuid": self._gpu_uuid(),

307 + "gpu_uuid": get_current_device_uuid(),

308 },

309 )

310 logger.info_once(

@@ -322,12 +322,6 @@ class WeightCacheDaemon:

322 logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank)

323 send_msg(conn, {"status": "ok"})

324

325 - def _gpu_uuid(self) -> str:

326 - props = torch.cuda.get_device_properties(

327 - torch.accelerator.current_device_index()

328 - )

329 - return str(props.uuid)

330 -

325

326 def _run_daemon(

327 tp_rank: int,

@@ -359,10 +353,6 @@ def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:

353

354

355 def main() -> None:

362 - from vllm.engine.arg_utils import EngineArgs

363 - from vllm.utils.argparse_utils import FlexibleArgumentParser

364 - from vllm.utils.network_utils import get_distributed_init_method, get_open_port

365 -

356 parser = FlexibleArgumentParser(

357 description="Launch weight cache daemons (one per TP rank)."

358 )

vllm/model_executor/model_loader/weight_cache/ipc_loader.py

+8 -19

@@ -12,6 +12,10 @@ import torch.nn as nn

13 from vllm.config import ModelConfig, VllmConfig

14 from vllm.config.load import LoadConfig

15 +from vllm.distributed import (

16 + get_tensor_model_parallel_rank,

17 + get_tensor_model_parallel_world_size,

18 +)

19 from vllm.logger import init_logger

20 from vllm.model_executor.model_loader.base_loader import BaseModelLoader

21 from vllm.model_executor.model_loader.default_loader import DefaultModelLoader

@@ -27,7 +31,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (

31 WeightCacheUnavailableError,

32 check_ipc_platform_support,

33 check_ipc_quant_support,

30 - get_physical_device_id,

34 + get_current_device_uuid,

35 get_socket_path,

36 recv_msg,

37 send_msg,

@@ -56,7 +60,7 @@ class IpcModelLoader(BaseModelLoader):

60 Extra config keys (via --model-loader-extra-config):

62 - socket_path: explicit daemon socket path. Defaults to a per-GPU path

59 - derived from the physical GPU id.

63 + derived from the physical GPU uuid.

64 - socket_dir: directory containing the daemon sockets.

65 - mode: "zero_copy" (default) or "copy".

66 - fallback: fall back to disk loading when the daemon is unavailable or

@@ -261,11 +265,6 @@ class IpcModelLoader(BaseModelLoader):

265 def _fetch_entries(

266 self, model_config: ModelConfig

267 ) -> tuple[dict[str, TensorEntry], dict[str, str]]:

264 - from vllm.distributed import (

265 - get_tensor_model_parallel_rank,

266 - get_tensor_model_parallel_world_size,

267 - )

268 -

268 cache_config = WeightCacheKey.from_model_config(

269 model_config,

270 tp_size=get_tensor_model_parallel_world_size(),

@@ -317,22 +316,12 @@ class IpcModelLoader(BaseModelLoader):

316 def _resolve_socket_path(self) -> str:

317 if self.socket_path is not None:

318 return self.socket_path

320 - device_index = torch.accelerator.current_device_index()

321 - gpu_id = get_physical_device_id(device_index)

322 - if gpu_id is None:

323 - raise WeightCacheUnavailableError(

324 - "Cannot infer the physical GPU id from CUDA_VISIBLE_DEVICES; "

325 - "pass socket_path via --model-loader-extra-config"

326 - )

327 - return get_socket_path(gpu_id, self.socket_dir)

319 + return get_socket_path(get_current_device_uuid(), self.socket_dir)

320

321 def _check_gpu_uuid(self, daemon_uuid: str | None) -> None:

322 if daemon_uuid is None:

323 return

332 - props = torch.cuda.get_device_properties(

333 - torch.accelerator.current_device_index()

334 - )

335 - local_uuid = str(props.uuid)

324 + local_uuid = get_current_device_uuid()

325 if daemon_uuid != local_uuid:

326 raise CacheConfigMismatchError(

327 f"Daemon GPU {daemon_uuid} != engine GPU {local_uuid}; "

vllm/model_executor/model_loader/weight_cache/protocol.py

+11 -27

@@ -23,14 +23,19 @@ from dataclasses import dataclass, fields

23 from typing import Any

25 import torch

26 +from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor

27 +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME

29 import vllm.version

30 from vllm.config import ModelConfig

31 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase

32 +from vllm.model_executor.model_loader.weight_utils import (

33 + filter_duplicate_safetensors_files,

34 +)

35 from vllm.platforms import current_platform

36 from vllm.utils.hashing import safe_hash

33 -SOCKET_NAME_TEMPLATE = "vllm_weight_cache_gpu{gpu_id}.sock"

38 +SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}.sock"

39 SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}"

41 _LEN_STRUCT = struct.Struct("!Q")

@@ -102,20 +107,9 @@ def check_ipc_quant_support(model: torch.nn.Module) -> None:

107 )

108

109

105 -def get_physical_device_id(device_index: int) -> int | None:

106 - """Map a local CUDA device index to the physical GPU id.

107 -

108 - Returns None if CUDA_VISIBLE_DEVICES contains non-integer entries

109 - (e.g. GPU UUIDs), in which case an explicit socket path is required.

110 - """

111 - visible = os.environ.get("CUDA_VISIBLE_DEVICES")

112 - if not visible:

113 - return device_index

114 - entries = visible.split(",")

115 - try:

116 - return int(entries[device_index])

117 - except (IndexError, ValueError):

118 - return None

110 +def get_current_device_uuid() -> str:

111 + """UUID of the physical GPU backing the current accelerator device."""

112 + return current_platform.get_device_uuid(torch.accelerator.current_device_index())

113

114

115 def get_socket_dir(socket_dir: str | None = None) -> str:

@@ -132,9 +126,9 @@ def get_socket_dir(socket_dir: str | None = None) -> str:

126 )

127

128

135 -def get_socket_path(gpu_id: int, socket_dir: str | None = None) -> str:

129 +def get_socket_path(gpu_uuid: str, socket_dir: str | None = None) -> str:

130 directory = get_socket_dir(socket_dir)

137 - return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_id=gpu_id))

131 + return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_uuid=gpu_uuid))

132

133

134 def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None:

@@ -245,12 +239,6 @@ def hash_checkpoint(model: str) -> str | None:

239 """

240 if not os.path.isdir(model):

241 return None

248 - from transformers.utils import SAFE_WEIGHTS_INDEX_NAME

249 -

250 - from vllm.model_executor.model_loader.weight_utils import (

251 - filter_duplicate_safetensors_files,

252 - )

253 -

242 files = glob.glob(os.path.join(model, "*.safetensors"))

243 if os.path.isfile(os.path.join(model, SAFE_WEIGHTS_INDEX_NAME)):

244 files = filter_duplicate_safetensors_files(

@@ -337,8 +325,6 @@ class TensorEntry:

325

326 @classmethod

327 def from_tensor(cls, tensor: torch.Tensor, kind: str) -> "TensorEntry":

340 - from torch.multiprocessing.reductions import reduce_tensor

341 -

328 tensor = tensor.detach()

329 if tensor.is_cuda:

330 _, ipc_args = reduce_tensor(tensor)

@@ -349,8 +335,6 @@ class TensorEntry:

335 if self.ipc_args is None:

336 assert self.cpu_tensor is not None

337 return self.cpu_tensor

352 - from torch.multiprocessing.reductions import rebuild_cuda_tensor

353 -

338 args = list(self.ipc_args)

339 # Index 6 of the args from reduce_tensor is the device index. It must

340 # be retargeted to the local index since the daemon and the engine may