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

Fast Start (#54921)

Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: liusy58 <liusy58@smail.nju.edu.cn> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

siyu committed Sep 4, 2026 at 21:32 UTC a69e75b9b6d6a26d90b6790e2eb75b3419a47c16

16 files changed +1529 -13

tests/model_executor/model_loader/test_weight_cache.py new

+153

@@ -0,0 +1,153 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""End-to-end tests for the IPC weight cache loader.

4 +

5 +A cold start (no weight cache daemon, so the loader falls back to disk) and

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

7 +identical outputs.

8 +"""

9 +

10 +import shutil

11 +import subprocess

12 +import sys

13 +import tempfile

14 +import threading

15 +import time

16 +

17 +from vllm import SamplingParams

18 +

19 +MODEL = "Qwen/Qwen3.5-0.8B"

20 +PROMPTS = [

21 + "Hello, my name is",

22 + "The capital of France is",

23 +]

24 +DAEMON_TIMEOUT_S = 600

25 +

26 +

27 +class WeightCacheDaemon:

28 + """Context manager running the real weight cache daemon as a subprocess."""

29 +

30 + _READY_MARKER = "Weight cache daemon READY"

31 +

32 + def __init__(self, model: str, tp_size: int):

33 + # Short base path: Unix socket paths are limited to ~107 characters.

34 + self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")

35 + self._cmd = [

36 + sys.executable,

37 + "-m",

38 + "vllm.model_executor.model_loader.weight_cache.daemon",

39 + "--model",

40 + model,

41 + "--tensor-parallel-size",

42 + str(tp_size),

43 + "--weight-cache-socket-dir",

44 + self.socket_dir,

45 + "--enforce-eager",

46 + ]

47 + self._proc: subprocess.Popen | None = None

48 + self._lines: list[str] = []

49 +

50 + def __enter__(self) -> "WeightCacheDaemon":

51 + self._proc = subprocess.Popen(

52 + self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True

53 + )

54 + threading.Thread(target=self._drain_stderr, daemon=True).start()

55 + try:

56 + self._wait_ready(DAEMON_TIMEOUT_S)

57 + except Exception:

58 + self._stop()

59 + raise

60 + return self

61 +

62 + def __exit__(self, *exc_info) -> None:

63 + self._stop()

64 +

65 + def _drain_stderr(self) -> None:

66 + assert self._proc is not None and self._proc.stderr is not None

67 + for line in self._proc.stderr:

68 + self._lines.append(line)

69 +

70 + def _logs(self) -> str:

71 + return "".join(self._lines)

72 +

73 + def _wait_ready(self, timeout_s: float) -> None:

74 + assert self._proc is not None

75 + deadline = time.monotonic() + timeout_s

76 + while time.monotonic() < deadline:

77 + if self._proc.poll() is not None:

78 + raise RuntimeError(

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

80 + f"{self._logs()}"

81 + )

82 + if self._READY_MARKER in self._logs():

83 + return

84 + time.sleep(1.0)

85 + raise TimeoutError(

86 + f"Weight cache daemon not ready after {timeout_s}s:\n{self._logs()}"

87 + )

88 +

89 + def _stop(self) -> None:

90 + assert self._proc is not None

91 + self._proc.terminate()

92 + try:

93 + self._proc.wait(timeout=60)

94 + except subprocess.TimeoutExpired:

95 + self._proc.kill()

96 + self._proc.wait()

97 + shutil.rmtree(self.socket_dir, ignore_errors=True)

98 +

99 +

100 +def generate(vllm_runner, socket_dir: str | None, fallback: bool, **kwargs):

101 + extra_config = (

102 + {} if socket_dir is None else {"socket_dir": socket_dir, "fallback": fallback}

103 + )

104 + with vllm_runner(

105 + MODEL,

106 + load_format="auto" if socket_dir is None else "ipc_cache",

107 + model_loader_extra_config=extra_config,

108 + # Greedy outputs are compared across engine restarts; cap the batch

109 + # size so scheduling cannot change the results.

110 + max_num_seqs=1,

111 + **kwargs,

112 + ) as llm:

113 + sampling_params = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True)

114 + return llm.generate(PROMPTS, sampling_params)

115 +

116 +

117 +def test_ipc_cache_cold_start_and_warm_restart(vllm_runner):

118 + """Cold start falls back to disk; warm restarts load weights via CUDA IPC.

119 +

120 + All runs must produce outputs identical to a default-loader baseline. The

121 + warm runs disable the disk fallback, so they only pass if the weights

122 + really came from the daemon.

123 + """

124 + llm_kwargs = dict(

125 + gpu_memory_utilization=0.3,

126 + enforce_eager=True,

127 + # Qwen3.5 is a hybrid mamba model; its mamba cache mode requires

128 + # chunked prefill.

129 + enable_chunked_prefill=True,

130 + )

131 +

132 + # Baseline: plain disk loading with the default loader.

133 + baseline_outputs = generate(vllm_runner, None, fallback=True, **llm_kwargs)

134 + assert all(text for _, texts in baseline_outputs for text in texts)

135 +

136 + # Cold start: no daemon is serving, so the loader falls back to disk.

137 + with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:

138 + cold_outputs = generate(

139 + vllm_runner, empty_socket_dir, fallback=True, **llm_kwargs

140 + )

141 +

142 + with WeightCacheDaemon(MODEL, tp_size=1) as daemon:

143 + warm_outputs = generate(

144 + vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs

145 + )

146 + # Warm restart: a second engine lifetime against the same daemon.

147 + restart_outputs = generate(

148 + vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs

149 + )

150 +

151 + assert cold_outputs == baseline_outputs

152 + assert warm_outputs == baseline_outputs

153 + assert restart_outputs == baseline_outputs

tools/pre_commit/check_forbidden_imports.py

@@ -48,6 +48,7 @@ CHECK_IMPORTS = {

48 "vllm/distributed/device_communicators/shm_object_storage.py",

49 "vllm/distributed/weight_transfer/ipc_engine.py",

50 "vllm/distributed/weight_transfer/clients.py",

51 + "vllm/model_executor/model_loader/weight_cache/protocol.py",

52 "tests/distributed/test_shm_broadcast.py",

53 "tests/distributed/test_weight_transfer.py",

54 "vllm/utils/hashing.py",

vllm/config/load.py

@@ -38,6 +38,10 @@ class LoadConfig:

38 - "instanttensor" will load the Safetensors weights on CUDA devices using

39 InstantTensor, which enables distributed loading with pipelined prefetching

40 and fast direct I/O.

41 + - "ipc_cache" will map post-quantized weights from a local weight cache

42 + daemon via CUDA IPC for fast engine restarts. See

43 + `vllm/model_executor/model_loader/weight_cache/daemon.py` for how to

44 + launch the daemon.

45 - "npcache" will load the weights in pytorch format and store a numpy cache

46 to speed up the loading.

47 - "dummy" will initialize the weights with random values, which is mainly

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+27 -11

@@ -27,7 +27,11 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

27 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (

28 SharedExperts,

29 )

30 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs

30 +from vllm.model_executor.utils import (

31 + is_weights_pre_processed,

32 + replace_parameter,

33 + set_weight_attrs,

34 +)

35 from vllm.platforms import current_platform

37 if TYPE_CHECKING:

@@ -43,6 +47,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

48 # --8<-- [end:unquantized_fused_moe]

50 + supports_pre_processed_weights = True

51 +

52 def __init__(self, moe: FusedMoEConfig):

53 super().__init__(moe)

54 self.unquantized_backend, self.experts_cls = select_unquantized_moe_backend(

@@ -153,16 +159,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

159 # optional w{13,2}_bias references and SwiGLU gate params. Since

160 # weight updates mutate those bias tensors in place, the kernel

161 # does not need to be re-built.

156 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)

157 - assert self.moe_quant_config is not None

158 - assert self.experts_cls is not None

159 - self.moe_kernel = make_unquantized_moe_kernel(

160 - quant_config=self.moe_quant_config,

161 - moe_config=self.moe,

162 - backend=self.unquantized_backend,

163 - experts_cls=self.experts_cls,

164 - routing_tables=layer._expert_routing_tables(),

165 - )

162 + self._init_moe_kernel(layer)

163

164 if self.unquantized_backend == UnquantizedMoeBackend.CPU:

165 # The CPU experts need the layer itself for the setup that

@@ -170,11 +167,30 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

167 # it only sees the two weight tensors: padding and prepacking

168 # into the grouped-gemm layout (bias included), and capturing

169 # the router config that monolithic apply() cannot carry.

170 + assert self.moe_kernel is not None

171 self.moe_kernel.fused_experts.process_weights_after_loading(layer)

172

173 + def _init_moe_kernel(self, layer: "RoutedExperts") -> None:

174 + """Build the MoE kernel from the layer's current (shuffled) weights."""

175 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

176 + assert self.moe_quant_config is not None

177 + assert self.experts_cls is not None

178 + self.moe_kernel = make_unquantized_moe_kernel(

179 + quant_config=self.moe_quant_config,

180 + moe_config=self.moe,

181 + backend=self.unquantized_backend,

182 + experts_cls=self.experts_cls,

183 + routing_tables=layer._expert_routing_tables(),

184 + )

185 +

186 def process_weights_after_loading(self, layer: "RoutedExperts") -> None:

187 super().process_weights_after_loading(layer)

188

189 + if is_weights_pre_processed():

190 + # Weights are already in runtime format; rebuild the kernel only.

191 + self._init_moe_kernel(layer)

192 + return

193 +

194 # Padding may allocate a same-shaped strided view. Copy it back instead

195 # of rebinding `.data`, since CUDA graphs capture the parameter storage

196 # address.

vllm/model_executor/layers/linear.py

@@ -163,6 +163,8 @@ class LinearMethodBase(QuantizeMethodBase):

163 class UnquantizedLinearMethod(LinearMethodBase):

164 """Linear method without quantization."""

165

166 + supports_pre_processed_weights = True

167 +

168 def __init__(self) -> None:

169 config = get_current_vllm_config_or_none()

170 linear_backend = (

vllm/model_executor/layers/quantization/base_config.py

@@ -27,6 +27,13 @@ class QuantizeMethodBase(ABC):

27 in process_weights_after_loading, reducing peak memory during loading.

28 """

30 + supports_pre_processed_weights: bool = False

31 + """

32 + Whether ``process_weights_after_loading`` supports running under

33 + ``weights_already_processed``. Methods must skip tensor transforms in

34 + that mode; the loader driver rejects methods that do not declare support.

35 + """

36 +

37 @abstractmethod

38 def create_weights(

39 self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs

vllm/model_executor/layers/quantization/fp8.py

+31 -1

@@ -74,7 +74,11 @@ from vllm.model_executor.parameter import (

74 BlockQuantScaleParameter,

75 PerTensorScaleParameter,

76 )

77 -from vllm.model_executor.utils import replace_parameter, set_weight_attrs

77 +from vllm.model_executor.utils import (

78 + is_weights_pre_processed,

79 + replace_parameter,

80 + set_weight_attrs,

81 +)

82 from vllm.platforms import current_platform

83 from vllm.utils.deep_gemm import (

84 is_deep_gemm_supported,

@@ -248,6 +252,8 @@ class Fp8LinearMethod(LinearMethodBase):

252 quant_config: The quantization config.

253 """

254

255 + supports_pre_processed_weights = True

256 +

257 def __init__(self, quant_config: Fp8Config):

258 self.quant_config = quant_config

259 self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False)

@@ -366,6 +372,19 @@ class Fp8LinearMethod(LinearMethodBase):

372 self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)

