Fast Start

tests/model_executor/model_loader/test_weight_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for the IPC weight cache loader.
+
+A cold start (no weight cache daemon, so the loader falls back to disk) and
+warm restarts (weights mapped from the daemon via CUDA IPC) must both serve
+identical outputs.
+"""
+
+import shutil
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+
+from vllm import SamplingParams
+
+MODEL = "Qwen/Qwen3.5-0.8B"
+PROMPTS = [
+    "Hello, my name is",
+    "The capital of France is",
+]
+DAEMON_TIMEOUT_S = 600
+
+
+class WeightCacheDaemon:
+    """Context manager running the real weight cache daemon as a subprocess."""
+
+    _READY_MARKER = "Weight cache daemon READY"
+
+    def __init__(self, model: str, tp_size: int):
+        # Short base path: Unix socket paths are limited to ~107 characters.
+        self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
+        self._cmd = [
+            sys.executable,
+            "-m",
+            "vllm.model_executor.model_loader.weight_cache.daemon",
+            "--model",
+            model,
+            "--tensor-parallel-size",
+            str(tp_size),
+            "--weight-cache-socket-dir",
+            self.socket_dir,
+            "--enforce-eager",
+        ]
+        self._proc: subprocess.Popen | None = None
+        self._lines: list[str] = []
+
+    def __enter__(self) -> "WeightCacheDaemon":
+        self._proc = subprocess.Popen(
+            self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True
+        )
+        threading.Thread(target=self._drain_stderr, daemon=True).start()
+        try:
+            self._wait_ready(DAEMON_TIMEOUT_S)
+        except Exception:
+            self._stop()
+            raise
+        return self
+
+    def __exit__(self, *exc_info) -> None:
+        self._stop()
+
+    def _drain_stderr(self) -> None:
+        assert self._proc is not None and self._proc.stderr is not None
+        for line in self._proc.stderr:
+            self._lines.append(line)
+
+    def _logs(self) -> str:
+        return "".join(self._lines)
+
+    def _wait_ready(self, timeout_s: float) -> None:
+        assert self._proc is not None
+        deadline = time.monotonic() + timeout_s
+        while time.monotonic() < deadline:
+            if self._proc.poll() is not None:
+                raise RuntimeError(
+                    f"Weight cache daemon exited with {self._proc.returncode}:\n"
+                    f"{self._logs()}"
+                )
+            if self._READY_MARKER in self._logs():
+                return
+            time.sleep(1.0)
+        raise TimeoutError(
+            f"Weight cache daemon not ready after {timeout_s}s:\n{self._logs()}"
+        )
+
+    def _stop(self) -> None:
+        assert self._proc is not None
+        self._proc.terminate()
+        try:
+            self._proc.wait(timeout=60)
+        except subprocess.TimeoutExpired:
+            self._proc.kill()
+            self._proc.wait()
+        shutil.rmtree(self.socket_dir, ignore_errors=True)
+
+
+def generate(vllm_runner, socket_dir: str | None, fallback: bool, **kwargs):
+    extra_config = (
+        {} if socket_dir is None else {"socket_dir": socket_dir, "fallback": fallback}
+    )
+    with vllm_runner(
+        MODEL,
+        load_format="auto" if socket_dir is None else "ipc_cache",
+        model_loader_extra_config=extra_config,
+        # Greedy outputs are compared across engine restarts; cap the batch
+        # size so scheduling cannot change the results.
+        max_num_seqs=1,
+        **kwargs,
+    ) as llm:
+        sampling_params = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True)
+        return llm.generate(PROMPTS, sampling_params)
+
+
+def test_ipc_cache_cold_start_and_warm_restart(vllm_runner):
+    """Cold start falls back to disk; warm restarts load weights via CUDA IPC.
+
+    All runs must produce outputs identical to a default-loader baseline. The
+    warm runs disable the disk fallback, so they only pass if the weights
+    really came from the daemon.
+    """
+    llm_kwargs = dict(
+        gpu_memory_utilization=0.3,
+        enforce_eager=True,
+        # Qwen3.5 is a hybrid mamba model; its mamba cache mode requires
+        # chunked prefill.
+        enable_chunked_prefill=True,
+    )
+
+    # Baseline: plain disk loading with the default loader.
+    baseline_outputs = generate(vllm_runner, None, fallback=True, **llm_kwargs)
+    assert all(text for _, texts in baseline_outputs for text in texts)
+
+    # Cold start: no daemon is serving, so the loader falls back to disk.
+    with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:
+        cold_outputs = generate(
+            vllm_runner, empty_socket_dir, fallback=True, **llm_kwargs
+        )
+
+    with WeightCacheDaemon(MODEL, tp_size=1) as daemon:
+        warm_outputs = generate(
+            vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs
+        )
+        # Warm restart: a second engine lifetime against the same daemon.
+        restart_outputs = generate(
+            vllm_runner, daemon.socket_dir, fallback=False, **llm_kwargs
+        )
+
+    assert cold_outputs == baseline_outputs
+    assert warm_outputs == baseline_outputs
+    assert restart_outputs == baseline_outputs

tools/pre_commit/check_forbidden_imports.py

@@ -48,6 +48,7 @@ class ForbiddenImport:
             "vllm/distributed/device_communicators/shm_object_storage.py",
             "vllm/distributed/weight_transfer/ipc_engine.py",
             "vllm/distributed/weight_transfer/clients.py",
+            "vllm/model_executor/model_loader/weight_cache/protocol.py",
             "tests/distributed/test_shm_broadcast.py",
             "tests/distributed/test_weight_transfer.py",
             "vllm/utils/hashing.py",

vllm/config/load.py

@@ -38,6 +38,10 @@ class LoadConfig:
     - "instanttensor" will load the Safetensors weights on CUDA devices using
       InstantTensor, which enables distributed loading with pipelined prefetching
       and fast direct I/O.
+    - "ipc_cache" will map post-quantized weights from a local weight cache
+      daemon via CUDA IPC for fast engine restarts. See
+      `vllm/model_executor/model_loader/weight_cache/daemon.py` for how to
+      launch the daemon.
     - "npcache" will load the weights in pytorch format and store a numpy cache
       to speed up the loading.
     - "dummy" will initialize the weights with random values, which is mainly

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

@@ -27,7 +27,11 @@
 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
     SharedExperts,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    is_weights_pre_processed,
+    replace_parameter,
+    set_weight_attrs,
+)
 from vllm.platforms import current_platform
 
 if TYPE_CHECKING:
@@ -43,6 +47,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
 
     # --8<-- [end:unquantized_fused_moe]
 
+    supports_pre_processed_weights = True
+
     def __init__(self, moe: FusedMoEConfig):
         super().__init__(moe)
         self.unquantized_backend, self.experts_cls = select_unquantized_moe_backend(
@@ -153,28 +159,38 @@ def _setup_kernel(
             # optional w{13,2}_bias references and SwiGLU gate params. Since
             # weight updates mutate those bias tensors in place, the kernel
             # does not need to be re-built.
-            self.moe_quant_config = self.get_fused_moe_quant_config(layer)
-            assert self.moe_quant_config is not None
-            assert self.experts_cls is not None
-            self.moe_kernel = make_unquantized_moe_kernel(
-                quant_config=self.moe_quant_config,
-                moe_config=self.moe,
-                backend=self.unquantized_backend,
-                experts_cls=self.experts_cls,
-                routing_tables=layer._expert_routing_tables(),
-            )
+            self._init_moe_kernel(layer)
 
             if self.unquantized_backend == UnquantizedMoeBackend.CPU:
                 # The CPU experts need the layer itself for the setup that
                 # convert_to_unquantized_kernel_format cannot express, since
                 # it only sees the two weight tensors: padding and prepacking
                 # into the grouped-gemm layout (bias included), and capturing
                 # the router config that monolithic apply() cannot carry.
+                assert self.moe_kernel is not None
                 self.moe_kernel.fused_experts.process_weights_after_loading(layer)
 
+    def _init_moe_kernel(self, layer: "RoutedExperts") -> None:
+        """Build the MoE kernel from the layer's current (shuffled) weights."""
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+        assert self.experts_cls is not None
+        self.moe_kernel = make_unquantized_moe_kernel(
+            quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            backend=self.unquantized_backend,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._expert_routing_tables(),
+        )
+
     def process_weights_after_loading(self, layer: "RoutedExperts") -> None:
         super().process_weights_after_loading(layer)
 
