weight cache deamon

docs/serving/weight-cache.md

@@ -0,0 +1,121 @@
+# Weight Cache Daemon (Fast Recovery)
+
+The weight cache daemon keeps a model's **post-quantized, tensor-parallel-sharded
+weights resident in GPU memory** in a small, long-lived process per rank. When an
+engine starts (or restarts), each engine rank maps those weights **zero-copy via
+CUDA IPC** instead of reading, dtype-converting, and quantizing checkpoints from
+disk. This turns a multi-minute weight load into a sub-second attach, which makes
+engine restarts (config changes, crashes, rolling upgrades) fast.
+
+::: warning Linux + NVIDIA GPU only
+The daemon relies on CUDA IPC handles and a POSIX parent-death signal. It is a
+single-node, NVIDIA-GPU feature. On other platforms `--weight-cache-mode` should
+be left `off`.
+:::
+
+## How it works
+
+- One daemon runs per **global rank** (`mapping.rank`) and owns that rank's GPU.
+- The daemon loads the model once, runs `process_weights_after_loading`, then
+  exports every parameter and persistent buffer as a CUDA IPC handle over a
+  per-rank Unix domain socket (`/tmp/tokenspeed_weight_cache_rank{rank}.sock`).
+- An engine rank connects, sends a **`CacheConfig` fingerprint** (model path,
+  architecture, parallelism topology, quant method + config hash, dtype,
+  revision, device capability, torch version), and only maps the weights if the
+  fingerprint matches exactly. Any mismatch falls back to a disk load (client
+  mode) or hard-errors (daemon mode).
+- The engine initializes the model on the meta device (no allocation) and swaps
+  in the imported IPC tensors, so no weight bytes are copied.
+
+### Supported quantization
+
+Zero-copy sharing only exports raw tensor data, so it is correct **only** when
+`process_weights_after_loading` is fully captured by that data. TokenSpeed
+enforces an allowlist:
+
+| Quantization | Supported |
+| --- | --- |
+| Unquantized (`bf16`/`fp16`) | Yes |
+| Block-wise FP8 (`weight_block_size` set) | Yes |
+| Per-tensor FP8, AWQ, GPTQ, Marlin, NVFP4, … | No — hard error |
+
+Methods that transpose/repack weights or stamp Python-side metadata the
+meta-initialized engine cannot reproduce are rejected up front rather than
+silently serving wrong numerics. For an unsupported model, disable the cache
+with `--weight-cache-mode off`.
+
+## Usage
+
+### Daemon mode (engine-managed)
+
+The engine launches the daemons for you, waits until they finish loading, then
+attaches every rank via IPC. Use this for the first start on a fresh host:
+
+```bash
+tokenspeed serve <model> \
+  --attn-tp-size 8 \
+  --weight-cache-mode daemon
+```
+
+The daemons keep running after the engine attaches, so a later restart of the
+engine reattaches in under a second.
+
+### Client mode (external daemons)
+
+Start the daemons out-of-band once, then point one or more short-lived engine
+processes at them. Use this when you restart the engine frequently and want the
+weights to survive across restarts:
+
+```bash
+# 1. Launch the daemons once (blocks until every rank is ready).
+python -m tokenspeed.runtime.weight_cache.daemon \
+  --model-path <model> \
+  --attn-tp-size 8
+
+# 2. Start (and later restart) the engine against the running daemons.
+tokenspeed serve <model> \
+  --attn-tp-size 8 \
+  --weight-cache-mode client
+```
+
+In client mode, if no daemon socket is present the engine falls back to a normal
+disk load instead of failing.
+
+## Parameters
+
+| Parameter | Use |
+| --- | --- |
+| `--weight-cache-mode` | `off` (default), `daemon` (engine launches daemons), or `client` (attach to pre-running daemons). |
+| `--weight-cache-socket` | Override the per-rank Unix socket path. Defaults to `/tmp/tokenspeed_weight_cache_rank{rank}.sock`. |
+
+The standalone launcher (`python -m tokenspeed.runtime.weight_cache.daemon`)
+accepts the parallelism topology directly:
+
+| Flag | Use |
+| --- | --- |
+| `--model-path` | Model to load and cache. |
+| `--attn-tp-size` / `--dense-tp-size` / `--moe-tp-size` | Layer-family tensor parallel sizes (mirror the engine's `Mapping`). |
+| `--ep-size` / `--dp-size` | Expert- and data-parallel sizes. |
+| `--nnodes` / `--node-rank` / `--base-gpu-id` / `--gpu-id-step` | Multi-node / GPU placement. |
+| `--rank` | Run a single rank's daemon (omit to launch all local ranks). |
+| `--load-format` / `--dtype` / `--quantization` / `--revision` | Weight load options; must match the engine to pass the fingerprint check. |
+| `--force` | Kill and take over a wedged daemon that still holds the socket. |
+
+## Operational notes
+
+- **Topology must match.** The engine and its daemons must use the same
+  parallelism sizes, dtype, quantization, and model revision, or the fingerprint
+  check will reject the attach.
+- **Memory.** Weights live in the daemon's GPU memory; the engine shares them
+  read-only. The engine therefore skips the CPU weight backup used by
+  `release_memory_occupation`, so sleep/wake that offloads weights is not
+  combined with the weight cache.
+- **Allocator.** CUDA IPC is incompatible with `expandable_segments`; the daemon
+  refuses to start if that allocator mode is set.
+- **Lifecycle.** Each daemon installs a parent-death signal in daemon mode and
+  writes a `*.ready` file recording its PID. Stale `*.sock`/`*.ready` files from a
+  crashed daemon are cleaned up automatically on the next launch; a still-running
+  daemon is left untouched unless `--force` is passed.
+- **Multi-node.** `--weight-cache-mode daemon` is single-node only. For
+  multi-node, pre-launch daemons on each node with the standalone launcher and
+  start the engine with `--weight-cache-mode client`.

python/tokenspeed/runtime/configs/load_config.py

@@ -38,6 +38,9 @@ class LoadFormat(str, enum.Enum):
     SHARDED_STATE = "sharded_state"
     MISTRAL = "mistral"
     EXTENSIBLE = "extensible"
+    # Load post-quantized, TP-sharded weights from a running weight cache
+    # daemon via CUDA IPC zero-copy mapping instead of from disk.
+    IPC_CACHE = "ipc_cache"
 
 
 @dataclass
@@ -72,6 +75,12 @@ class LoadConfig:
 
     ext_yaml: str | None = None
 
+    # Weight cache daemon (CUDA IPC zero-copy weight loading). ``weight_cache_mode``
+    # is one of "off" | "daemon" | "client"; ``weight_cache_socket`` optionally
+    # overrides the auto-derived per-rank Unix socket path used to reach the daemon.
+    weight_cache_mode: str = "off"
+    weight_cache_socket: str | None = None
+
     def __post_init__(self) -> None:
         model_loader_extra_config = self.model_loader_extra_config or {}
         if isinstance(model_loader_extra_config, str):

python/tokenspeed/runtime/entrypoints/engine.py

@@ -519,6 +519,47 @@ def launch_phase_sigquit_handler(signum, frame):
     mp.set_start_method("spawn", force=True)
 
 
+def _launch_weight_cache_daemons(server_args: ServerArgs) -> list:
+    """Spawn one weight cache daemon per local rank and wait until they are ready.
+
+    Returns the list of daemon subprocess handles (already loaded and serving
+    IPC handles). The daemons keep running independently of this call; each
+    installs a parent-death signal so it exits with the engine.
+    """
+    from tokenspeed.runtime.weight_cache.daemon import launch_weight_cache_daemons
+
+    mapping = server_args.mapping
+    if mapping.nnodes > 1:
+        raise ValueError(
+            "--weight-cache-mode daemon currently supports single-node "
+            "deployments only. For multi-node, pre-launch daemons on each node "
+            "with `python -m tokenspeed.runtime.weight_cache.daemon ...` and "
+            "start the engine with --weight-cache-mode client."
+        )
+
+    logger.info("Launching weight cache daemons (mode=daemon) ...")
+    return launch_weight_cache_daemons(
+        model_path=server_args.model,
+        attn_tp_size=mapping.attn.tp_size,
+        dense_tp_size=mapping.dense.tp_size,
+        moe_tp_size=mapping.moe.tp_size,
+        ep_size=mapping.moe.ep_size,
+        dp_size=mapping.attn.dp_size,
+        nnodes=mapping.nnodes,
+        node_rank=server_args.node_rank,
+        base_gpu_id=server_args.base_gpu_id,
+        gpu_id_step=server_args.gpu_id_step,
+        load_format=server_args.load_format,
+        dtype=server_args.dtype,
+        quantization=server_args.quantization,
+        trust_remote_code=server_args.trust_remote_code,
+        revision=server_args.revision,
+        download_dir=server_args.download_dir,
+        device=server_args.device,
+        wait=True,
+    )
+
+
 def _launch_subprocesses(
     server_args: ServerArgs, port_args: PortArgs | None = None
 ) -> tuple[AsyncLLM, None, dict]:
@@ -539,6 +580,21 @@ def _launch_subprocesses(
         server_args.model, server_args.tokenizer
     )
 
+    # Weight cache daemon (fast engine recovery via CUDA IPC). In "daemon" mode
+    # the engine spawns one weight cache daemon per local rank and blocks until
+    # they finish loading from disk; every engine rank then maps their
+    # post-quantized weights via zero-copy IPC. In "client" mode the daemons are
+    # assumed to be pre-running (started out-of-band), so nothing is launched
+    # here — the engine ranks connect to the existing sockets.
+    weight_cache_daemon_procs = None
+    if getattr(server_args, "weight_cache_mode", "off") == "daemon":
+        weight_cache_daemon_procs = _launch_weight_cache_daemons(server_args)
+        logger.info(
+            "%d weight cache daemon(s) ready; engine ranks will map weights "
+            "via CUDA IPC.",
+            len(weight_cache_daemon_procs),
+        )
+
     scheduler_procs = []
     if not server_args.mapping.attn.has_dp:
         # Launch tensor parallel scheduler processes

python/tokenspeed/runtime/execution/weight_loader.py

@@ -73,18 +73,28 @@ def load_model(
 
         set_cuda_arch()
 
+        # Weight cache daemon (CUDA IPC zero-copy loading). When enabled, weights
+        # are mapped from a running daemon's GPU memory instead of read from disk.
+        weight_cache_mode = getattr(server_args, "weight_cache_mode", "off")
+        use_weight_cache = weight_cache_mode != "off"
+
         # Create load config
         load_config = LoadConfig(
-            load_format=server_args.load_format,
+            load_format=("ipc_cache" if use_weight_cache else server_args.load_format),
             download_dir=server_args.download_dir,
             ext_yaml=server_args.ext_yaml,
             weight_loader_prefetch_checkpoints=server_args.weight_loader_prefetch_checkpoints,
             weight_loader_prefetch_num_threads=server_args.weight_loader_prefetch_num_threads,
+            weight_cache_mode=weight_cache_mode,
+            weight_cache_socket=getattr(server_args, "weight_cache_socket", None),
         )
 
-        # Load model with memory saver context. Tag as "weights" with CPU backup
-        # so release_memory_occupation offloads (and restores) them byte-exact.
-        with memory_saver_adapter.region(tag="weights", enable_cpu_backup=True):
+        # In zero-copy IPC mode the weights already live in the daemon's GPU
+        # memory and are shared read-only, so a CPU backup would both waste host
+        # memory and (on restore) detach the engine from the shared mapping.
+        with memory_saver_adapter.region(
+            tag="weights", enable_cpu_backup=not use_weight_cache
+        ):
             model = get_model(
                 model_config=model_config,
                 load_config=load_config,

python/tokenspeed/runtime/model_loader/loader.py

@@ -662,4 +662,19 @@ def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
     if load_config.load_format == LoadFormat.EXTENSIBLE:
         return ExtensibleModelLoader(load_config)
 
+    if load_config.load_format == LoadFormat.IPC_CACHE:
+        # Imported lazily so the common load paths don't pull in torch IPC and
+        # the weight cache stack on every get_model_loader() call.
+        from tokenspeed.runtime.weight_cache.ipc_loader import IpcModelLoader
+
+        return IpcModelLoader(
+            load_config,
+            socket_path=load_config.weight_cache_socket,
+            weight_cache_mode=(
+                load_config.weight_cache_mode
+                if load_config.weight_cache_mode != "off"
+                else "client"
+            ),
+        )
+
     return DefaultModelLoader(load_config)

python/tokenspeed/runtime/utils/server_args.py

@@ -110,6 +110,12 @@ class ServerArgs:
     revision: str | None = None
     language_model_only: bool = False
 
+    # Weight cache daemon (CUDA IPC zero-copy weight loading for fast engine
+    # recovery). "off" disables it; "daemon" launches per-rank daemons and
+    # loads via IPC; "client" connects to pre-running daemons (restart path).
+    weight_cache_mode: str = "off"
+    weight_cache_socket: str | None = None
+
     # Direct SMG msgpack ZMQ path. When enabled, the scheduler skips the pickle
     # PULL/PUSH IPC and instead connects to SMG (which binds the handshake/input/
     # output sockets) over the msgpack wire. Default OFF;
@@ -995,6 +1001,28 @@ def add_cli_args(parser: argparse.ArgumentParser):
             "a numpy cache to speed up the loading. "
             '"dummy" will initialize the weights with random values.',
         )
+        parser.add_argument(
+            "--weight-cache-mode",
+            type=str,
+            default=ServerArgs.weight_cache_mode,
+            choices=["off", "daemon", "client"],
+            help="Weight cache daemon mode for fast engine recovery via CUDA "
+            "IPC zero-copy weight loading. "
+            '"off" (default) loads weights from disk normally. '
+            '"daemon" launches one weight cache daemon per rank, then maps '
+            "their post-quantized weights via IPC (first start). "
+            '"client" connects to pre-running daemons and maps their weights '
+            "via IPC (engine restart); falls back to disk load only when no "
+            "daemon socket is present.",
+        )
+        parser.add_argument(
+            "--weight-cache-socket",
+            type=str,
+            default=ServerArgs.weight_cache_socket,
+            help="Override the auto-derived per-rank Unix socket path used to "
+            "reach the weight cache daemon. Defaults to "
+            "/tmp/tokenspeed_weight_cache_rank{rank}.sock.",
+        )
         parser.add_argument(
             "--trust-remote-code",
             action=argparse.BooleanOptionalAction,

python/tokenspeed/runtime/weight_cache/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Intentionally light: importing a weight_cache submodule (e.g.
+# ``tokenspeed.runtime.weight_cache.protocol``) executes this package __init__
+# first. Eagerly re-exporting daemon/ipc_loader here would pull in torch and the
+# model loader on that cheap protocol import, re-introducing the circular-import
+# and startup-cost problems the local-import layout avoids. Import the concrete
+# symbols from their submodules instead, e.g.
+#     from tokenspeed.runtime.weight_cache.protocol import CacheConfig
+#     from tokenspeed.runtime.weight_cache.daemon import launch_weight_cache_daemons
+#     from tokenspeed.runtime.weight_cache.ipc_loader import IpcModelLoader