373

374 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

375 + if is_weights_pre_processed():

376 + # Weights are already in runtime format; the only state tensor

377 + # export cannot carry is the `input_scale = None` stamp that dynamic

378 + # activation quantization relies on, since create_weights never

379 + # registered that parameter.

380 + if not self.act_q_static and not hasattr(layer, "input_scale"):

381 + layer.input_scale = None

382 + # Marlin reads marlin_input_dtype in apply_weights; it lives on the

383 + # method (not exported with the weights), so restore it here too.

384 + if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):

385 + self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype

386 + return

387 +

388 if self.use_marlin:

389 if not self.block_quant:

390 # Canonicalize to (K, N) for the kernel.

@@ -472,6 +491,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

491 quant_config: The quantization config.

492 """

493

494 + supports_pre_processed_weights = True

495 +

496 def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):

497 super().__init__(layer.moe_config)

498 self.quant_config = quant_config

@@ -713,6 +734,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

734 replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)

735 replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)

736

737 + self._init_moe_kernel(layer)

738 +

739 + def _init_moe_kernel(self, layer: RoutedExperts) -> None:

740 + """Build the MoE kernel from the layer's current (converted) weights."""

741 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

742 assert self.moe_quant_config is not None

743 assert self.experts_cls is not None

@@ -725,6 +750,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

750 )

751

752 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

753 + if is_weights_pre_processed():

754 + # Weights are already in kernel format; rebuild the kernel only.

755 + self._init_moe_kernel(layer)

756 + return

757 +

758 # Allow for accessing weights and scales in standard way.

759 w13 = layer.w13_weight

760 w2 = layer.w2_weight

vllm/model_executor/layers/quantization/kv_cache.py

+17

@@ -8,6 +8,7 @@ from vllm.model_executor.layers.quantization.base_config import (

8 QuantizationConfig,

9 QuantizeMethodBase,

10 )

11 +from vllm.model_executor.utils import is_weights_pre_processed

12 from vllm.platforms import current_platform

13 from vllm.utils.torch_utils import is_quantized_kv_cache

14 from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales

@@ -51,6 +52,8 @@ class BaseKVCacheMethod(QuantizeMethodBase):

52 quant_config: the appropriate QuantizationConfig

53 """

55 + supports_pre_processed_weights = True

56 +

57 def __init__(self, quant_config: QuantizationConfig):

58 self.quant_config = quant_config

@@ -72,6 +75,20 @@ class BaseKVCacheMethod(QuantizeMethodBase):

75 raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")

77 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

78 + if is_weights_pre_processed():

79 + # Scale buffers were exported with post-processed values; drop the

80 + # consumed scale placeholders and rebuild the host-side copies.

81 + for name in ("k_scale", "v_scale", "q_scale", "prob_scale"):

82 + if hasattr(layer, name):

83 + delattr(layer, name)

84 + if hasattr(layer, "_k_scale"):

85 + layer._k_scale_float = layer._k_scale.item()

86 + layer._v_scale_float = layer._v_scale.item()

87 + layer._q_scale_float = layer._q_scale.item()

88 + layer._k_scale_cpu = layer._k_scale.detach().to("cpu")

89 + layer._v_scale_cpu = layer._v_scale.detach().to("cpu")

90 + return

91 +

92 # skip if there are no weights to process (for example, weight reloading)

93 if not hasattr(layer, "q_scale"):

94 assert not hasattr(layer, "k_scale")

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -36,6 +36,8 @@ DEFAULT_VOCAB_PADDING_SIZE = 64

36 class UnquantizedEmbeddingMethod(QuantizeMethodBase):

37 """Unquantized method for embeddings."""

39 + supports_pre_processed_weights = True

40 +

41 def create_weights(

42 self,

43 layer: torch.nn.Module,

vllm/model_executor/model_loader/__init__.py

@@ -24,6 +24,7 @@ from vllm.model_executor.model_loader.utils import (

24 get_model_architecture,

25 get_model_cls,

26 )

27 +from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader

29 logger = init_logger(__name__)

@@ -35,6 +36,7 @@ LoadFormats = Literal[

36 "dummy",

37 "fastsafetensors",

38 "instanttensor",

39 + "ipc_cache",

40 "mistral",

41 "modelexpress",

42 "npcache",

@@ -51,6 +53,7 @@ _LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] = {

53 "dummy": DummyModelLoader,

54 "fastsafetensors": DefaultModelLoader,

55 "instanttensor": DefaultModelLoader,

56 + "ipc_cache": IpcModelLoader,

57 "mistral": DefaultModelLoader,

58 "modelexpress": ModelExpressModelLoader,

59 "npcache": DefaultModelLoader,

@@ -150,6 +153,7 @@ __all__ = [

153 "ModelExpressModelLoader",

154 "DefaultModelLoader",

155 "DummyModelLoader",

156 + "IpcModelLoader",

157 "RunaiModelStreamerLoader",

158 "ShardedStateLoader",

159 "TensorizerLoader",

vllm/model_executor/model_loader/utils.py

+16 -1

@@ -26,6 +26,7 @@ from vllm.model_executor.model_loader.reload import (

26 )

27 from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings

28 from vllm.model_executor.models.interfaces import SupportsQuant

29 +from vllm.model_executor.utils import is_weights_pre_processed

30 from vllm.tracing import instrument

31 from vllm.utils.mem_utils import release_device_memory_under_pressure

32 from vllm.utils.platform_utils import is_pin_memory_available

@@ -97,13 +98,27 @@ def initialize_model(

98 def process_weights_after_loading(

99 model: nn.Module, model_config: ModelConfig, target_device: torch.device

100 ) -> None:

101 + """Post-process loaded weights into runtime format.