+        if is_weights_pre_processed():
+            # Weights are already in runtime format; rebuild the kernel only.
+            self._init_moe_kernel(layer)
+            return
+
         # Padding may allocate a same-shaped strided view. Copy it back instead
         # of rebinding `.data`, since CUDA graphs capture the parameter storage
         # address.

vllm/model_executor/layers/linear.py

@@ -163,6 +163,8 @@ def apply(
 class UnquantizedLinearMethod(LinearMethodBase):
     """Linear method without quantization."""
 
+    supports_pre_processed_weights = True
+
     def __init__(self) -> None:
         config = get_current_vllm_config_or_none()
         linear_backend = (

vllm/model_executor/layers/quantization/base_config.py

@@ -27,6 +27,13 @@ class QuantizeMethodBase(ABC):
     in process_weights_after_loading, reducing peak memory during loading.
     """
 
+    supports_pre_processed_weights: bool = False
+    """
+    Whether ``process_weights_after_loading`` supports running under
+    ``weights_already_processed``. Methods must skip tensor transforms in
+    that mode; the loader driver rejects methods that do not declare support.
+    """
+
     @abstractmethod
     def create_weights(
         self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs

vllm/model_executor/layers/quantization/fp8.py

@@ -74,7 +74,11 @@
     BlockQuantScaleParameter,
     PerTensorScaleParameter,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    is_weights_pre_processed,
+    replace_parameter,
+    set_weight_attrs,
+)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import (
     is_deep_gemm_supported,
@@ -248,6 +252,8 @@ class Fp8LinearMethod(LinearMethodBase):
         quant_config: The quantization config.
     """
 
+    supports_pre_processed_weights = True
+
     def __init__(self, quant_config: Fp8Config):
         self.quant_config = quant_config
         self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False)
@@ -366,6 +372,19 @@ def create_weights(
         self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if is_weights_pre_processed():
+            # Weights are already in runtime format; the only state tensor
+            # export cannot carry is the `input_scale = None` stamp that dynamic
+            # activation quantization relies on, since create_weights never
+            # registered that parameter.
+            if not self.act_q_static and not hasattr(layer, "input_scale"):
+                layer.input_scale = None
+            # Marlin reads marlin_input_dtype in apply_weights; it lives on the
+            # method (not exported with the weights), so restore it here too.
+            if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):
+                self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
+            return
+
         if self.use_marlin:
             if not self.block_quant:
                 # Canonicalize to (K, N) for the kernel.
@@ -472,6 +491,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         quant_config: The quantization config.
     """
 
+    supports_pre_processed_weights = True
+
     def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
         super().__init__(layer.moe_config)
         self.quant_config = quant_config
@@ -713,6 +734,10 @@ def _setup_kernel(
         replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
         replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
 
+        self._init_moe_kernel(layer)
+
+    def _init_moe_kernel(self, layer: RoutedExperts) -> None:
+        """Build the MoE kernel from the layer's current (converted) weights."""
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         assert self.moe_quant_config is not None
         assert self.experts_cls is not None
@@ -725,6 +750,11 @@ def _setup_kernel(
         )
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
+        if is_weights_pre_processed():
+            # Weights are already in kernel format; rebuild the kernel only.
+            self._init_moe_kernel(layer)
+            return
+
         # Allow for accessing weights and scales in standard way.
         w13 = layer.w13_weight
         w2 = layer.w2_weight

vllm/model_executor/layers/quantization/kv_cache.py

@@ -8,6 +8,7 @@
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from vllm.model_executor.utils import is_weights_pre_processed
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales
@@ -51,6 +52,8 @@ class BaseKVCacheMethod(QuantizeMethodBase):
         quant_config: the appropriate QuantizationConfig
     """
 
+    supports_pre_processed_weights = True
+
     def __init__(self, quant_config: QuantizationConfig):
         self.quant_config = quant_config
 
@@ -72,6 +75,20 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor:
         raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if is_weights_pre_processed():
+            # Scale buffers were exported with post-processed values; drop the
+            # consumed scale placeholders and rebuild the host-side copies.
+            for name in ("k_scale", "v_scale", "q_scale", "prob_scale"):
+                if hasattr(layer, name):
+                    delattr(layer, name)
+            if hasattr(layer, "_k_scale"):
+                layer._k_scale_float = layer._k_scale.item()
+                layer._v_scale_float = layer._v_scale.item()
+                layer._q_scale_float = layer._q_scale.item()
+                layer._k_scale_cpu = layer._k_scale.detach().to("cpu")
+                layer._v_scale_cpu = layer._v_scale.detach().to("cpu")
+            return
+
         # skip if there are no weights to process (for example, weight reloading)
         if not hasattr(layer, "q_scale"):
             assert not hasattr(layer, "k_scale")

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -36,6 +36,8 @@
 class UnquantizedEmbeddingMethod(QuantizeMethodBase):
     """Unquantized method for embeddings."""
 
+    supports_pre_processed_weights = True
+
     def create_weights(
         self,
         layer: torch.nn.Module,