102 +

103 + Under ``weights_already_processed`` (weight cache IPC loader), quant

104 + methods skip tensor transforms and must declare

105 + ``supports_pre_processed_weights``, otherwise this raises ``RuntimeError``.

106 + """

107 # Reclaim memory when an explicit lm_head has been

108 # loaded, but it is identical to the input embeddings.

109 maybe_retie_word_embeddings(model, model_config)

110

104 - for _, module in model.named_modules():

111 + for name, module in model.named_modules():

112 quant_method = getattr(module, "quant_method", None)

113 if isinstance(quant_method, QuantizeMethodBase):

114 + if (

115 + is_weights_pre_processed()

116 + and not quant_method.supports_pre_processed_weights

117 + ):

118 + raise RuntimeError(

119 + f"layer {name or '<root>'}: {type(quant_method).__name__} "

120 + "does not support pre-processed weights"

121 + )

122 # When quant methods need to process weights after loading

123 # (for repacking, quantizing, etc), they expect parameters

124 # to be on the global target device. This scope is for the

vllm/model_executor/model_loader/weight_cache/__init__.py new

+24

@@ -0,0 +1,24 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader

5 +from vllm.model_executor.model_loader.weight_cache.protocol import (

6 + CacheConfigMismatchError,

7 + TensorEntry,

8 + UnsupportedQuantForIPCError,

9 + WeightCacheKey,

10 + WeightCacheUnavailableError,

11 + check_ipc_quant_support,

12 + is_ipc_quant_supported,

13 +)

14 +

15 +__all__ = [

16 + "CacheConfigMismatchError",

17 + "IpcModelLoader",

18 + "TensorEntry",

19 + "UnsupportedQuantForIPCError",

20 + "WeightCacheKey",

21 + "WeightCacheUnavailableError",

22 + "check_ipc_quant_support",

23 + "is_ipc_quant_supported",

24 +]

vllm/model_executor/model_loader/weight_cache/daemon.py new

+412

@@ -0,0 +1,412 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Weight cache daemon for fast engine restarts.

4 +

5 +One daemon process per GPU holds the post-quantized, TP-sharded weights of its

6 +rank in GPU memory and serves CUDA IPC handles to vLLM engines over a Unix

7 +domain socket. Restarting engines map the weights via zero-copy IPC instead of

8 +reloading from disk.

9 +

10 +Launch one daemon per TP rank with a single command:

11 +

12 + python -m vllm.model_executor.model_loader.weight_cache.daemon \\

13 + --model /path/to/model --tensor-parallel-size 4

14 +

15 +Engines then load from the daemons with:

16 +

17 + vllm serve /path/to/model --tensor-parallel-size 4 \\

18 + --load-format ipc_cache

19 +

20 +Only tensor parallelism is supported; pipeline, data, and expert parallelism

21 +are rejected at launch.

22 +"""

23 +

24 +import contextlib

25 +import fcntl

26 +import multiprocessing

27 +import os

28 +import queue

29 +import signal

30 +import socket

31 +import sys

32 +from collections.abc import Callable

33 +

34 +import torch

35 +

36 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

37 +from vllm.distributed import (

38 + ensure_model_parallel_initialized,

39 + init_distributed_environment,

40 +)

41 +from vllm.logger import init_logger

42 +from vllm.model_executor.model_loader.weight_cache.protocol import (

43 + TensorEntry,

44 + WeightCacheKey,

45 + WeightCacheUnavailableError,

46 + check_ipc_quant_support,

47 + ensure_private_socket_dir,

48 + get_physical_device_id,

49 + get_socket_path,

50 + recv_msg,

51 + send_msg,

52 + verify_peer_is_owner,

53 +)

54 +from vllm.platforms import current_platform

55 +

56 +logger = init_logger(__name__)

57 +

58 +

59 +def _report_ready(message: str) -> None:

60 + """Write a readiness message straight to the original stderr descriptor.

61 +

62 + Loading the model pulls in FlashInfer's CuTeDSL JIT compiler, which swaps

63 + ``sys.stdout``/``sys.stderr`` for in-memory buffers while compiling kernels

64 + across worker threads. That save/restore races on the interpreter-global

65 + streams and can leave them (and vLLM's logging handler, which caches the

66 + original stream object) detached, silently swallowing everything logged

67 + afterwards -- including the daemon's readiness announcement. Since operators

68 + rely on that line to know the daemon is serving, write it directly to file

69 + descriptor 2 so it bypasses the logging machinery entirely.

70 + """

71 + with contextlib.suppress(OSError):

72 + os.write(2, (message + "\n").encode())

73 +

74 +

75 +def export_entries(

76 + model: torch.nn.Module,

77 +) -> tuple[dict[str, TensorEntry], dict[str, str]]:

78 + """Export a model's tensors, preserving tied-parameter aliases.

79 +

80 + ``named_parameters``/``named_buffers`` are iterated with

81 + ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing

82 + storage with ``embed_tokens.weight``) are not silently dropped. Each unique

83 + tensor is exported once; every additional name that refers to the same

84 + tensor object is recorded in the returned alias map so the client can

85 + re-establish the shared identity instead of allocating uninitialized

86 + memory for it.

87 +

88 + Returns:

89 + A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to

90 + its ``TensorEntry`` and ``aliases`` maps each duplicate name to its

91 + canonical name.

92 + """

93 + entries: dict[str, TensorEntry] = {}

94 + aliases: dict[str, str] = {}

95 + canonical_by_id: dict[int, str] = {}

96 +

97 + def _add(name: str, tensor: torch.Tensor, kind: str) -> None:

98 + canonical = canonical_by_id.get(id(tensor))

99 + if canonical is not None:

100 + aliases[name] = canonical

101 + return

102 + canonical_by_id[id(tensor)] = name

103 + entries[name] = TensorEntry.from_tensor(tensor, kind)

104 +

105 + for name, param in model.named_parameters(remove_duplicate=False):

106 + _add(name, param, "param")

107 + # named_buffers includes non-persistent buffers (e.g. rotary embedding

108 + # caches) that state_dict would miss.

109 + for name, buffer in model.named_buffers(remove_duplicate=False):

110 + if name in entries or name in aliases:

111 + continue

112 + _add(name, buffer, "buffer")

113 + return entries, aliases

114 +

115 +

116 +class WeightCacheDaemon:

117 + """Per-GPU process that loads one TP shard and serves CUDA IPC handles."""

118 +

119 + def __init__(

120 + self,

121 + vllm_config: VllmConfig,

122 + tp_rank: int,

123 + distributed_init_method: str,

124 + socket_dir: str | None = None,

125 + ):

126 + self.vllm_config = vllm_config

127 + self.tp_rank = tp_rank

128 + self.distributed_init_method = distributed_init_method

129 + self.socket_dir = socket_dir

130 + self.model: torch.nn.Module | None = None

131 + self.entries: dict[str, TensorEntry] = {}

132 + self.aliases: dict[str, str] = {}

133 + # Fingerprint before loading: process_weights_after_loading may

134 + # mutate hf_config.quantization_config.

135 + self.cache_config = WeightCacheKey.from_model_config(

136 + vllm_config.model_config,

137 + tp_size=vllm_config.parallel_config.tensor_parallel_size,

138 + tp_rank=tp_rank,

139 + )

140 +

141 + def load_model(self) -> None:

142 + from vllm.model_executor.model_loader import get_model

143 +

144 + tp_size = self.cache_config.tp_size

145 + torch.accelerator.set_device_index(self.tp_rank)

146 + init_distributed_environment(

147 + world_size=tp_size,

148 + rank=self.tp_rank,

149 + distributed_init_method=self.distributed_init_method,

150 + local_rank=self.tp_rank,

151 + backend=current_platform.dist_backend,

152 + )

153 + with set_current_vllm_config(self.vllm_config):

154 + ensure_model_parallel_initialized(tp_size, 1)

155 + self.model = get_model(vllm_config=self.vllm_config)

156 + self._export_entries()

157 + logger.info(

158 + "Weight cache daemon rank %d cached %d tensors",

159 + self.tp_rank,

160 + len(self.entries),

161 + )

162 +

163 + def _export_entries(self) -> None:

164 + assert self.model is not None

165 + self.entries, self.aliases = export_entries(self.model)

166 +

167 + def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None:

168 + """Serve requests until terminated.

169 +

170 + The socket is only bound once the model is fully cached, so clients

171 + get a connection error (and fall back to disk) until the daemon is

172 + ready.

173 +

174 + Args:

175 + ready_callback: Invoked once the socket is bound and listening,

176 + so the launcher can report overall readiness.

177 + """

178 + socket_path = self._socket_path()

179 + ensure_private_socket_dir(

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

181 + )

182 + # Hold an exclusive per-GPU lock for the daemon's lifetime so a second

183 + # daemon cannot remove this daemon's live socket and hijack the path.

184 + lock_fd = self._acquire_gpu_lock(socket_path)

185 + if os.path.exists(socket_path):

186 + os.unlink(socket_path)

187 + server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

188 + server.bind(socket_path)

189 + os.chmod(socket_path, 0o600)

190 + server.listen()

191 + logger.info(

192 + "Weight cache daemon rank %d serving on %s", self.tp_rank, socket_path

193 + )

194 + print(

195 + f"Weight cache daemon rank {self.tp_rank} ready: serving on {socket_path}"

196 + )

197 + if ready_callback is not None:

198 + ready_callback()

199 + try:

200 + while True:

201 + conn, _ = server.accept()

202 + with conn:

203 + try:

204 + verify_peer_is_owner(conn)

205 + self._handle_connection(conn)

206 + except (ConnectionError, EOFError):

207 + logger.warning("Client disconnected mid-request")

208 + except Exception:

209 + # A single malformed or malicious request must not take

210 + # down the daemon for every other engine on this GPU.

211 + logger.exception(

212 + "Error handling weight cache client; continuing"

213 + )

214 + finally:

215 + server.close()

216 + if os.path.exists(socket_path):

217 + os.unlink(socket_path)

218 + os.close(lock_fd)

219 +

220 + def _acquire_gpu_lock(self, socket_path: str) -> int:

221 + """Take an exclusive lock guarding this GPU's socket path.

222 +

223 + The lock is advisory and released automatically when the daemon exits

224 + (or crashes), so a stale socket is only ever removed by whoever owns

225 + the lock. A running daemon holding it makes a second daemon fail fast

226 + instead of clobbering the live socket.

227 + """

228 + lock_fd = os.open(f"{socket_path}.lock", os.O_CREAT | os.O_RDWR, 0o600)

229 + try:

230 + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)

231 + except OSError as e:

232 + os.close(lock_fd)

233 + raise WeightCacheUnavailableError(

234 + f"Another weight cache daemon already owns {socket_path}"

235 + ) from e

236 + return lock_fd

237 +

238 + def _socket_path(self) -> str:

239 + device_index = torch.accelerator.current_device_index()

240 + gpu_id = get_physical_device_id(device_index)

241 + if gpu_id is None:

242 + gpu_id = device_index

243 + return get_socket_path(gpu_id, self.socket_dir)

244 +

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

246 + request = recv_msg(conn)

247 + cmd = request.get("cmd")

248 + if cmd == "get_state":

249 + self._handle_get_state(conn, request)

250 + elif cmd == "release":

251 + self._handle_release(conn)

252 + else:

253 + send_msg(conn, {"status": "error", "message": f"Unknown command {cmd!r}"})

254 +

255 + def _handle_get_state(self, conn: socket.socket, request: dict) -> None:

256 + client_config = request.get("cache_config")

257 + if not isinstance(client_config, WeightCacheKey):

258 + send_msg(conn, {"status": "error", "message": "Missing cache_config"})

259 + return

260 + mismatched = self.cache_config.mismatched_fields(client_config)

261 + if mismatched:

262 + logger.warning("WeightCacheKey mismatch on fields: %s", mismatched)

263 + send_msg(conn, {"status": "mismatch", "fields": mismatched})

264 + return

265 + if not self.entries:

266 + send_msg(conn, {"status": "error", "message": "Weights were released"})

267 + return

268 + send_msg(

269 + conn,

270 + {

271 + "status": "ok",

272 + "entries": self.entries,

273 + "aliases": self.aliases,

274 + "gpu_uuid": self._gpu_uuid(),

275 + },

276 + )

277 +

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

279 + self.entries.clear()

280 + self.aliases.clear()

281 + self.model = None

282 + torch.accelerator.empty_cache()

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

284 + send_msg(conn, {"status": "ok"})

285 +

286 + def _gpu_uuid(self) -> str:

287 + props = torch.cuda.get_device_properties(

288 + torch.accelerator.current_device_index()

289 + )

290 + return str(props.uuid)

291 +

292 +

293 +def _run_daemon(

294 + tp_rank: int,

295 + vllm_config: VllmConfig,

296 + distributed_init_method: str,

297 + socket_dir: str | None,

298 + ready_queue: "multiprocessing.Queue[int]",

299 +) -> None:

300 + daemon = WeightCacheDaemon(

301 + vllm_config, tp_rank, distributed_init_method, socket_dir

302 + )

303 + daemon.load_model()

304 + daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank))

305 +

306 +

307 +def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:

308 + """Reject every parallelism mode other than tensor parallelism."""

309 + unsupported = {

310 + "pipeline parallelism": parallel_config.pipeline_parallel_size > 1,

311 + "data parallelism": parallel_config.data_parallel_size > 1,

312 + "expert parallelism": parallel_config.enable_expert_parallel,

313 + }

314 + for name, enabled in unsupported.items():

315 + if enabled:

316 + raise ValueError(

317 + f"The weight cache daemon only supports tensor parallelism; "

318 + f"{name} is not supported"

319 + )

320 +

321 +

322 +def main() -> None:

323 + from vllm.engine.arg_utils import EngineArgs

324 + from vllm.utils.argparse_utils import FlexibleArgumentParser

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

326 +

327 + parser = FlexibleArgumentParser(

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

329 + )

330 + EngineArgs.add_cli_args(parser)

331 + parser.add_argument(

332 + "--weight-cache-socket-dir",

333 + type=str,

334 + default=None,

335 + help="Directory for the daemon Unix sockets (default: tempdir).",

336 + )

337 + args = parser.parse_args()

338 + engine_args = EngineArgs.from_cli_args(args)

339 + vllm_config = engine_args.create_engine_config()

340 + if vllm_config.load_config.load_format == "ipc_cache":

341 + raise ValueError(

342 + "The weight cache daemon itself must load from disk; use the "

343 + "default --load-format"

344 + )

345 + # Checked before loading anything: an unsupported quantization method would

346 + # otherwise only surface in the engine, after a full load.

347 + check_ipc_quant_support(vllm_config.model_config, where="daemon")

348 + parallel_config = vllm_config.parallel_config

349 + _reject_unsupported_parallelism(parallel_config)

350 + tp_size = parallel_config.tensor_parallel_size

351 +

352 + distributed_init_method = get_distributed_init_method("127.0.0.1", get_open_port())

353 + ctx = multiprocessing.get_context("spawn")

354 + ready_queue: multiprocessing.Queue[int] = ctx.Queue()

355 + procs = [

356 + ctx.Process(

357 + target=_run_daemon,

358 + args=(

359 + rank,

360 + vllm_config,

361 + distributed_init_method,

362 + args.weight_cache_socket_dir,

363 + ready_queue,

364 + ),

365 + name=f"vllm-weight-cache-daemon-{rank}",

366 + )

367 + for rank in range(tp_size)

368 + ]

369 + for proc in procs:

370 + proc.start()

371 +

372 + def _shutdown(signum, frame):

373 + for proc in procs:

374 + proc.terminate()

375 +

376 + signal.signal(signal.SIGINT, _shutdown)

377 + signal.signal(signal.SIGTERM, _shutdown)

378 +

379 + ready_ranks: set[int] = set()

380 + while len(ready_ranks) < tp_size:

381 + try:

382 + ready_ranks.add(ready_queue.get(timeout=1.0))

383 + except queue.Empty:

384 + dead = [p for p in procs if p.exitcode is not None]

385 + if dead:

386 + logger.error(

387 + "Weight cache daemon rank(s) exited during startup "

388 + "(exitcodes=%s); shutting down.",

389 + [p.exitcode for p in dead],

390 + )

391 + for proc in procs:

392 + proc.terminate()

393 + for proc in procs:

394 + proc.join()

395 + sys.exit(max((p.exitcode or 0) for p in procs))

396 + logger.info_once(

397 + "===== Weight cache daemon READY: all %d ranks serving in %s =====",

398 + tp_size,

399 + args.weight_cache_socket_dir or "the default socket dir",

400 + )

401 + _report_ready(

402 + f"===== Weight cache daemon READY: all {tp_size} rank(s) serving in "

403 + f"{args.weight_cache_socket_dir or 'the default socket dir'} ====="

404 + )

405 +

406 + for proc in procs:

407 + proc.join()

408 + sys.exit(max(proc.exitcode or 0 for proc in procs))

409 +

410 +

411 +if __name__ == "__main__":

412 + main()

vllm/model_executor/model_loader/weight_cache/ipc_loader.py new

+394

@@ -0,0 +1,394 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""IPC model loader: maps post-quantized weights from a local weight cache

4 +daemon via CUDA IPC instead of loading from disk."""

5 +

6 +import dataclasses

7 +import socket

8 +from copy import copy

9 +

10 +import torch

11 +import torch.nn as nn

12 +

13 +from vllm.config import ModelConfig, VllmConfig

14 +from vllm.config.load import LoadConfig

15 +from vllm.logger import init_logger

16 +from vllm.model_executor.model_loader.base_loader import BaseModelLoader

17 +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader

18 +from vllm.model_executor.model_loader.utils import (

19 + initialize_model,

20 + process_weights_after_loading,

21 +)

22 +from vllm.model_executor.model_loader.weight_cache.protocol import (

23 + CacheConfigMismatchError,

24 + TensorEntry,

25 + UnsupportedQuantForIPCError,

26 + WeightCacheKey,

27 + WeightCacheUnavailableError,

28 + check_ipc_quant_support,

29 + get_physical_device_id,

30 + get_socket_path,

31 + recv_msg,

32 + send_msg,

33 + verify_socket_owner,

34 +)

35 +from vllm.model_executor.utils import weights_already_processed

36 +from vllm.tracing import instrument

37 +from vllm.utils.torch_utils import set_default_torch_dtype

38 +

39 +logger = init_logger(__name__)

40 +

41 +_CONNECT_TIMEOUT_S = 5.0

42 +_STATE_TIMEOUT_S = 300.0

43 +

44 +

45 +class IpcModelLoader(BaseModelLoader):

46 + """Loads a model by mapping the weight cache daemon's tensors via CUDA IPC.

47 +

48 + The model is initialized on the meta device and every parameter/buffer is

49 + replaced by the daemon's post-quantized tensor, so

50 + process_weights_after_loading is skipped entirely. In "zero_copy" mode the

51 + engine shares the daemon's GPU memory; in "copy" mode the tensors are

52 + cloned into engine-owned memory and the daemon is asked to release its

53 + cache afterwards.

54 +

55 + Extra config keys (via --model-loader-extra-config):

56 +

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

58 + derived from the physical GPU id.

59 + - socket_dir: directory containing the daemon sockets.

60 + - mode: "zero_copy" (default) or "copy".

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

62 + the fingerprints mismatch (default: True).

63 + - connect_timeout_s: socket connect timeout (default: 5.0).

64 + - state_timeout_s: timeout for the weight-transfer request (default: 300.0).

65 +

66 + Note: in zero-copy mode the weights live in the daemon's CUDA IPC

67 + allocations, so sleep mode (CuMemAllocator weight offloading) must not be

68 + used with this loader.

69 + """

70 +

71 + def __init__(self, load_config: LoadConfig):

72 + super().__init__(load_config)

73 + extra_config = copy(load_config.model_loader_extra_config or {})

74 + self.socket_path: str | None = extra_config.pop("socket_path", None)

75 + self.socket_dir: str | None = extra_config.pop("socket_dir", None)

76 + self.mode: str = extra_config.pop("mode", "zero_copy")

77 + self.fallback: bool = extra_config.pop("fallback", True)

78 + self.connect_timeout_s: float = float(

79 + extra_config.pop("connect_timeout_s", _CONNECT_TIMEOUT_S)

80 + )

81 + self.state_timeout_s: float = float(

82 + extra_config.pop("state_timeout_s", _STATE_TIMEOUT_S)

83 + )

84 + if self.mode not in ("zero_copy", "copy"):

85 + raise ValueError(

86 + f"Invalid weight cache mode {self.mode!r}, "

87 + "expected 'zero_copy' or 'copy'"

88 + )

89 + if extra_config:

90 + raise ValueError(

91 + f"Unexpected extra config keys for load format "

92 + f"{load_config.load_format}: {sorted(extra_config)}"

93 + )

94 +

95 + def download_model(self, model_config: ModelConfig) -> None:

96 + DefaultModelLoader(self._fallback_load_config()).download_model(model_config)

97 +

98 + def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:

99 + """Best-effort in-place reload for an already-initialized model.

100 +

101 + Copies daemon tensors into matching parameters/buffers. The model is

102 + expected to already be in the post-quantized layout (e.g. previously

103 + loaded through this loader).

104 + """

105 + device_index = torch.accelerator.current_device_index()

106 + entries, _ = self._fetch_entries(model_config)

107 + params = dict(model.named_parameters())

108 + buffers = dict(model.named_buffers())

109 + for name, entry in entries.items():

110 + target = params.get(name, buffers.get(name))

111 + source = entry.rebuild(device_index)

112 + if target is None or target.shape != source.shape:

113 + logger.warning("Skipping mismatched cached tensor %s", name)

114 + continue

115 + target.data.copy_(source)

116 +

117 + @instrument(span_name="Load model")

118 + def load_model(

119 + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""

120 + ) -> nn.Module:

121 + # Unsupported quantization is a permanent misconfiguration rather than a

122 + # transient daemon outage, so it is raised even when fallback is on.

123 + self._check_supported(vllm_config, model_config)

124 + state_fetched = False

125 + try:

126 + entries, aliases = self._fetch_entries(model_config)

127 + state_fetched = True

128 + return self._build_model(

129 + vllm_config, model_config, prefix, entries, aliases

130 + )

131 + except (WeightCacheUnavailableError, CacheConfigMismatchError) as e:

132 + if not self.fallback:

133 + raise

134 + logger.warning(

135 + "Weight cache unusable (%s); falling back to disk loading", e

136 + )

137 + except Exception:

138 + if not self.fallback:

139 + raise

140 + logger.exception(

141 + "Weight cache IPC loading failed; falling back to disk loading"

142 + )

143 + # _build_model failed after fetching state without reaching its

144 + # copy-mode release, so the daemon still holds the full cache;

145 + # release it so the disk fallback does not OOM against it.

146 + if state_fetched and self.mode == "copy":

147 + self._send_release()

148 + torch.accelerator.empty_cache()

149 + return self._fallback_load(vllm_config, model_config, prefix)

150 +

151 + def _build_model(

152 + self,

153 + vllm_config: VllmConfig,

154 + model_config: ModelConfig,

155 + prefix: str,

156 + entries: dict[str, TensorEntry],

157 + aliases: dict[str, str],

158 + ) -> nn.Module:

159 + device_config = vllm_config.device_config

160 + load_device = (

161 + device_config.device

162 + if self.load_config.device is None

163 + else self.load_config.device

164 + )

165 + target_device = torch.device(load_device)

166 + device_index = (

167 + target_device.index

168 + if target_device.index is not None

169 + else torch.accelerator.current_device_index()

170 + )

171 + with set_default_torch_dtype(model_config.dtype):

172 + with torch.device("meta"):

173 + model = initialize_model(

174 + vllm_config=vllm_config,

175 + model_config=model_config,

176 + prefix=prefix,

177 + )

178 + self._apply_entries(model, entries, aliases, device_index)

179 + # The daemon exports tensors that already went through

180 + # process_weights_after_loading; re-run it in pre-processed mode

181 + # so quant methods only rebuild Python-side state (e.g. the MoE

182 + # kernel). Leftovers are materialized afterwards so that

183 + # placeholders the daemon-side post-processing consumed are

184 + # dropped rather than filled with uninitialized memory.

185 + with weights_already_processed():

186 + process_weights_after_loading(model, model_config, target_device)

187 + _materialize_remaining_meta_tensors(

188 + model, torch.device(target_device.type, device_index)

189 + )

190 + if self.mode == "copy":

191 + self._send_release()

192 + logger.info(

193 + "Mapped %d tensors from the weight cache daemon (%s mode)",

194 + len(entries),

195 + self.mode,

196 + )

197 + return model.eval()

198 +

199 + @staticmethod

200 + def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None:

201 + check_ipc_quant_support(model_config, where="engine")

202 + cache_dtype = vllm_config.cache_config.cache_dtype

203 + if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"):

204 + # BaseKVCacheMethod.process_weights_after_loading turns the loaded

205 + # k/v scale parameters into plain float attributes. For fp8 cache

206 + # dtypes those are rebuilt from the exported scale buffers when

207 + # process_weights_after_loading runs in pre-processed mode; other

208 + # quantized cache dtypes are not verified.

209 + raise UnsupportedQuantForIPCError(

210 + f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not "

211 + "supported by the weight cache; use --kv-cache-dtype auto."

212 + )

213 +

214 + def _apply_entries(

215 + self,

216 + model: nn.Module,

217 + entries: dict[str, TensorEntry],

218 + aliases: dict[str, str],

219 + device_index: int,

220 + ) -> None:

221 + # remove_duplicate=False keeps tied module aliases reachable by name:

222 + # a tied lm_head *is* the embedding module, so the deduplicated view

223 + # would not contain "lm_head" at all.

224 + modules = dict(model.named_modules(remove_duplicate=False))

225 + registered: dict[str, torch.Tensor] = {}

226 +

227 + def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None:

228 + module_name, _, leaf = name.rpartition(".")

229 + module = modules.get(module_name)

230 + if module is None:

231 + raise RuntimeError(f"Cached tensor {name} has no matching module")

232 + # Replace via registration rather than param.data assignment,

233 + # which fails for meta tensors. Entries may also introduce

234 + # post-quantization tensors absent from the meta model.

235 + module._parameters.pop(leaf, None)

236 + module._buffers.pop(leaf, None)

237 + if is_param:

238 + obj: torch.Tensor = (

239 + tensor

240 + if isinstance(tensor, nn.Parameter)

241 + else nn.Parameter(tensor, requires_grad=False)

242 + )

243 + module.register_parameter(leaf, obj)

244 + else:

245 + obj = tensor

246 + module.register_buffer(leaf, obj)

247 + registered[name] = obj

248 +

249 + for name, entry in entries.items():

250 + tensor = entry.rebuild(device_index)

251 + if self.mode == "copy":

252 + tensor = tensor.clone()

253 + _register(name, tensor, entry.kind == "param")

254 +

255 + # Re-establish tied-weight aliases by registering the *same* object the

256 + # canonical name resolved to, so parameter identity (and the tie) is

257 + # preserved instead of allocating uninitialized memory.

258 + for alias_name, canonical_name in aliases.items():

259 + obj = registered.get(canonical_name)

260 + if obj is None:

261 + logger.warning(

262 + "Cached alias %s references missing canonical tensor %s",

263 + alias_name,

264 + canonical_name,

265 + )

266 + continue

267 + _register(alias_name, obj, isinstance(obj, nn.Parameter))

268 +

269 + def _fetch_entries(

270 + self, model_config: ModelConfig

271 + ) -> tuple[dict[str, TensorEntry], dict[str, str]]:

272 + from vllm.distributed import (

273 + get_tensor_model_parallel_rank,

274 + get_tensor_model_parallel_world_size,

275 + )

276 +

277 + cache_config = WeightCacheKey.from_model_config(

278 + model_config,

279 + tp_size=get_tensor_model_parallel_world_size(),

280 + tp_rank=get_tensor_model_parallel_rank(),

281 + )

282 + return self._request_state(cache_config)

283 +

284 + def _request_state(

285 + self, cache_config: WeightCacheKey

286 + ) -> tuple[dict[str, TensorEntry], dict[str, str]]:

287 + with self._connect(self.state_timeout_s) as conn:

288 + send_msg(conn, {"cmd": "get_state", "cache_config": cache_config})

289 + response = recv_msg(conn)

290 + status = response.get("status")

291 + if status == "mismatch":

292 + raise CacheConfigMismatchError(

293 + f"WeightCacheKey mismatch on fields: {response.get('fields')}"

294 + )

295 + if status != "ok":

296 + raise WeightCacheUnavailableError(

297 + f"Weight cache daemon error: {response.get('message')}"

298 + )

299 + self._check_gpu_uuid(response.get("gpu_uuid"))

300 + return response["entries"], response.get("aliases", {})

301 +

302 + def _connect(self, timeout: float) -> socket.socket:

303 + socket_path = self._resolve_socket_path()

304 + # The auto-derived per-user directory is locked to 0700 and checked

305 + # strictly. When the operator explicitly configures a path they own the

306 + # trust decision, so only ownership/symlink safety is enforced.

307 + strict_perms = self.socket_path is None and self.socket_dir is None

308 + try:

309 + verify_socket_owner(socket_path, strict_perms=strict_perms)

310 + except OSError as e:

311 + raise WeightCacheUnavailableError(

312 + f"Weight cache socket {socket_path} is unavailable: {e}"

313 + ) from e

314 + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

315 + sock.settimeout(timeout)

316 + try:

317 + sock.connect(socket_path)

318 + except OSError as e:

319 + sock.close()

320 + raise WeightCacheUnavailableError(

321 + f"Cannot connect to weight cache daemon at {socket_path}: {e}"

322 + ) from e

323 + return sock

324 +

325 + def _resolve_socket_path(self) -> str:

326 + if self.socket_path is not None:

327 + return self.socket_path

328 + device_index = torch.accelerator.current_device_index()

329 + gpu_id = get_physical_device_id(device_index)

330 + if gpu_id is None:

331 + raise WeightCacheUnavailableError(

332 + "Cannot infer the physical GPU id from CUDA_VISIBLE_DEVICES; "

333 + "pass socket_path via --model-loader-extra-config"

334 + )

335 + return get_socket_path(gpu_id, self.socket_dir)

336 +

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

338 + if daemon_uuid is None:

339 + return

340 + props = torch.cuda.get_device_properties(

341 + torch.accelerator.current_device_index()

342 + )

343 + local_uuid = str(props.uuid)

344 + if daemon_uuid != local_uuid:

345 + raise CacheConfigMismatchError(

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

347 + "check the socket path / GPU mapping"

348 + )

349 +

350 + def _send_release(self) -> None:

351 + try:

352 + with self._connect(self.connect_timeout_s) as conn:

353 + send_msg(conn, {"cmd": "release"})

354 + recv_msg(conn)

355 + except (WeightCacheUnavailableError, ConnectionError, OSError):

356 + logger.warning("Failed to ask the weight cache daemon to release")

357 +

358 + def _fallback_load_config(self) -> LoadConfig:

359 + # DefaultModelLoader must not see load_format="ipc_cache" or the ipc

360 + # extra config keys.

361 + return dataclasses.replace(

362 + self.load_config, load_format="auto", model_loader_extra_config={}

363 + )

364 +

365 + def _fallback_load(

366 + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str

367 + ) -> nn.Module:

368 + loader = DefaultModelLoader(self._fallback_load_config())

369 + return loader.load_model(

370 + vllm_config=vllm_config, model_config=model_config, prefix=prefix

371 + )

372 +

373 +

374 +def _materialize_remaining_meta_tensors(model: nn.Module, device: torch.device) -> None:

375 + """Allocate any tensors the daemon did not provide.

376 +

377 + These are typically parameters removed on the daemon side by

378 + process_weights_after_loading; they are not expected to be read at

379 + runtime, so they are left uninitialized.

380 + """

381 + for module_name, module in model.named_modules():

382 + for leaf, param in list(module._parameters.items()):

383 + if param is not None and param.device.type == "meta":

384 + logger.warning(

385 + "Materializing empty parameter %s.%s missing from the weight cache",

386 + module_name,

387 + leaf,

388 + )

389 + module._parameters[leaf] = nn.Parameter(

390 + torch.empty_like(param, device=device), requires_grad=False

391 + )

392 + for leaf, buffer in list(module._buffers.items()):

393 + if buffer is not None and buffer.device.type == "meta":

394 + module._buffers[leaf] = torch.empty_like(buffer, device=device)

vllm/model_executor/model_loader/weight_cache/protocol.py new

+412

@@ -0,0 +1,412 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""WeightCacheKey fingerprinting and socket protocol for the weight cache daemon.

4 +

5 +The protocol uses pickle over a Unix domain socket and is only intended for

6 +communication between trusted local processes owned by the same user. The

7 +sockets live in a per-user private directory (mode 0700) and the daemon

8 +restricts the socket file permissions to the owner (0600). Both the daemon and

9 +the engine verify that the directory and socket are owned by the current user

10 +and are not group/world accessible before trusting them, so a different local

11 +user cannot pre-plant a malicious socket at a predictable path.

12 +"""

13 +

14 +import glob

15 +import json

16 +import os

17 +import pickle

18 +import socket

19 +import stat

20 +import struct

21 +import tempfile

22 +from dataclasses import dataclass, fields

23 +from typing import Any

24 +

25 +import torch

26 +

27 +import vllm.version

28 +from vllm.config import ModelConfig

29 +from vllm.utils.hashing import safe_hash

30 +

31 +SOCKET_NAME_TEMPLATE = "vllm_weight_cache_gpu{gpu_id}.sock"

32 +SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}"

33 +

34 +_LEN_STRUCT = struct.Struct("!Q")

35 +# Sanity bound for a single message. IPC handles are tiny; only small

36 +# non-CUDA tensors are ever shipped by value.

37 +MAX_MSG_SIZE = 1 << 34

38 +

39 +

40 +def _current_uid() -> int:

41 + getuid = getattr(os, "getuid", None)

42 + return getuid() if getuid is not None else -1

43 +

44 +

45 +class WeightCacheUnavailableError(Exception):

46 + """Raised when no weight cache daemon is reachable or usable."""

47 +

48 +

49 +class CacheConfigMismatchError(Exception):

50 + """Raised when the daemon's cached weights don't match the engine."""

51 +

52 +

53 +class UnsupportedQuantForIPCError(Exception):

54 + """Raised when a quantization method is not verified for IPC weight sharing."""

55 +

56 +

57 +# The daemon exports tensor data only, so sharing is correct just for methods

58 +# whose post-load effect is either fully captured by that data or rebuilt when

59 +# process_weights_after_loading runs under weights_already_processed. Anything

60 +# absent from this registry hard-errors: a method that silently repacks into

61 +# shapes the client cannot reproduce or stamps unrebuildable Python-side state

62 +# would serve wrong numerics. Extend it only after an end-to-end check against

63 +# a disk-loaded baseline.

64 +#

65 +# This check is config-only so the daemon fails fast before a slow full load,

66 +# rather than surfacing the incompatibility only after the engine has built the

67 +# model (where process_weights_after_loading's supports_pre_processed_weights

68 +# guard would eventually catch it).

69 +

70 +

71 +def _quant_config_field(quant_config: Any, key: str) -> Any:

72 + if quant_config is None:

73 + return None

74 + if isinstance(quant_config, dict):

75 + return quant_config.get(key)

76 + return getattr(quant_config, key, None)

77 +

78 +

79 +def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:

80 + """Only block-wise FP8 is verified.

81 +

82 + Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes

83 + ``layer.weight`` during post-processing -- a shape the client's

84 + meta-initialized model cannot reproduce.

85 + """

86 + return _quant_config_field(quant_config, "weight_block_size") is not None

87 +

88 +

89 +# quantization name -> predicate(quant_config) -> True when verified safe.

90 +IPC_QUANT_ALLOWLIST: dict[str | None, Any] = {

91 + None: lambda _quant_config: True, # unquantized

92 + "fp8": _fp8_round_trips_via_ipc,

93 +}

94 +

95 +

96 +def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool:

97 + predicate = IPC_QUANT_ALLOWLIST.get(quantization)

98 + return False if predicate is None else bool(predicate(quant_config))

99 +

100 +

101 +def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None:

102 + """Hard-error unless the model's quantization is verified for IPC sharing.

103 +

104 + Args:

105 + model_config: Model configuration to inspect.

106 + where: Short tag ("daemon"/"engine") used in the error message.

107 +

108 + Raises:

109 + UnsupportedQuantForIPCError: If the quantization method is not on the

110 + verified allowlist.

111 + """

112 + quantization = model_config.quantization

113 + # Prefer the canonical, nested-aware config (multimodal models keep it under

114 + # text_config); fall back to the raw hf_config for older code paths.

115 + quant_config = getattr(

116 + getattr(model_config, "model_arch_config", None), "quantization_config", None

117 + )

118 + if quant_config is None:

119 + quant_config = getattr(model_config.hf_config, "quantization_config", None)

120 + if is_ipc_quant_supported(quantization, quant_config):

121 + return

122 + verified = ", ".join(

123 + "unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST

124 + )

125 + raise UnsupportedQuantForIPCError(

126 + f"[weight_cache:{where}] quantization {quantization!r} is not verified "

127 + f"for CUDA IPC weight sharing: its post-load processing may repack "

128 + f"weights into shapes the client cannot reproduce or stamp Python-side "

129 + f"state that tensor export cannot carry, which would silently serve "

130 + f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size "

131 + f"set, i.e. block-wise). Use the default --load-format for this model."

132 + )

133 +

134 +

135 +def get_physical_device_id(device_index: int) -> int | None:

136 + """Map a local CUDA device index to the physical GPU id.

137 +

138 + Returns None if CUDA_VISIBLE_DEVICES contains non-integer entries

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

140 + """

141 + visible = os.environ.get("CUDA_VISIBLE_DEVICES")

142 + if not visible:

143 + return device_index

144 + entries = visible.split(",")

145 + try:

146 + return int(entries[device_index])

147 + except (IndexError, ValueError):

148 + return None

149 +

150 +

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

152 + """Return the directory that holds the daemon sockets.

153 +

154 + When no explicit directory is given, use a per-user private directory

155 + under the system temp dir so its path is unpredictable to other users and

156 + can be locked down to mode 0700.

157 + """

158 + if socket_dir is not None:

159 + return socket_dir

160 + return os.path.join(

161 + tempfile.gettempdir(), SOCKET_DIR_TEMPLATE.format(uid=_current_uid())

162 + )

163 +

164 +

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

166 + directory = get_socket_dir(socket_dir)

167 + return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_id=gpu_id))

168 +

169 +

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

171 + """Create the socket directory (if needed) locked down to the owner.

172 +

173 + Called by the daemon before binding. Existing directories are re-checked

174 + and, for the auto-derived path, tightened so a pre-existing world-writable

175 + directory is rejected.

176 + """

177 + os.makedirs(directory, mode=0o700, exist_ok=True)

178 + if strict_perms:

179 + os.chmod(directory, 0o700)

180 + verify_private_dir(directory, strict_perms=strict_perms)

181 +

182 +

183 +def verify_private_dir(directory: str, strict_perms: bool = True) -> None:

184 + """Verify a directory is a real dir owned by us and not world/group readable.

185 +

186 + When ``strict_perms`` is False the group/world permission bits are not

187 + checked; this is used for directories the operator explicitly configured

188 + (they own the trust decision), while the auto-derived per-user directory is

189 + always checked strictly.

190 + """

191 + info = os.lstat(directory)

192 + if stat.S_ISLNK(info.st_mode):

193 + raise WeightCacheUnavailableError(

194 + f"Refusing to use symlinked socket directory {directory}"

195 + )

196 + if not stat.S_ISDIR(info.st_mode):

197 + raise WeightCacheUnavailableError(f"{directory} is not a directory")

198 + uid = _current_uid()

199 + if uid != -1 and info.st_uid != uid:

200 + raise WeightCacheUnavailableError(

201 + f"Socket directory {directory} is not owned by the current user"

202 + )

203 + if strict_perms and info.st_mode & 0o077:

204 + raise WeightCacheUnavailableError(

205 + f"Socket directory {directory} is group/world accessible"

206 + )

207 +

208 +

209 +def verify_socket_owner(socket_path: str, strict_perms: bool = True) -> None:

210 + """Verify the socket lives in a private dir and is owned by the current user.

211 +

212 + Called by the engine before connecting so it never talks to a socket a

213 + different user could have planted.

214 + """

215 + verify_private_dir(os.path.dirname(socket_path), strict_perms=strict_perms)

216 + info = os.lstat(socket_path)

217 + if stat.S_ISLNK(info.st_mode):

218 + raise WeightCacheUnavailableError(

219 + f"Refusing to connect to symlinked socket {socket_path}"

220 + )

221 + uid = _current_uid()

222 + if uid != -1 and info.st_uid != uid:

223 + raise WeightCacheUnavailableError(

224 + f"Socket {socket_path} is not owned by the current user"

225 + )

226 +

227 +

228 +def verify_peer_is_owner(conn: socket.socket) -> None:

229 + """Best-effort check that the connecting peer runs as the current user.

230 +

231 + Uses SO_PEERCRED where available (Linux). Silently returns on platforms

232 + that do not expose peer credentials.

233 + """

234 + so_peercred = getattr(socket, "SO_PEERCRED", None)

235 + if so_peercred is None:

236 + return

237 + try:

238 + creds = conn.getsockopt(socket.SOL_SOCKET, so_peercred, struct.calcsize("3i"))

239 + _, peer_uid, _ = struct.unpack("3i", creds)

240 + except OSError:

241 + return

242 + uid = _current_uid()

243 + if uid != -1 and peer_uid != uid:

244 + raise PermissionError(f"Rejecting weight cache connection from uid {peer_uid}")

245 +

246 +

247 +def _hash_quant_config(quant_config: Any) -> str:

248 + if quant_config is None:

249 + return ""

250 + if hasattr(quant_config, "to_dict"):

251 + quant_config = quant_config.to_dict()

252 + payload = json.dumps(quant_config, sort_keys=True, default=str)

253 + return safe_hash(payload.encode(), usedforsecurity=False).hexdigest()

254 +

255 +

256 +def _safetensors_header(path: str) -> bytes:

257 + """Return the raw safetensors header (length prefix + JSON) of a file.

258 +

259 + The header carries tensor names, dtypes, shapes and byte offsets, so it is

260 + a content fingerprint of the shard without reading any weight bytes.

261 + """

262 + with open(path, "rb") as f:

263 + size_bytes = f.read(8)

264 + (header_len,) = struct.unpack("<Q", size_bytes)

265 + return size_bytes + f.read(header_len)

266 +

267 +

268 +def hash_checkpoint(model: str) -> str | None:

269 + """Fingerprint checkpoint content from local safetensors metadata.

270 +

271 + Hashes each shard's safetensors header so a daemon and an engine pointing

272 + at identical weights in different directories produce the same key. Returns

273 + None when local safetensors files can't be located (e.g. an undownloaded

274 + Hugging Face repo id), leaving the caller to fall back to the model path.

275 + """

276 + if not os.path.isdir(model):

277 + return None

278 + from transformers.utils import SAFE_WEIGHTS_INDEX_NAME

279 +

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

281 + filter_duplicate_safetensors_files,

282 + )

283 +

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

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

286 + files = filter_duplicate_safetensors_files(

287 + files, model, SAFE_WEIGHTS_INDEX_NAME

288 + )

289 + if not files:

290 + return None

291 + hasher = safe_hash(b"", usedforsecurity=False)

292 + for path in sorted(files, key=os.path.basename):

293 + hasher.update(os.path.basename(path).encode())

294 + hasher.update(_safetensors_header(path))

295 + return hasher.hexdigest()

296 +

297 +

298 +@dataclass(frozen=True)

299 +class WeightCacheKey:

300 + """Fingerprint of the cached weights.

301 +

302 + Any mismatch between the daemon's and the engine's fingerprint means the

303 + cached weights cannot be reused and the engine must load from disk.

304 + """

305 +

306 + checkpoint: str

307 + model_arch: str

308 + tp_size: int

309 + tp_rank: int

310 + dtype: str

311 + quantization: str | None

312 + quant_config_hash: str

313 + revision: str | None

314 + vllm_version: str

315 +

316 + @classmethod

317 + def from_model_config(

318 + cls, model_config: ModelConfig, tp_size: int, tp_rank: int

319 + ) -> "WeightCacheKey":

320 + """Build the fingerprint for a model configuration.

321 +

322 + Must be called before weight loading: process_weights_after_loading

323 + may mutate hf_config.quantization_config, which would change the hash

324 + between the daemon and the engine.

325 +

326 + The checkpoint is identified by a hash of its safetensors metadata when

327 + the weights are available locally, so a daemon and engine referencing

328 + identical weights in different directories still match; otherwise it

329 + falls back to the model path.

330 + """

331 + hf_config = model_config.hf_config

332 + arch = ",".join(getattr(hf_config, "architectures", None) or [])

333 + quant_config = getattr(hf_config, "quantization_config", None)

334 + checkpoint = hash_checkpoint(model_config.model) or model_config.model

335 + return cls(

336 + checkpoint=checkpoint,

337 + model_arch=arch,

338 + tp_size=tp_size,

339 + tp_rank=tp_rank,

340 + dtype=str(model_config.dtype),

341 + quantization=model_config.quantization,

342 + quant_config_hash=_hash_quant_config(quant_config),

343 + revision=model_config.revision,

344 + vllm_version=vllm.version.__version__,

345 + )

346 +

347 + def mismatched_fields(self, other: "WeightCacheKey") -> list[str]:

348 + return [

349 + f.name

350 + for f in fields(self)

351 + if getattr(self, f.name) != getattr(other, f.name)

352 + ]

353 +

354 +

355 +@dataclass

356 +class TensorEntry:

357 + """A single cached tensor.

358 +

359 + CUDA tensors are exported as `torch.multiprocessing` reduction args

360 + (CUDA IPC handles); non-CUDA tensors are shipped by value.

361 + """

362 +

363 + kind: str

364 + """Either "param" or "buffer"."""

365 + ipc_args: tuple | None = None

366 + cpu_tensor: torch.Tensor | None = None

367 +

368 + @classmethod

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

370 + from torch.multiprocessing.reductions import reduce_tensor

371 +

372 + tensor = tensor.detach()

373 + if tensor.is_cuda:

374 + _, ipc_args = reduce_tensor(tensor)

375 + return cls(kind=kind, ipc_args=ipc_args)

376 + return cls(kind=kind, cpu_tensor=tensor.cpu())

377 +

378 + def rebuild(self, device_index: int) -> torch.Tensor:

379 + if self.ipc_args is None:

380 + assert self.cpu_tensor is not None

381 + return self.cpu_tensor

382 + from torch.multiprocessing.reductions import rebuild_cuda_tensor

383 +

384 + args = list(self.ipc_args)

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

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

387 + # have different CUDA_VISIBLE_DEVICES mappings.

388 + args[6] = device_index

389 + return rebuild_cuda_tensor(*args)

390 +

391 +

392 +def send_msg(sock: socket.socket, obj: Any) -> None:

393 + payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)

394 + sock.sendall(_LEN_STRUCT.pack(len(payload)))

395 + sock.sendall(payload)

396 +

397 +

398 +def recv_msg(sock: socket.socket) -> Any:

399 + (length,) = _LEN_STRUCT.unpack(_recv_exact(sock, _LEN_STRUCT.size))

400 + if length > MAX_MSG_SIZE:

401 + raise ValueError(f"Message size {length} exceeds limit {MAX_MSG_SIZE}")

402 + return pickle.loads(_recv_exact(sock, length))

403 +

404 +

405 +def _recv_exact(sock: socket.socket, num_bytes: int) -> bytes:

406 + buf = bytearray()

407 + while len(buf) < num_bytes:

408 + chunk = sock.recv(num_bytes - len(buf))

409 + if not chunk:

410 + raise ConnectionError("Socket closed while receiving message")

411 + buf.extend(chunk)

412 + return bytes(buf)

vllm/model_executor/utils.py

+23

@@ -3,12 +3,35 @@

3 """Utils for model executor."""

5 import copy

6 +from contextlib import contextmanager

7 +from contextvars import ContextVar

8 from typing import Any

10 import torch

12 from vllm.utils.torch_utils import is_torch_equal_or_newer

14 +_weights_pre_processed: ContextVar[bool] = ContextVar(

15 + "weights_pre_processed", default=False

16 +)

17 +

18 +

19 +@contextmanager

20 +def weights_already_processed():

21 + """Mark that weights are already in post-processed (runtime) format, so

22 + ``process_weights_after_loading`` skips tensor transforms (used by the

23 + weight cache IPC loader).

24 + """

25 + token = _weights_pre_processed.set(True)

26 + try:

27 + yield

28 + finally:

29 + _weights_pre_processed.reset(token)

30 +

31 +

32 +def is_weights_pre_processed() -> bool:

33 + return _weights_pre_processed.get()

34 +

36 def set_weight_attrs(

37 weight: torch.Tensor,