[FEAT] Support fast engine recovery through weight cache (#27139)

Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com>
Co-authored-by: liusy58 <liusy58@linux.alibaba.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Mike Qiu
2026-07-25 14:31:21 +08:00
committed by GitHub
co-authored by liusy58 Alex Nails
parent bdd0698541
commit f9c14e6bd4
18 changed files with 3404 additions and 2 deletions
+6
View File
@@ -34,6 +34,7 @@ class LoadFormat(str, enum.Enum):
FASTSAFETENSORS = "fastsafetensors"
PRIVATE = "private"
RUNAI_STREAMER = "runai_streamer"
IPC_CACHE = "ipc_cache"
@dataclass
@@ -104,6 +105,11 @@ class LoadConfig:
# For multi-layer MTP
draft_model_idx: Optional[int] = None
# Weight cache daemon options
weight_cache_mode: str = "off" # "off", "daemon", "client"
weight_cache_socket: Optional[str] = None # Path to daemon socket (for client mode)
fallback_load_format: Union[str, "LoadFormat"] = LoadFormat.AUTO
def __post_init__(self):
model_loader_extra_config = self.model_loader_extra_config or {}
if isinstance(model_loader_extra_config, str):
+234 -2
View File
@@ -27,6 +27,8 @@ import multiprocessing as mp
import os
import random
import signal
import subprocess
import sys
import tempfile
import threading
import time
@@ -110,7 +112,12 @@ from sglang.srt.utils import (
set_ulimit,
)
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
from sglang.srt.utils.network import get_zmq_socket, is_port_available
from sglang.srt.utils.network import (
NetworkAddress,
get_free_port,
get_zmq_socket,
is_port_available,
)
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils.watchdog import SubprocessWatchdog
from sglang.version import __version__
@@ -238,6 +245,7 @@ class Engine(EngineScoreMixin, EngineBase):
port_args,
scheduler_init_result,
subprocess_watchdog,
weight_cache_daemon_procs,
) = self._launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=self.init_tokenizer_manager_func,
@@ -247,6 +255,10 @@ class Engine(EngineScoreMixin, EngineBase):
self.tokenizer_manager = tokenizer_manager
self.template_manager = template_manager
self._scheduler_init_result = scheduler_init_result
# Engine-spawned weight cache daemons owned by *this* instance (empty
# unless --weight-cache-mode daemon). Kept per-instance so two Engines
# in one process each reap only their own daemons in shutdown().
self._weight_cache_daemon_procs = weight_cache_daemon_procs
if tokenizer_manager is not None:
tokenizer_manager._subprocess_watchdog = subprocess_watchdog
self.port_args = port_args
@@ -588,6 +600,208 @@ class Engine(EngineScoreMixin, EngineBase):
ret = self.loop.run_until_complete(generator.__anext__())
return ret
@classmethod
def _launch_weight_cache_daemons(cls, server_args: ServerArgs):
"""Launch weight cache daemon processes for this node's PP×TP ranks.
All daemon processes join the same NCCL distributed group so that
TP-sharded model loading works correctly. Each daemon holds its
rank's weight shard in GPU memory and serves IPC handles.
Lifecycle: these daemons are *co-terminal* with the engine. They are
children of this process (kill_itself_when_parent_died installs
PR_SET_PDEATHSIG) and are gracefully reaped in ``shutdown()``. They do
NOT persist across engine restarts, so ``--weight-cache-mode daemon``
on its own does not deliver a faster restart -- the first start is in
fact slower (disk-load into the daemon plus the IPC handshake). The
fast-recovery story is the standalone launcher
(``python -m sglang.srt.weight_cache.daemon``) plus
``--weight-cache-mode client``, where the daemon outlives the engine.
"""
if server_args.dp_size > 1:
raise ValueError(
"Weight cache daemon mode does not support dp_size > 1. "
"Please set --dp-size 1 when using --weight-cache-mode daemon."
)
# Multi-node needs an explicit rendezvous address; otherwise each node
# picks its own local 127.0.0.1 port (below) and the per-node daemons
# can never form the joint process group.
if server_args.nnodes > 1 and not server_args.dist_init_addr:
raise ValueError(
"Multi-node weight cache daemons (nnodes > 1) require "
"--dist-init-addr so all nodes rendezvous at the same endpoint."
)
tp_size = server_args.tp_size
pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
_calculate_rank_ranges(
server_args.nnodes,
server_args.pp_size,
tp_size,
server_args.node_rank,
)
)
# Build the distributed init method (multi-node uses the user-provided
# dist_init_addr so all nodes reach the same endpoint).
if server_args.dist_init_addr:
host, port = server_args.dist_init_addr.rsplit(":", 1)
dist_init_method = f"tcp://{host}:{port}"
else:
# Fresh free port for the daemons' own rendezvous, not the engine's
# nccl_port: a pinned --nccl-port would otherwise collide with the
# engine's own NCCL TCPStore.
dist_init_method = NetworkAddress("127.0.0.1", get_free_port()).to_tcp()
num_daemons = len(pp_rank_range) * len(tp_rank_range)
daemon_procs = []
logger.info(
f"Launching {num_daemons} weight cache daemon(s) on node "
f"{server_args.node_rank} for model={server_args.model_path}, "
f"pp_ranks={pp_rank_range.start}..{pp_rank_range.stop - 1}, "
f"tp_ranks={tp_rank_range.start}..{tp_rank_range.stop - 1}, "
f"dist_init_method={dist_init_method}"
)
# Validate and clean up stale .ready/.sock files from prior runs.
# If a daemon is still alive at this rank, raise instead of clobbering.
from sglang.srt.weight_cache.protocol import (
cleanup_stale_daemon_files,
compute_global_rank,
compute_local_gpu_id,
get_ready_path,
)
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
cleanup_stale_daemon_files(global_rank)
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
gpu_id = compute_local_gpu_id(
pp_rank,
tp_rank,
pp_size_per_node,
tp_size_per_node,
base_gpu_id=server_args.base_gpu_id,
gpu_id_step=server_args.gpu_id_step,
)
cmd = [
sys.executable,
"-m",
"sglang.srt.weight_cache.daemon",
"--model-path",
server_args.model_path,
"--gpu-id",
str(gpu_id),
"--tp-size",
str(tp_size),
"--tp-rank",
str(tp_rank),
"--pp-size",
str(server_args.pp_size),
"--pp-rank",
str(pp_rank),
"--dp-size",
"1",
"--ep-size",
str(server_args.ep_size),
"--load-format",
server_args.load_format,
"--dtype",
server_args.dtype,
"--dist-init-method",
dist_init_method,
]
if server_args.quantization:
cmd += ["--quantization", server_args.quantization]
if (
server_args.model_loader_extra_config
and server_args.model_loader_extra_config != "{}"
):
cmd += [
"--model-loader-extra-config",
server_args.model_loader_extra_config,
]
if server_args.trust_remote_code:
cmd += ["--trust-remote-code"]
if server_args.revision:
cmd += ["--revision", server_args.revision]
proc = subprocess.Popen(cmd)
daemon_procs.append(proc)
# Wait for all daemons to be ready (ready file exists). On any failure
# (readiness timeout or a daemon exiting early) terminate the siblings
# we already spawned before propagating, so a partial launch does not
# leak GPU-resident daemons.
timeout = server_args.weight_cache_timeout
check_interval = 2
start_time = time.time()
try:
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
ready_path = get_ready_path(global_rank)
while not os.path.exists(ready_path):
time.sleep(check_interval)
if time.time() - start_time > timeout:
raise TimeoutError(
f"Weight cache daemon for pp_rank={pp_rank} "
f"tp_rank={tp_rank} did not become ready "
f"within {timeout}s"
)
# Check if daemon process is still alive
for p in daemon_procs:
if p.poll() is not None:
raise RuntimeError(
f"Weight cache daemon (pid={p.pid}) exited prematurely "
f"with code {p.returncode}"
)
logger.info(
f"Weight cache daemon for pp_rank={pp_rank} "
f"tp_rank={tp_rank} is ready"
)
except BaseException:
cls._terminate_weight_cache_daemons(daemon_procs)
raise
logger.info(
f"All {num_daemons} weight cache daemons on node "
f"{server_args.node_rank} are ready"
)
return daemon_procs
@staticmethod
def _terminate_weight_cache_daemons(procs, timeout: float = 10.0):
"""Gracefully stop engine-spawned weight cache daemons.
Send SIGTERM first so each daemon's signal handler can unlink its
``.sock``/``.ready`` files, then SIGKILL any straggler. This matters
because ``shutdown()`` otherwise reaps children via
``kill_process_tree`` (SIGKILL), which would skip that cleanup and
leave stale files that make the next client-mode boot fail with a
confusing "socket exists but connection refused" instead of a clean
"no daemon" path.
"""
if not procs:
return
for p in procs:
if p.poll() is None:
p.terminate() # SIGTERM -> daemon cleanup handler runs
for p in procs:
try:
p.wait(timeout=timeout)
except subprocess.TimeoutExpired:
logger.warning(
f"Weight cache daemon (pid={p.pid}) did not exit within "
f"{timeout}s of SIGTERM; sending SIGKILL."
)
p.kill()
@classmethod
def _launch_scheduler_processes(
cls,
@@ -779,7 +993,7 @@ class Engine(EngineScoreMixin, EngineBase):
"""Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.
Returns:
Tuple of (tokenizer_manager, template_manager, port_args, scheduler_init_result, subprocess_watchdog).
Tuple of (tokenizer_manager, template_manager, port_args, scheduler_init_result, subprocess_watchdog, weight_cache_daemon_procs).
"""
# Configure global environment
configure_logger(server_args)
@@ -820,6 +1034,13 @@ class Engine(EngineScoreMixin, EngineBase):
):
resolve_auto_parsers(server_args)
# Launch daemons (daemon mode only). Handles are threaded back to the
# owning Engine instance (not a class attr) so two Engines in one process
# don't clobber each other's daemon list.
weight_cache_daemon_procs: List = []
if server_args.weight_cache_mode == "daemon":
weight_cache_daemon_procs = cls._launch_weight_cache_daemons(server_args)
# Launch scheduler processes
scheduler_init_result, scheduler_procs = cls._launch_scheduler_processes(
server_args, port_args, run_scheduler_process_func
@@ -846,6 +1067,7 @@ class Engine(EngineScoreMixin, EngineBase):
port_args,
scheduler_init_result,
None,
weight_cache_daemon_procs,
)
launch_dummy_health_check_server(
@@ -859,6 +1081,7 @@ class Engine(EngineScoreMixin, EngineBase):
port_args,
scheduler_init_result,
None,
weight_cache_daemon_procs,
)
# Launch detokenizer process(es) — optionally fronted by a router when
@@ -906,6 +1129,7 @@ class Engine(EngineScoreMixin, EngineBase):
port_args,
scheduler_init_result,
subprocess_watchdog,
weight_cache_daemon_procs,
)
def shutdown(self):
@@ -923,6 +1147,14 @@ class Engine(EngineScoreMixin, EngineBase):
send_to_rpc.close(linger=0)
self.send_to_rpc = None
# Gracefully stop weight cache daemons *before* the blanket
# kill_process_tree below, so their SIGTERM handlers can unlink the
# .sock/.ready files instead of being SIGKILLed and leaving stale state.
daemon_procs = getattr(self, "_weight_cache_daemon_procs", None)
if daemon_procs:
self._terminate_weight_cache_daemons(daemon_procs)
self._weight_cache_daemon_procs = []
kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60)
def __enter__(self):
@@ -2675,6 +2675,7 @@ def launch_server(
port_args,
scheduler_init_result,
subprocess_watchdog,
_weight_cache_daemon_procs,
) = Engine._launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=init_tokenizer_manager_func,
@@ -183,6 +183,22 @@ class SchedulerWeightUpdaterManager:
parameter = self.tp_worker.get_weights_by_name(recv_req)
return GetWeightsByNameReqOutput(parameter=parameter)
def _assert_weight_cache_inactive(self, op: str) -> None:
"""Reject freeing/restoring model weights while the CUDA IPC weight
cache is active: the weights are shared with the daemon via CUDA IPC, so
freeing them would leave the daemon and every peer pointing at released
memory.
"""
mode = self.tp_worker.model_runner.server_args.weight_cache_mode
if mode != "off":
raise RuntimeError(
f"[weight_cache] {op} of model weights is not supported while the "
f"weight cache is active (--weight-cache-mode {mode}): the weights "
f"are shared with the daemon via CUDA IPC, so freeing them would "
f"corrupt the daemon's master copy and every co-attached engine. "
f"Restart with --weight-cache-mode off to use this operation."
)
def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
assert (
self.is_fully_idle()
@@ -215,6 +231,7 @@ class SchedulerWeightUpdaterManager:
self.flush_cache()
if GPU_MEMORY_TYPE_WEIGHTS in tags:
self._assert_weight_cache_inactive("release_memory_occupation")
self.stashed_model_static_state = _export_static_state(
self.tp_worker.model_runner.model
)
@@ -241,6 +258,7 @@ class SchedulerWeightUpdaterManager:
self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH)
if GPU_MEMORY_TYPE_WEIGHTS in tags:
self._assert_weight_cache_inactive("resume_memory_occupation")
self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
torch.distributed.barrier(self.tp_cpu_group)
_import_static_state(
@@ -128,6 +128,7 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import (
load_kv_cache_scales,
load_model_with_memory_saver,
maybe_downgrade_dtype_for_legacy_gpu,
maybe_enable_ipc_weight_cache,
maybe_register_debug_tensor_dump_hook,
maybe_trigger_remote_instance_nccl_send_group,
report_online_quantization,
@@ -987,6 +988,18 @@ class ModelRunner:
remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine,
remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id,
draft_model_idx=self.draft_model_idx,
weight_cache_mode=self.server_args.weight_cache_mode,
weight_cache_socket=self.server_args.weight_cache_socket,
)
# If the weight cache is enabled, override the load format to IPC_CACHE
# and derive the per-rank daemon socket. Idempotent across reloads.
maybe_enable_ipc_weight_cache(
load_config=self.load_config,
server_args=self.server_args,
tp_size=self.ps.tp_size,
pp_rank=self.ps.pp_rank,
tp_rank=self.ps.tp_rank,
)
if self.device == "cpu":
self.model_config = adjust_config_with_unaligned_cpu_tp(
@@ -172,6 +172,8 @@ def build_load_config(
remote_instance_weight_transporter_engine: Any,
remote_instance_weight_transporter_session_id: str,
draft_model_idx: Optional[int],
weight_cache_mode: str,
weight_cache_socket: Optional[str],
) -> LoadConfig:
from sglang.srt.configs.modelopt_config import ModelOptConfig
@@ -199,9 +201,45 @@ def build_load_config(
modelopt_config=modelopt_config,
rl_quant_profile=server_args.rl_quant_profile,
draft_model_idx=draft_model_idx,
weight_cache_mode=weight_cache_mode,
weight_cache_socket=weight_cache_socket,
)
def maybe_enable_ipc_weight_cache(
*,
load_config: LoadConfig,
server_args: ServerArgs,
tp_size: int,
pp_rank: int,
tp_rank: int,
) -> None:
"""Switch ``load_config`` onto the IPC weight-cache path, in place.
Overrides the load format to ``IPC_CACHE`` (remembering the original as the
disk fallback) and derives the per-rank daemon socket if unset. Idempotent:
the format swap is guarded on ``!= IPC_CACHE`` so a second call (e.g. a
weight reload) can't overwrite the captured fallback format.
"""
if server_args.weight_cache_mode == "off":
return
if load_config.load_format != LoadFormat.IPC_CACHE:
load_config.fallback_load_format = load_config.load_format
load_config.load_format = LoadFormat.IPC_CACHE
# Compute socket path using global rank (tp_size * pp_rank + tp_rank) so
# each daemon has a unique socket even across PP stages and nodes.
if load_config.weight_cache_socket is None:
from sglang.srt.weight_cache.protocol import (
compute_global_rank,
get_socket_path,
)
global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
load_config.weight_cache_socket = get_socket_path(global_rank=global_rank)
def load_model_with_memory_saver(
*,
server_args: ServerArgs,
@@ -218,6 +256,17 @@ def load_model_with_memory_saver(
enable_cpu_backup = server_args.enable_weights_cpu_backup or (
is_draft_worker and server_args.enable_draft_weights_cpu_backup
)
# In zero-copy IPC mode, the weights are shared with the daemon via
# CUDA IPC and must not be offloaded/reloaded by the memory saver.
is_ipc_zero_copy = server_args.weight_cache_mode != "off"
if is_ipc_zero_copy and enable_cpu_backup:
logger.warning(
"[ModelRunner] Disabling weights CPU backup in zero-copy IPC mode — "
"IPC-mapped weights cannot be offloaded to CPU."
)
enable_cpu_backup = False
remote_instance_weight_info = None
with memory_saver_adapter.region(
GPU_MEMORY_TYPE_WEIGHTS,
@@ -123,6 +123,21 @@ class WeightUpdater:
logger.error(message)
return False, message
def _assert_weight_cache_inactive(self: WeightUpdater, op: str) -> None:
"""Reject weight mutations while the CUDA IPC weight cache is active:
param.data is the daemon's master copy shared with every co-attached
engine, so an in-place update would silently corrupt them all.
"""
mode = self.get_model_runner().server_args.weight_cache_mode
if mode != "off":
raise RuntimeError(
f"[weight_cache] {op} is not supported while the weight cache is "
f"active (--weight-cache-mode {mode}): model weights are shared "
f"with the daemon via CUDA IPC, so mutating them in place would "
f"corrupt the daemon's master copy and every co-attached engine. "
f"Restart with --weight-cache-mode off to use this operation."
)
def update_weights_from_disk(
self: WeightUpdater,
model_path: str,
@@ -131,6 +146,7 @@ class WeightUpdater:
recapture_cuda_graph: bool = False,
) -> tuple[bool, str]:
"""Update engine weights in-place from the disk."""
self._assert_weight_cache_inactive("update_weights_from_disk")
error = _unsupported_derived_weight_cache_error()
if error is not None:
return False, error
@@ -220,6 +236,7 @@ class WeightUpdater:
dtype: the data type of the parameter to be updated.
shape: the shape of the parameter to be updated.
"""
self._assert_weight_cache_inactive("update_weights_from_distributed")
error = _unsupported_derived_weight_cache_error()
if error is not None:
return False, error
@@ -309,6 +326,7 @@ class WeightUpdater:
return False, error
monkey_patch_torch_reductions()
self._assert_weight_cache_inactive("update_weights_from_tensor")
if load_format == "flattened_bucket":
# Handle flattened bucket format
return self._update_weights_from_flattened_bucket(
@@ -368,6 +386,7 @@ class WeightUpdater:
def update_weights_from_ipc(self: WeightUpdater, recv_req):
"""Update weights from IPC for checkpoint-engine integration."""
self._assert_weight_cache_inactive("update_weights_from_ipc")
error = _unsupported_derived_weight_cache_error()
if error is not None:
return False, error
+22
View File
@@ -3319,4 +3319,26 @@ def get_model_loader(
if load_config.load_format == LoadFormat.RUNAI_STREAMER:
return RunaiModelStreamerLoader(load_config)
if load_config.load_format == LoadFormat.IPC_CACHE:
from sglang.srt.weight_cache.ipc_loader import IpcModelLoader
from sglang.srt.weight_cache.protocol import (
compute_global_rank,
get_socket_path,
)
if load_config.weight_cache_socket:
socket_path = load_config.weight_cache_socket
else:
from sglang.srt.runtime_context import get_parallel
ps = get_parallel()
global_rank = compute_global_rank(ps.tp_size, ps.pp_rank, ps.tp_rank)
socket_path = get_socket_path(global_rank=global_rank)
return IpcModelLoader(
load_config=load_config,
socket_path=socket_path,
weight_cache_mode=load_config.weight_cache_mode,
fallback_load_format=load_config.fallback_load_format,
)
return DefaultModelLoader(load_config)
+1
View File
@@ -52,6 +52,7 @@ def launch_server(
port_args,
scheduler_init_result,
subprocess_watchdog,
_weight_cache_daemon_procs,
) = RayEngine._launch_subprocesses(
server_args,
init_tokenizer_manager_func=init_tokenizer_manager_func,
+59
View File
@@ -124,6 +124,10 @@ LOAD_FORMAT_CHOICES = [
"private",
"runai_streamer",
]
# NOTE: LoadFormat.IPC_CACHE intentionally has no public --load-format choice.
# It is an internal dispatch format set automatically by ModelRunner when the
# weight cache is enabled (weight_cache_mode != "off"). Exposing it as a CLI
# choice let users create contradictory combos (see _handle_load_format).
# TODO: this list should likely contain only methods that support online quantization, or that support using custom quantization classes compatible with a given `quant_method` in config.json.
# Some of the choices here do NOT support online quantization.
@@ -3302,6 +3306,36 @@ class ServerArgs:
NS("exec.features"),
] = False
weight_cache_mode: A[
str,
Arg(
help="Weight cache mode. 'off': normal disk loading. "
"'daemon': launch weight cache daemon (holds weights in GPU memory). "
"Engine-spawned daemons are co-terminal with the engine and do NOT "
"persist across restarts, so this alone does not speed up restart "
"(the first start is slower). For fast recovery, run the standalone "
"daemon (python -m sglang.srt.weight_cache.daemon) and connect with "
"'client'. 'client': connect to existing daemon and load via IPC.",
choices=["off", "daemon", "client"],
),
NS("model"),
] = "off"
weight_cache_socket: A[
Optional[str],
Arg(
help="Unix socket path for weight cache daemon (client mode)."
"If not set, uses /tmp/sglang_weight_cache_rank{global_rank}.sock",
),
NS("model"),
] = None
weight_cache_timeout: A[
int,
Arg(
help="Timeout in seconds for weight cache daemon readiness (default: 1800).",
),
NS("model"),
] = 1800
# -------------------------------------------------------------------------
# Custom hooks, probe, and plugins
# -------------------------------------------------------------------------
@@ -6926,6 +6960,31 @@ class ServerArgs:
self.validate_transfer_engine()
)
# "ipc_cache" is an internal-only load format: ModelRunner sets it
# automatically when the weight cache is enabled, and it is not a public
# --load-format choice. Setting it directly is always wrong (no daemon is
# launched, and fallback_load_format inherits a nonsensical format), so
# reject it and point at the knob (defense-in-depth; the CLI already
# rejects it via LOAD_FORMAT_CHOICES).
if self.load_format == "ipc_cache":
raise ValueError(
"load_format='ipc_cache' is an internal-only format and must not "
"be set directly. Enable the weight cache via --weight-cache-mode "
"client (connect to an existing daemon) or daemon (launch one); "
"that selects IPC loading automatically."
)
# Speculative decoding loads an extra draft model whose weights the
# daemon does not export, so refuse the combination up front instead of
# failing deep inside draft-worker load (draft-model daemon TBD).
if self.weight_cache_mode != "off" and self.speculative_algorithm is not None:
raise ValueError(
"--weight-cache-mode is not supported together with speculative "
"decoding (--speculative-algorithm): the weight cache daemon does "
"not export the draft model's weights. Disable one of them "
"(--weight-cache-mode off) for this configuration."
)
def _is_mistral_native_format(self) -> bool:
"""True iff the checkpoint requires load_format=mistral.
@@ -0,0 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
#
# Intentionally empty: importing any weight_cache submodule (e.g.
# ``sglang.srt.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 refactor removed. Import the concrete
# symbols from their submodules instead, e.g.
# from sglang.srt.weight_cache.protocol import CacheConfig
# from sglang.srt.weight_cache.daemon import launch_weight_cache_daemons
# from sglang.srt.weight_cache.ipc_loader import IpcModelLoader
+944
View File
@@ -0,0 +1,944 @@
# SPDX-License-Identifier: Apache-2.0
"""Weight Cache Daemon — a persistent process that holds post-quantized,
TP-sharded model weights in GPU memory and serves them via CUDA IPC handles.
Each GPU runs one daemon process for its TP rank. The daemon:
1. Loads model weights from disk (full pipeline: disk → TP shard → quantize)
2. Exports every parameter/buffer as a CUDA IPC handle
3. Serves handles over a Unix socket to requesting engine processes
4. Validates CacheConfig compatibility before serving
Usage:
# Single-node: launch all TP rank daemons with a single command:
python -m sglang.srt.weight_cache.daemon \
--model-path /path/to/model --tp-size 4 \
--load-format auto --dtype auto --quantization fp8
# Multi-node: run on each node with --nnodes and --node-rank:
# Node 0:
python -m sglang.srt.weight_cache.daemon \
--model-path /path/to/model --tp-size 16 \
--nnodes 2 --node-rank 0 \
--dist-init-method tcp://node0-ip:29500
# Node 1:
python -m sglang.srt.weight_cache.daemon \
--model-path /path/to/model --tp-size 16 \
--nnodes 2 --node-rank 1 \
--dist-init-method tcp://node0-ip:29500
# Or launch a single daemon for a specific rank:
python -m sglang.srt.weight_cache.daemon \
--model-path /path/to/model \
--gpu-id 0 --tp-size 4 --tp-rank 0 \
--dist-init-method tcp://127.0.0.1:29500
"""
import logging
import os
import signal
import socket
import time
from typing import Any, Dict, Optional
import torch
import torch.distributed as dist
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.platforms import current_platform
from sglang.srt.utils import MultiprocessingSerializer
from .protocol import (
CacheConfig,
check_ipc_quant_support,
cleanup_stale_daemon_files,
compute_env_stamp,
compute_global_rank,
compute_local_gpu_id,
get_quant_method_name,
get_ready_path,
get_socket_path,
hash_quant_config,
recv_msg,
send_msg,
)
logger = logging.getLogger(__name__)
# Per-connection timeout for the serial serve loop. A client exchange is tiny
# (a config dict + IPC handle metadata), so this generous bound never trips a
# healthy client, yet guarantees one hung/dead peer can't stall the other
# engine ranks indefinitely.
CLIENT_CONNECTION_TIMEOUT = 30.0
class WeightCacheDaemon:
"""Persistent GPU weight cache for a single TP rank.
Holds the complete post-quantization state_dict in GPU memory and
serves CUDA IPC handles to engine processes via Unix socket.
"""
def __init__(
self,
model_path: str,
gpu_id: int,
tp_size: int = 1,
tp_rank: int = 0,
pp_size: int = 1,
pp_rank: int = 0,
dp_size: int = 1,
ep_size: int = 1,
load_format: str = "auto",
dtype: str = "auto",
quantization: Optional[str] = None,
model_loader_extra_config: str = "{}",
trust_remote_code: bool = False,
revision: Optional[str] = None,
dist_init_method: Optional[str] = None,
):
self.model_path = model_path
self.gpu_id = gpu_id
self.tp_size = tp_size
self.tp_rank = tp_rank
self.pp_size = pp_size
self.pp_rank = pp_rank
self.dp_size = dp_size
self.ep_size = ep_size
self.load_format = load_format
self.dtype = dtype
self.quantization = quantization
self.model_loader_extra_config = model_loader_extra_config
self.trust_remote_code = trust_remote_code
self.revision = revision
self.dist_init_method = dist_init_method
self.socket_path = get_socket_path(
compute_global_rank(tp_size, pp_rank, tp_rank)
)
self.ready_path = get_ready_path(compute_global_rank(tp_size, pp_rank, tp_rank))
self.model = None
self.config: Optional[CacheConfig] = None
# name -> {"handle": base64_str, "shape": list, "dtype": str, "is_param": bool}
self.state_entries: Dict[str, Dict[str, Any]] = {}
def _init_distributed(self, server_args, model_config):
"""Initialize the distributed backend required for model loading.
Uses the same world_size/rank formula as the engine:
world_size = tp_size * pp_size
rank = tp_size * pp_rank + tp_rank
"""
from sglang.srt.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
model_parallel_is_initialized,
)
if model_parallel_is_initialized():
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id}] "
f"Distributed already initialized, skipping"
)
return
# Initialize distributed environment
import torch.distributed as dist
if not dist.is_initialized():
if self.dist_init_method is None:
# Fallback: auto-assign a port. This only works for single-process.
import socket as sock_mod
with sock_mod.socket(sock_mod.AF_INET, sock_mod.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
free_port = s.getsockname()[1]
self.dist_init_method = f"tcp://127.0.0.1:{free_port}"
init_distributed_environment(
world_size=self.tp_size * self.pp_size,
rank=compute_global_rank(self.tp_size, self.pp_rank, self.tp_rank),
distributed_init_method=self.dist_init_method,
local_rank=self.gpu_id,
backend=current_platform.get_torch_distributed_backend_str(),
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
pipeline_model_parallel_size=self.pp_size,
expert_model_parallel_size=self.ep_size,
)
# Initialize DP attention state (required by some models like Qwen3 MoE)
from sglang.srt.layers.dp_attention import initialize_dp_attention
initialize_dp_attention(server_args, model_config)
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id} tp_rank={self.tp_rank}] "
f"Distributed backend initialized (tp_size={self.tp_size}, "
f"pp_size={self.pp_size}, "
f"world_size={self.tp_size * self.pp_size})"
)
def load(self):
"""Full loading pipeline: disk → TP shard → quantize → export IPC handles."""
# CUDA IPC weight sharing relies on torch's _share_cuda_ handle export,
# which only exists on CUDA-alike platforms (CUDA / ROCm). Fail loud here
# instead of dying deep inside the export with an opaque error.
if not current_platform.is_cuda_alike():
raise RuntimeError(
f"[WeightCacheDaemon] the weight cache daemon requires a CUDA-alike "
f"platform (CUDA or ROCm) for CUDA IPC weight sharing, but the "
f"active platform device type is {current_platform.device_type!r}. "
f"Disable the weight cache (--weight-cache-mode off)."
)
# expandable_segments makes torch's caching allocator hand out memory
# that cannot be exported via _share_cuda_, so the IPC export below would
# die mid-way with an opaque CUDA error. Fail fast with an actionable
# message before touching the device.
self._assert_ipc_compatible_allocator()
current_platform.set_device(current_platform.get_device(self.gpu_id))
# Reduce thread contention during multi-process loading
torch.set_num_threads(1)
# Lazy imports to avoid circular dependencies and speed up startup
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.model_loader.loader import get_model_loader
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import ServerArgs
server_args = ServerArgs(
model_path=self.model_path,
dtype=self.dtype,
quantization=self.quantization,
trust_remote_code=self.trust_remote_code,
tp_size=self.tp_size,
pp_size=self.pp_size,
dp_size=self.dp_size,
ep_size=self.ep_size,
load_format=self.load_format,
model_loader_extra_config=self.model_loader_extra_config,
)
get_context().set_server_args(server_args)
# Initialize distributed backend for model loading
# (must be done after server_args and model_config are available)
# Build model config first, then init distributed
model_config = ModelConfig(
model_path=self.model_path,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
dtype=self.dtype,
quantization=self.quantization,
)
# Build cache config fingerprint BEFORE loading the model.
# Loading may mutate hf_config.quantization_config (e.g. via
# process_weights_after_loading), which would produce a different
# hash than what the engine computes from the original config.
# ModelConfig always exposes hf_config/quantization directly;
# quantization_config is the only genuinely-optional attribute.
quant_config = getattr(model_config.hf_config, "quantization_config", None)
quant_method = get_quant_method_name(
self.quantization or model_config.quantization
)
if not quant_method and quant_config is not None:
quant_method = get_quant_method_name(quant_config)
self.config = CacheConfig(
model_path=self.model_path,
model_arch=(
model_config.hf_config.architectures[0]
if model_config.hf_config.architectures
else ""
),
tp_size=self.tp_size,
tp_rank=self.tp_rank,
pp_size=self.pp_size,
pp_rank=self.pp_rank,
dp_size=self.dp_size,
ep_size=self.ep_size,
quant_method=quant_method,
quant_config_hash=hash_quant_config(quant_config),
dtype=str(model_config.dtype),
revision=self.revision or "",
**compute_env_stamp(),
)
# Refuse to serve quant methods not verified to round-trip through pure
# IPC tensor export. Checked before loading so an unsupported model
# fails fast instead of after minutes of disk I/O.
check_ipc_quant_support(quant_method, quant_config, where="daemon")
# Initialize distributed backend (requires server_args + model_config)
self._init_distributed(server_args, model_config)
# Build load config
load_config = LoadConfig(
load_format=self.load_format,
model_loader_extra_config=self.model_loader_extra_config,
tp_rank=self.tp_rank,
)
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id} tp_rank={self.tp_rank}] "
f"Loading model from disk: {self.model_path}"
)
tic = time.perf_counter()
# Load model using DefaultModelLoader (includes TP sharding + quant post-process)
loader = get_model_loader(load_config=load_config, model_config=model_config)
self.model = loader.load_model(
model_config=model_config,
device_config=DeviceConfig(current_platform.device_type, self.gpu_id),
)
elapsed = time.perf_counter() - tic
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id} tp_rank={self.tp_rank}] "
f"Model loaded from disk in {elapsed:.2f}s"
)
# Ensure every post-processing kernel has retired before we export the
# memory: clients map these tensors read-only via IPC and would otherwise
# risk observing half-written weights.
current_platform.synchronize()
# Export all parameters and buffers as IPC handles
self._export_state()
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id} tp_rank={self.tp_rank}] "
f"Exported {len(self.state_entries)} tensors as IPC handles. "
f"Ready to serve."
)
@staticmethod
def _assert_ipc_compatible_allocator() -> None:
"""Reject allocator configs incompatible with CUDA IPC export.
The expandable-segments allocator returns memory that cannot be shared
through torch's _share_cuda_ handle, which would make the export fail
partway with an opaque error. Detect it up front and fail loud.
"""
for var in ("PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF"):
conf = os.environ.get(var, "")
for field in conf.split(","):
key, _, value = field.partition(":")
if (
key.strip() == "expandable_segments"
and value.strip().lower() == "true"
):
raise RuntimeError(
f"[WeightCacheDaemon] {var} sets expandable_segments:True, "
f"which is incompatible with CUDA IPC weight sharing: the "
f"expandable-segments allocator hands out memory that cannot "
f"be exported via _share_cuda_, so the IPC handle export "
f"would fail mid-way. Unset expandable_segments for the "
f"weight cache daemon process (it can stay enabled for the "
f"engine itself)."
)
def _export_state(self):
"""Export model parameters and buffers as CUDA IPC handles.
This includes both persistent buffers (in state_dict) and non-persistent
buffers (e.g. rotary embedding cos_sin_cache) so the engine can fully
reconstruct the model state via zero-copy IPC.
"""
self.state_entries.clear()
# remove_duplicate=False so tied weights are recognized as parameters
# under every name. state_dict() below emits both tied keys, and with the
# deduped set the duplicate would be mis-registered as a buffer, not a
# parameter, on the client.
param_names = set(
name for name, _ in self.model.named_parameters(remove_duplicate=False)
)
state_dict_names = set(self.model.state_dict().keys())
# Export all items from state_dict (parameters + persistent buffers)
for name, tensor in self.model.state_dict().items():
ipc_handle = MultiprocessingSerializer.serialize(
tensor.data, output_str=True
)
self.state_entries[name] = {
"handle": ipc_handle,
"shape": list(tensor.shape),
"dtype": str(tensor.dtype).replace("torch.", ""),
"is_param": name in param_names,
}
# Also export non-persistent buffers (not in state_dict but needed
# for inference, e.g. rotary embedding cos_sin_cache)
non_persistent_count = 0
for name, buf in self.model.named_buffers():
if name not in state_dict_names:
ipc_handle = MultiprocessingSerializer.serialize(
buf.data, output_str=True
)
self.state_entries[name] = {
"handle": ipc_handle,
"shape": list(buf.shape),
"dtype": str(buf.dtype).replace("torch.", ""),
"is_param": False,
}
non_persistent_count += 1
# Log total size
total_bytes = sum(
entry["handle"].__len__() if hasattr(entry["handle"], "__len__") else 0
for entry in self.state_entries.values()
)
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id}] "
f"Exported {len(self.state_entries)} tensors "
f"({non_persistent_count} non-persistent buffers), "
f"serialized handle size ~{total_bytes / 1024 / 1024:.1f} MB"
)
def serve(self):
"""Block and serve IPC handles over Unix socket."""
# Do NOT unlink an existing socket here: stale-file cleanup is the launch
# path's job (cleanup_stale_daemon_files refuses to remove a socket whose
# .ready still points at a live PID). A leftover live socket makes bind()
# fail loudly instead of silently stealing another daemon's socket.
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
old_umask = os.umask(0o177)
try:
sock.bind(self.socket_path)
finally:
os.umask(old_umask)
sock.listen(8)
sock.settimeout(1.0) # Allow periodic shutdown check
# Write ready file
with open(self.ready_path, "w") as f:
f.write(f"pid={os.getpid()}\n")
f.write(f"config={self.config.to_dict()}\n")
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id}] " f"Listening on {self.socket_path}"
)
self._running = True
def _signal_handler(signum, frame):
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id}] Received signal {signum}, shutting down"
)
self._running = False
signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)
try:
while self._running:
try:
conn, _ = sock.accept()
# The listen-socket timeout above only bounds accept(); the
# accepted connection is blocking by default. Since we serve
# connections serially, a client that connects but never sends
# (or dies mid-send) would block recv_msg forever and stall
# every other engine rank. Bound each exchange instead.
conn.settimeout(CLIENT_CONNECTION_TIMEOUT)
try:
self._handle_connection(conn)
except Exception as e:
logger.error(
f"[WeightCacheDaemon gpu={self.gpu_id}] "
f"Error handling connection: {e}",
exc_info=True,
)
finally:
conn.close()
except socket.timeout:
continue
finally:
sock.close()
if os.path.exists(self.socket_path):
os.unlink(self.socket_path)
if os.path.exists(self.ready_path):
os.unlink(self.ready_path)
logger.info(f"[WeightCacheDaemon gpu={self.gpu_id}] Shutdown complete")
def _handle_connection(self, conn: socket.socket):
"""Handle a single client connection."""
req = recv_msg(conn)
if req.get("type") == "query_config":
# Client asks for config without requesting handles
send_msg(conn, {"status": "ok", "config": self.config.to_dict()})
elif req.get("type") == "fetch_state":
# Client requests full state with IPC handles
engine_config = CacheConfig.from_dict(req["config"])
if not self.config.matches(engine_config):
# Log detailed mismatch info for debugging
daemon_dict = self.config.to_dict()
engine_dict = engine_config.to_dict()
mismatches = {
k: (daemon_dict.get(k), engine_dict.get(k))
for k in daemon_dict
if daemon_dict.get(k) != engine_dict.get(k)
}
logger.warning(
f"[WeightCacheDaemon gpu={self.gpu_id}] "
f"Config mismatch: {mismatches}"
)
send_msg(
conn, {"status": "mismatch", "daemon_config": self.config.to_dict()}
)
return
logger.info(
f"[WeightCacheDaemon gpu={self.gpu_id}] "
f"Serving {len(self.state_entries)} IPC handles to engine"
)
send_msg(
conn,
{
"status": "ok",
"config": self.config.to_dict(),
"entries": self.state_entries,
# PID so the client can watch daemon liveness: if this
# process dies while clients hold IPC mappings, their
# param.data (and any CUDA-graph-captured addresses) dangle.
"pid": os.getpid(),
},
)
elif req.get("type") == "ping":
send_msg(conn, {"status": "ok"})
else:
send_msg(
conn,
{
"status": "error",
"message": f"Unknown request type: {req.get('type')}",
},
)
def shutdown(self):
"""Release GPU memory and clean up."""
if dist.is_initialized():
dist.destroy_process_group()
if self.model is not None:
del self.model
self.model = None
self.state_entries.clear()
current_platform.empty_cache()
self._running = False
def run_weight_cache_daemon(
model_path: str,
gpu_id: int,
tp_size: int = 1,
tp_rank: int = 0,
pp_size: int = 1,
pp_rank: int = 0,
dp_size: int = 1,
ep_size: int = 1,
load_format: str = "auto",
dtype: str = "auto",
quantization: Optional[str] = None,
model_loader_extra_config: str = "{}",
trust_remote_code: bool = False,
revision: Optional[str] = None,
dist_init_method: Optional[str] = None,
):
"""Entry point for running a weight cache daemon process."""
logging.basicConfig(
level=logging.INFO,
format=f"%(asctime)s [Daemon gpu={gpu_id} tp_rank={tp_rank}] %(levelname)s %(message)s",
)
# Die if our parent (the engine or the standalone launcher that spawned us)
# dies, even on SIGKILL/OOM-kill. Without this an orphaned daemon keeps a
# full weight copy pinned in GPU memory and its live-PID .ready file blocks
# the next launch — the opposite of fast recovery.
from sglang.srt.utils import kill_itself_when_parent_died
kill_itself_when_parent_died()
daemon = WeightCacheDaemon(
model_path=model_path,
gpu_id=gpu_id,
tp_size=tp_size,
tp_rank=tp_rank,
pp_size=pp_size,
pp_rank=pp_rank,
dp_size=dp_size,
ep_size=ep_size,
load_format=load_format,
dtype=dtype,
quantization=quantization,
model_loader_extra_config=model_loader_extra_config,
trust_remote_code=trust_remote_code,
revision=revision,
dist_init_method=dist_init_method,
)
daemon.load()
daemon.serve()
def launch_weight_cache_daemons(
model_path: str,
tp_size: int = 1,
pp_size: int = 1,
dp_size: int = 1,
ep_size: int = 1,
nnodes: int = 1,
node_rank: int = 0,
base_gpu_id: int = 0,
gpu_id_step: int = 1,
load_format: str = "auto",
dtype: str = "auto",
quantization: Optional[str] = None,
model_loader_extra_config: str = "{}",
trust_remote_code: bool = False,
revision: Optional[str] = None,
dist_init_method: Optional[str] = None,
timeout: int = 1800,
force: bool = False,
):
"""Launch weight cache daemon processes for this node's PP×TP ranks.
For single-node (nnodes=1): spawns pp_size * tp_size daemons.
For multi-node (nnodes>1): spawns this node's share of PP×TP daemons,
mapping local gpu_id to the correct global (pp_rank, tp_rank).
Uses subprocess.Popen instead of multiprocessing.Process to avoid
initializing CUDA in the parent process, which can degrade CUDA IPC
performance in child processes.
Usage (single-node):
python -m sglang.srt.weight_cache.daemon \\
--model-path /path/to/model --tp-size 4
Usage (multi-node, run on each node):
# Node 0:
python -m sglang.srt.weight_cache.daemon \\
--model-path /path/to/model --tp-size 16 \\
--nnodes 2 --node-rank 0 \\
--dist-init-method tcp://node0-ip:29500
# Node 1:
python -m sglang.srt.weight_cache.daemon \\
--model-path /path/to/model --tp-size 16 \\
--nnodes 2 --node-rank 1 \\
--dist-init-method tcp://node0-ip:29500
"""
import socket as sock_mod
import subprocess
import sys
# Replicate _calculate_rank_ranges logic from engine.py
pp_size_per_node = max(pp_size // nnodes, 1)
nnodes_per_pp_rank = max(nnodes // pp_size, 1)
pp_rank_range = range(
pp_size_per_node * (node_rank // nnodes_per_pp_rank),
pp_size_per_node * (node_rank // nnodes_per_pp_rank + 1),
)
nnodes_per_tp_group = nnodes_per_pp_rank
tp_size_per_node = tp_size // nnodes_per_tp_group
tp_rank_range = range(
tp_size_per_node * (node_rank % nnodes_per_tp_group),
tp_size_per_node * (node_rank % nnodes_per_tp_group + 1),
)
if nnodes > 1 and dist_init_method is None:
raise ValueError(
"dist_init_method is required for multi-node weight cache daemons. "
"Use --dist-init-method tcp://<node0-ip>:<port> to specify the "
"rendezvous address accessible from all nodes."
)
# Auto-allocate a free port for the distributed init method
if dist_init_method is None:
with sock_mod.socket(sock_mod.AF_INET, sock_mod.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
free_port = s.getsockname()[1]
dist_init_method = f"tcp://127.0.0.1:{free_port}"
python_path = sys.executable
daemon_module = "sglang.srt.weight_cache.daemon"
# Validate and clean up stale .ready/.sock files from prior runs.
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
cleanup_stale_daemon_files(global_rank, force=force)
procs = []
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
gpu_id = compute_local_gpu_id(
pp_rank,
tp_rank,
pp_size_per_node,
tp_size_per_node,
base_gpu_id=base_gpu_id,
gpu_id_step=gpu_id_step,
)
cmd = [
python_path,
"-m",
daemon_module,
"--model-path",
model_path,
"--gpu-id",
str(gpu_id),
"--tp-size",
str(tp_size),
"--tp-rank",
str(tp_rank),
"--dp-size",
str(dp_size),
"--ep-size",
str(ep_size),
"--pp-size",
str(pp_size),
"--pp-rank",
str(pp_rank),
"--load-format",
load_format,
"--dtype",
dtype,
"--dist-init-method",
dist_init_method,
]
if quantization:
cmd += ["--quantization", quantization]
if model_loader_extra_config and model_loader_extra_config != "{}":
cmd += ["--model-loader-extra-config", model_loader_extra_config]
if trust_remote_code:
cmd += ["--trust-remote-code"]
if revision:
cmd += ["--revision", revision]
proc = subprocess.Popen(cmd)
procs.append(proc)
logger.info(
f"Launched weight cache daemon gpu={gpu_id} "
f"pp_rank={pp_rank} tp_rank={tp_rank} pid={proc.pid}"
)
# Wait for all daemons on this node to become ready
num_daemons = len(procs)
check_interval = 2
start_time = time.time()
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
global_rank = compute_global_rank(tp_size, pp_rank, tp_rank)
ready_path = get_ready_path(global_rank)
while not os.path.exists(ready_path):
time.sleep(check_interval)
if time.time() - start_time > timeout:
logger.error(
f"Weight cache daemon pp_rank={pp_rank} tp_rank={tp_rank} "
f"did not become ready within {timeout}s"
)
for p in procs:
p.terminate()
raise TimeoutError(
f"Weight cache daemon pp_rank={pp_rank} tp_rank={tp_rank} "
f"did not become ready within {timeout}s"
)
# Check if any daemon exited prematurely
for p in procs:
retcode = p.poll()
if retcode is not None:
logger.error(
f"Weight cache daemon exited prematurely "
f"with code {retcode}"
)
for other in procs:
if other.poll() is None:
other.terminate()
raise RuntimeError(
f"Weight cache daemon exited prematurely "
f"with code {retcode}"
)
logger.info(
f"Weight cache daemon pp_rank={pp_rank} tp_rank={tp_rank} is ready"
)
logger.info(
f"All {num_daemons} weight cache daemons on node {node_rank} are ready "
f"(pp_ranks={pp_rank_range.start}..{pp_rank_range.stop - 1}, "
f"tp_ranks={tp_rank_range.start}..{tp_rank_range.stop - 1}, "
f"dist_init_method={dist_init_method})"
)
# Monitor daemons — poll all of them and, the moment any one exits,
# terminate the rest and raise. A serial proc.wait() would not notice a
# mid-list death (e.g. procs[1] dying while procs[0] is still alive) until
# the earlier proc happened to exit, and it never surfaced the failure.
exited = None
try:
while exited is None:
for proc in procs:
if proc.poll() is not None:
exited = proc
break
else:
time.sleep(1)
continue
except KeyboardInterrupt:
logger.info("Received KeyboardInterrupt, shutting down daemons")
finally:
for proc in procs:
if proc.poll() is None:
proc.terminate()
for proc in procs:
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
proc.kill()
logger.info("All weight cache daemons have been terminated")
if exited is not None:
raise RuntimeError(
f"Weight cache daemon (pid={exited.pid}) exited with code "
f"{exited.returncode}; terminated the remaining daemons."
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="SGLang Weight Cache Daemon")
parser.add_argument("--model-path", required=True, help="Path to model weights")
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
parser.add_argument(
"--gpu-id",
type=int,
default=None,
help="GPU device ID for a single daemon. "
"If omitted, launches daemons for all TP ranks (0..tp_size-1).",
)
parser.add_argument(
"--tp-rank",
type=int,
default=None,
help="TP rank for a single daemon. "
"If omitted, launches daemons for all TP ranks.",
)
parser.add_argument("--dp-size", type=int, default=1, help="Data parallel size")
parser.add_argument("--ep-size", type=int, default=1, help="Expert parallel size")
parser.add_argument("--pp-size", type=int, default=1, help="Pipeline parallel size")
parser.add_argument("--pp-rank", type=int, default=0, help="Pipeline parallel rank")
parser.add_argument("--nnodes", type=int, default=1, help="Total number of nodes")
parser.add_argument(
"--base-gpu-id",
type=int,
default=0,
help="GPU id of this node's first rank (mirrors the engine's "
"--base-gpu-id). Used to place daemons on the same GPUs the engine "
"ranks will use.",
)
parser.add_argument(
"--gpu-id-step",
type=int,
default=1,
help="Stride between consecutive ranks' GPU ids (mirrors the engine's "
"--gpu-id-step).",
)
parser.add_argument(
"--node-rank",
type=int,
default=0,
help="Rank of this node (0-indexed). Required for multi-node.",
)
parser.add_argument("--load-format", default="auto", help="Weight load format")
parser.add_argument("--dtype", default="auto", help="Model dtype")
parser.add_argument("--quantization", default=None, help="Quantization method")
parser.add_argument(
"--model-loader-extra-config",
default="{}",
help="Extra config for model loader (JSON string)",
)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument("--revision", default=None, help="Model revision")
parser.add_argument(
"--dist-init-method",
default=None,
help="Distributed init method (e.g. tcp://node0-ip:29500). "
"Auto-assigned for single-node when launching all ranks. "
"Required for multi-node (nnodes > 1) and must be accessible "
"from all nodes. Also required for tp_size > 1 when launching "
"a single rank.",
)
parser.add_argument(
"--timeout",
type=int,
default=1800,
help="Timeout in seconds to wait for all daemons to become ready (default: 1800)",
)
parser.add_argument(
"--force",
action="store_true",
help="Take over a rank whose .ready file still points at a live PID by "
"killing that daemon (use to reclaim a wedged/orphaned daemon).",
)
args = parser.parse_args()
if args.gpu_id is not None or args.tp_rank is not None:
# Single-rank mode: launch one daemon for the specified rank
gpu_id = args.gpu_id if args.gpu_id is not None else args.tp_rank
tp_rank = args.tp_rank if args.tp_rank is not None else args.gpu_id
# Refuse to clobber a live daemon already holding this rank (mirrors the
# multi-rank launcher and the engine path, which the multi-rank spawns
# of this same entrypoint rely on). --force kills and takes over.
cleanup_stale_daemon_files(
compute_global_rank(args.tp_size, args.pp_rank, tp_rank),
force=args.force,
)
run_weight_cache_daemon(
model_path=args.model_path,
gpu_id=gpu_id,
tp_size=args.tp_size,
tp_rank=tp_rank,
pp_size=args.pp_size,
pp_rank=args.pp_rank,
dp_size=args.dp_size,
ep_size=args.ep_size,
load_format=args.load_format,
dtype=args.dtype,
quantization=args.quantization,
model_loader_extra_config=args.model_loader_extra_config,
trust_remote_code=args.trust_remote_code,
revision=args.revision,
dist_init_method=args.dist_init_method,
)
else:
# Multi-rank mode: launch daemons for this node's TP ranks
launch_weight_cache_daemons(
model_path=args.model_path,
tp_size=args.tp_size,
pp_size=args.pp_size,
dp_size=args.dp_size,
ep_size=args.ep_size,
nnodes=args.nnodes,
node_rank=args.node_rank,
base_gpu_id=args.base_gpu_id,
gpu_id_step=args.gpu_id_step,
load_format=args.load_format,
dtype=args.dtype,
quantization=args.quantization,
model_loader_extra_config=args.model_loader_extra_config,
trust_remote_code=args.trust_remote_code,
revision=args.revision,
dist_init_method=args.dist_init_method,
timeout=args.timeout,
force=args.force,
)
@@ -0,0 +1,565 @@
# SPDX-License-Identifier: Apache-2.0
"""IPC Model Loader — loads model weights from a Weight Cache Daemon via CUDA IPC.
Zero-copy mode: param.data points directly to IPC-mapped GPU memory. Only 1x GPU
memory needed — engine and daemon share the same physical GPU memory via CUDA IPC.
Engine depends on daemon staying alive.
"""
import logging
import os
import signal
import stat
import threading
import time
from typing import Optional
import torch
import torch.nn as nn
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.model_loader.loader import (
BaseModelLoader,
_initialize_model,
)
from sglang.srt.utils import MultiprocessingSerializer
from .protocol import (
CacheConfig,
check_ipc_quant_support,
compute_env_stamp,
get_quant_method_name,
hash_quant_config,
recv_msg,
send_msg,
)
logger = logging.getLogger(__name__)
# How often the client polls the serving daemon's PID for liveness.
_DAEMON_LIVENESS_POLL_INTERVAL = 5.0
class IpcModelLoader(BaseModelLoader):
"""Load model weights from a Weight Cache Daemon via CUDA IPC handles.
In daemon mode (weight_cache_mode="daemon"), the engine and daemon share
the same GPU. Falling back to disk loading would cause OOM because both
processes would hold weights on the same GPU. Therefore, daemon mode
raises an error if the daemon is unavailable instead of falling back.
In client mode, disk fallback is allowed ONLY when the daemon is genuinely
absent (its Unix socket file does not exist). Every other failure is a hard
error rather than a silent fallback, so a broken IPC path never masquerades
as a healthy (but slow, disk-loaded) server:
- socket file missing -> fall back to disk load
- connection refused -> raise (daemon crashed after binding)
- CacheConfig mismatch -> raise (do NOT disk-load on a shared GPU
holding a different config's weights;
also surfaces fingerprint drift bugs)
- any protocol / transfer error -> raise
See _fetch_from_cache for the authoritative fallback-vs-raise contract.
"""
def __init__(
self,
load_config: LoadConfig,
socket_path: str,
fallback_loader_cls=None,
weight_cache_mode: str = "client",
fallback_load_format: str = "auto",
):
super().__init__(load_config)
self.socket_path = socket_path
self.weight_cache_mode = weight_cache_mode
self._fallback_loader_cls = fallback_loader_cls
self._fallback_load_format = fallback_load_format
def load_model(
self,
*,
model_config,
device_config,
) -> nn.Module:
"""Load model weights from the weight cache daemon.
In daemon mode, raises RuntimeError if the daemon is unavailable
(fallback to disk loading would cause OOM on shared GPUs).
In client mode, falls back to DefaultModelLoader.
"""
tic = time.perf_counter()
# Hard-gate unsupported quant methods before touching the daemon, so an
# unsupported model fails explicitly instead of silently disk-loading
# (client mode) or serving wrong-numerics IPC weights. Checked here so
# it applies regardless of whether the daemon is reachable.
quant_method, engine_quant_config = self._resolve_engine_quant(model_config)
check_ipc_quant_support(quant_method, engine_quant_config, where="client")
# Try to fetch state from daemon
cache_data = self._fetch_from_cache(model_config)
if cache_data is None:
if self.weight_cache_mode == "daemon":
raise RuntimeError(
f"[IpcModelLoader] Weight cache daemon not available at "
f"{self.socket_path}. In daemon mode, fallback to disk "
f"loading is disabled because the daemon process already "
f"holds weights on the same GPU — loading from disk would "
f"cause OOM. Please ensure the weight cache daemon is "
f"running and the config matches."
)
logger.warning(
"[IpcModelLoader] Weight cache not available or config mismatch, "
"falling back to disk load"
)
return self._fallback_load(model_config, device_config)
entries = cache_data["entries"]
logger.info(
f"[IpcModelLoader] Fetched {len(entries)} IPC handles from daemon "
f"in {time.perf_counter() - tic:.2f}s"
)
from sglang.srt.model_loader.loader import (
_get_quantization_config,
)
quant_config = _get_quantization_config(model_config, self.load_config)
model = self._load_zero_copy_mode(
model_config,
device_config,
entries,
quant_config,
)
# Skip _post_load_weights: the daemon already ran
# process_weights_after_loading on the weights before exporting
# IPC handles. Running it again would double-process (e.g.,
# re-quantize already-quantized weights), corrupting tensor data.
# Rebuild stale tensor views. Some modules store tensor views as
# plain attributes (not parameters/buffers) during __init__. When
# the model is initialized on meta device and then weights are
# replaced via IPC mapping, these views still point to the old
# meta storage. We must recreate them from the now-valid tensors.
self._rebuild_stale_views(model)
# The model now points into the daemon's GPU memory via CUDA IPC. If the
# daemon dies, those pointers dangle, so watch it and fail loud.
self._start_daemon_liveness_watchdog(cache_data.get("pid"))
logger.info(
f"[IpcModelLoader] Loaded model via IPC (mode={self.weight_cache_mode}), "
f"total={time.perf_counter() - tic:.2f}s"
)
return model.eval()
def _start_daemon_liveness_watchdog(self, daemon_pid: Optional[int]) -> None:
"""Fail loud if the serving daemon dies while we hold its weights.
In both client and (engine-spawned) daemon mode, the model's param.data
points into the daemon's GPU memory via CUDA IPC, and CUDA graphs may
capture those addresses. If the daemon exits, the pointers dangle:
forward passes would read freed GPU memory -> illegal-address crashes or
silent garbage. There is no safe in-place recovery, so a background
thread polls the daemon PID and, on death, SIGKILLs this process with a
clear message instead of letting it serve corrupt results.
"""
if not daemon_pid or daemon_pid <= 0:
logger.warning(
"[IpcModelLoader] Daemon did not report a PID; skipping the "
"daemon-liveness watchdog. A daemon crash will not be detected."
)
return
def _daemon_alive(pid: int) -> bool:
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except PermissionError:
return True # exists but owned by another user
return True
def _watch() -> None:
while True:
time.sleep(_DAEMON_LIVENESS_POLL_INTERVAL)
if not _daemon_alive(daemon_pid):
logger.critical(
f"[IpcModelLoader] Weight cache daemon (pid={daemon_pid}) "
f"died while this engine holds its weights via CUDA IPC. "
f"The mapped weight pointers are now dangling; continuing "
f"would read freed GPU memory. Terminating this process."
)
os.kill(os.getpid(), signal.SIGKILL)
return
threading.Thread(
target=_watch, name="weight-cache-daemon-watchdog", daemon=True
).start()
logger.info(
f"[IpcModelLoader] Started daemon-liveness watchdog for pid={daemon_pid}"
)
def _resolve_engine_quant(self, model_config):
"""Return (quant_method, quant_config) matching the daemon's fingerprint.
Shared by the IPC allowlist gate and the CacheConfig fingerprint so the
two can never drift apart. ModelConfig always exposes
hf_config/quantization directly; quantization_config is the only
genuinely-optional attribute.
"""
quant_config = getattr(model_config.hf_config, "quantization_config", None)
quant_method = get_quant_method_name(model_config.quantization)
if not quant_method and quant_config is not None:
quant_method = get_quant_method_name(quant_config)
return quant_method, quant_config
@staticmethod
def _rebuild_stale_views(model):
"""Rebuild tensor views that went stale after IPC weight replacement.
RadixLinearAttention.conv_weights is a view of conv1d.weight created
during __init__. After IPC mapping replaces conv1d.weight with a new
tensor, the old view still points to meta-device storage. Recreate
it from the now-valid parameter.
"""
try:
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
except ImportError:
return
count = 0
for _, module in model.named_modules():
conv1d = getattr(module, "conv1d", None)
attn = getattr(module, "attn", None)
if conv1d is not None and isinstance(attn, RadixLinearAttention):
if hasattr(conv1d, "weight") and conv1d.weight is not None:
attn.conv_weights = conv1d.weight.view(
conv1d.weight.size(0), conv1d.weight.size(2)
)
if hasattr(conv1d, "bias") and conv1d.bias is not None:
attn.bias = conv1d.bias
count += 1
if count > 0:
logger.info(f"[IpcModelLoader] Rebuilt {count} stale conv_weights views")
@staticmethod
def _set_module_tensor(model, name, tensor, is_param=True):
"""Replace or register a parameter/buffer in the model by its full dotted name.
This is necessary because setting param.data on a meta-device tensor
raises a type mismatch error (meta and CUDA tensors have incompatible
dispatch keys). Instead, we walk the module tree and use setattr to
replace the entire parameter/buffer object.
If the attribute already exists as a parameter/buffer, it is replaced.
If it doesn't exist (e.g. post-quantization params like weight_scale),
it is registered as a new parameter or buffer.
"""
parts = name.split(".")
obj = model
for part in parts[:-1]:
obj = getattr(obj, part)
leaf_name = parts[-1]
if is_param:
# requires_grad=False: the IPC memory is shared/read-only and SGLang
# is inference-only, so autograd must never write into it.
new_param = nn.Parameter(tensor, requires_grad=False)
setattr(obj, leaf_name, new_param)
else:
# register_buffer raises KeyError if the name already exists as a
# parameter or plain attribute (not a buffer). This happens when
# process_weights_after_loading converts a parameter to a buffer
# (e.g. Mamba's A_log). Remove the old attribute first.
if leaf_name in obj._parameters:
del obj._parameters[leaf_name]
elif hasattr(obj, leaf_name) and leaf_name not in obj._buffers:
delattr(obj, leaf_name)
obj.register_buffer(leaf_name, tensor)
def _load_zero_copy_mode(
self,
model_config,
device_config,
entries,
quant_config,
) -> nn.Module:
"""Zero-copy load: map IPC tensors directly as param.data.
The model is initialized on the meta device (no memory allocation),
then each parameter's data is replaced with the IPC-mapped GPU tensor.
The engine and daemon share the same physical GPU memory via CUDA IPC.
"""
from sglang.srt.model_loader.utils import set_default_torch_dtype
# Initialize model on meta device to avoid any GPU/CPU memory allocation.
# This creates the model structure with the correct parameter shapes/dtypes
# but without allocating actual storage.
with set_default_torch_dtype(model_config.dtype):
with torch.device("meta"):
model = _initialize_model(
model_config,
self.load_config,
quant_config,
)
# Build lookup dicts of existing parameter/buffer names in the
# meta-device model. Post-quantization parameters (e.g. weight_scale
# from FP8) are created by process_weights_after_loading, which the
# daemon already ran. These params exist in the daemon's entries but
# NOT in the meta-device model — we must register them as new attrs.
# Use dicts (not sets) so we can do O(1) shape/dtype validation
# without re-traversing the model tree on every lookup.
# remove_duplicate=False mirrors the daemon's export (which keys tied
# weights under every name) so a tied parameter is recognized under all
# of its names here too.
existing_params = {
name: param
for name, param in model.named_parameters(remove_duplicate=False)
}
existing_buffers = {name: buf for name, buf in model.named_buffers()}
existing_names = set(existing_params) | set(existing_buffers)
imported_refs = []
imported_count = 0
mismatched = []
new_params_count = 0
map_tic = time.perf_counter()
# Iterate over ALL daemon entries (not just model params/buffers).
# This ensures post-quantization parameters (weight_scale, etc.)
# that were created by process_weights_after_loading are also mapped.
for name, entry in entries.items():
imported_tensor = MultiprocessingSerializer.deserialize(entry["handle"])
is_param = entry.get("is_param", True)
if name in existing_names:
# Existing parameter/buffer — validate shape/dtype
if name in existing_params:
ref_param = existing_params[name]
else:
ref_param = existing_buffers[name]
if (
imported_tensor.shape != ref_param.shape
or imported_tensor.dtype != ref_param.dtype
):
mismatched.append(
f" {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} "
f"vs model={ref_param.shape}/{ref_param.dtype}"
)
del imported_tensor
continue
# Replace or register the tensor in the model
self._set_module_tensor(model, name, imported_tensor, is_param=is_param)
imported_refs.append(imported_tensor)
imported_count += 1
if name not in existing_names:
new_params_count += 1
if mismatched:
raise RuntimeError(
f"[IpcModelLoader] {len(mismatched)} tensor(s) have shape/dtype "
f"mismatch between the IPC daemon and the meta-initialized model. "
f"The quantization method passed the IPC allowlist gate "
f"(check_ipc_quant_support), so this is NOT an unsupported-quant "
f"case — it indicates the daemon's weight fingerprint is "
f"incomplete or the daemon/client configs drifted (a bug to fix), "
f"not merely uninitialized weights:\n" + "\n".join(mismatched)
)
# After mapping every daemon entry, any tensor still on the meta device
# is one the daemon did NOT provide. Filling it with torch.empty() would
# hand the model uninitialized GPU memory — silently producing wrong
# output, the worst failure mode for a load path. Hard-error and list the
# offenders instead.
#
# The daemon exports the full state_dict AND non-persistent buffers
# (e.g. rotary embedding cos_sin_cache), so a correct setup leaves nothing
# on meta here. A non-empty list means the daemon's export is incomplete,
# or the model has a genuinely-recomputable buffer that must be recomputed
# explicitly (not filled with garbage) — add that handling here if needed.
still_on_meta_params = [
name
for name, param in model.named_parameters()
if param.device.type == "meta"
]
still_on_meta_buffers = [
name for name, buf in model.named_buffers() if buf.device.type == "meta"
]
if still_on_meta_params or still_on_meta_buffers:
raise RuntimeError(
f"[IpcModelLoader] After IPC mapping, "
f"{len(still_on_meta_params)} parameter(s) and "
f"{len(still_on_meta_buffers)} buffer(s) remain on the meta device "
f"— the daemon did not export them. Refusing to fill them with "
f"uninitialized memory, which would silently produce wrong output. "
f"This means the daemon's export is incomplete, or a recomputable "
f"buffer needs explicit recompute logic here.\n"
f" params: {still_on_meta_params[:10]}"
f"{'...' if len(still_on_meta_params) > 10 else ''}\n"
f" buffers: {still_on_meta_buffers[:10]}"
f"{'...' if len(still_on_meta_buffers) > 10 else ''}"
)
map_elapsed = time.perf_counter() - map_tic
# Stash IPC refs on the model to prevent GC (which would unmap the memory)
if imported_refs:
model._ipc_imported_tensors = imported_refs
logger.info(
f"[IpcModelLoader] Zero-copy: mapped {imported_count} tensors "
f"({new_params_count} new post-quant), time={map_elapsed:.3f}s"
)
return model
def _fetch_from_cache(self, model_config) -> Optional[dict]:
"""Connect to daemon, validate config, fetch IPC handles.
Returns the daemon response dict on success, None if the daemon is
genuinely absent (socket file doesn't exist). Raises on all other
failures so they are never silently swallowed as a disk-load fallback.
"""
import socket as socket_mod
# Only connect to a real socket node owned by us: reject a symlink, a
# plain file, or another user's socket planted at this /tmp path. An
# absent socket means no daemon -> fall back to disk (return None).
try:
st = os.lstat(self.socket_path)
except FileNotFoundError:
logger.info(
f"[IpcModelLoader] Daemon socket not found at {self.socket_path}."
)
return None
if not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid():
raise RuntimeError(
f"[IpcModelLoader] Refusing to connect: {self.socket_path} is not "
f"a socket owned by this user."
)
sock = socket_mod.socket(socket_mod.AF_UNIX, socket_mod.SOCK_STREAM)
try:
sock.settimeout(30)
sock.connect(self.socket_path)
except FileNotFoundError:
# Raced: socket removed between lstat and connect -> treat as absent.
sock.close()
return None
except ConnectionRefusedError:
sock.close()
raise RuntimeError(
f"[IpcModelLoader] Daemon socket exists at {self.socket_path} but "
f"refused the connection. The daemon may have crashed after "
f"creating the socket. Check daemon logs."
)
except Exception as e:
sock.close()
raise RuntimeError(
f"[IpcModelLoader] Failed to connect to daemon at "
f"{self.socket_path}: {e}"
) from e
try:
# Build engine's config fingerprint
from sglang.srt.runtime_context import get_parallel
ps = get_parallel()
tp_size = ps.tp_size
tp_rank = ps.tp_rank
pp_size = ps.pp_size
pp_rank = ps.pp_rank
ep_size = ps.moe_ep_size
from sglang.srt.runtime_context import get_server_args
dp_size = get_server_args().dp_size
quant_method, quant_config = self._resolve_engine_quant(model_config)
engine_config = CacheConfig(
model_path=model_config.model_path,
model_arch=(
model_config.hf_config.architectures[0]
if model_config.hf_config.architectures
else ""
),
tp_size=tp_size,
tp_rank=tp_rank,
pp_size=pp_size,
pp_rank=pp_rank,
dp_size=dp_size,
ep_size=ep_size,
quant_method=quant_method,
quant_config_hash=hash_quant_config(quant_config),
dtype=str(model_config.dtype),
revision=model_config.revision or "",
**compute_env_stamp(),
)
logger.info(
f"[IpcModelLoader] Requesting weights from daemon at "
f"{self.socket_path} with config: "
f"model={engine_config.model_path}, "
f"arch={engine_config.model_arch}, "
f"tp={engine_config.tp_size}/{engine_config.tp_rank}, "
f"quant={engine_config.quant_method}, "
f"dtype={engine_config.dtype}"
)
send_msg(sock, {"type": "fetch_state", "config": engine_config.to_dict()})
result = recv_msg(sock)
if result.get("status") != "ok":
daemon_config = result.get("daemon_config", {})
raise RuntimeError(
f"[IpcModelLoader] Daemon config mismatch!\n"
f" Engine config: {engine_config.to_dict()}\n"
f" Daemon config: {daemon_config}"
)
return result
except RuntimeError:
raise
except Exception as e:
raise RuntimeError(
f"[IpcModelLoader] Error communicating with daemon at "
f"{self.socket_path}: {e}"
) from e
finally:
sock.close()
def _fallback_load(self, model_config, device_config) -> nn.Module:
"""Fall back to DefaultModelLoader for disk-based loading."""
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.model_loader.loader import DefaultModelLoader
fallback_config = LoadConfig(
load_format=self._fallback_load_format,
download_dir=self.load_config.download_dir,
model_loader_extra_config=self.load_config.model_loader_extra_config,
tp_rank=self.load_config.tp_rank,
)
loader_cls = self._fallback_loader_cls or DefaultModelLoader
fallback = loader_cls(fallback_config)
return fallback.load_model(
model_config=model_config, device_config=device_config
)
def download_model(self, model_config) -> None:
"""No-op: daemon handles its own model downloading."""
pass
+380
View File
@@ -0,0 +1,380 @@
# SPDX-License-Identifier: Apache-2.0
"""Protocol definitions for the weight cache daemon.
Defines CacheConfig for validation and socket message protocol helpers.
"""
import hashlib
import json
import logging
import os
import pickle
import signal
import struct
from typing import Any, Dict, Optional
import msgspec
from sglang.srt.utils.common import safe_pickle_loads
logger = logging.getLogger(__name__)
# Socket path template for weight cache daemons (keyed by global rank
# = tp_size * pp_rank + tp_rank, so multi-node / multi-PP don't collide)
WEIGHT_CACHE_SOCKET_TEMPLATE = "/tmp/sglang_weight_cache_rank{global_rank}.sock"
# Ready file template — daemon writes this after loading completes
WEIGHT_CACHE_READY_TEMPLATE = "/tmp/sglang_weight_cache_rank{global_rank}.ready"
class CacheConfig(msgspec.Struct):
"""Fingerprint of the cached weights. Used to validate compatibility
between a daemon's cached state and a requesting engine process.
Any mismatch triggers a fallback to disk loading.
"""
model_path: str
model_arch: str
tp_size: int
tp_rank: int
pp_size: int
pp_rank: int
dp_size: int
ep_size: int
quant_method: str # e.g. "fp8", "gptq_marlin", "" for unquantized
quant_config_hash: str # SHA-256 hash of quantization config
dtype: str # e.g. "torch.float16"
revision: str # model revision the weights were loaded from ("" if unset)
# Environment stamp: a daemon and a client that ran different post-processing
# branches (different GPU compute capability or torch/kernel version) can
# produce incompatible weights that would map cleanly yet serve garbage.
# Comparing these turns that into a clean mismatch. See compute_env_stamp().
device_capability: str # local compute capability, e.g. "8.0" ("" if N/A)
torch_version: str # torch.__version__ of the process that built the weights
def matches(self, other: "CacheConfig") -> bool:
"""Check if two configs are compatible for weight sharing."""
return self == other
def to_dict(self) -> Dict[str, Any]:
return {f: getattr(self, f) for f in self.__struct_fields__}
@classmethod
def from_dict(cls, d: Dict[str, Any]) -> "CacheConfig":
return cls(**d)
def hash_quant_config(quant_config: Any) -> str:
"""Compute a stable hash of the quantization config.
Avoids str()/repr() on arbitrary objects because those embed memory
addresses (e.g. "at 0x7f..."), producing different hashes across
processes and causing permanent config mismatch.
"""
if quant_config is None:
return ""
try:
if hasattr(quant_config, "to_dict"):
config_str = json.dumps(quant_config.to_dict(), sort_keys=True)
elif isinstance(quant_config, dict):
config_str = json.dumps(quant_config, sort_keys=True)
elif hasattr(quant_config, "__dict__"):
config_str = (
type(quant_config).__name__
+ ":"
+ json.dumps(
{
k: v
for k, v in sorted(quant_config.__dict__.items())
if not k.startswith("_")
and isinstance(
v, (str, int, float, bool, type(None), list, dict)
)
},
sort_keys=True,
)
)
else:
config_str = type(quant_config).__name__
return hashlib.sha256(config_str.encode()).hexdigest()
except Exception:
config_str = type(quant_config).__name__
return hashlib.sha256(config_str.encode()).hexdigest()
def get_quant_method_name(quant_config: Any) -> str:
"""Extract the quantization method name from config."""
if quant_config is None:
return ""
if isinstance(quant_config, str):
return quant_config
if hasattr(quant_config, "get_name"):
return quant_config.get_name()
if hasattr(quant_config, "name"):
return quant_config.name
return type(quant_config).__name__
# ---------------------------------------------------------------------------
# IPC quantization-method allowlist
# ---------------------------------------------------------------------------
#
# CUDA IPC zero-copy sharing exports ONLY raw tensor data, so it is correct only
# when process_weights_after_loading's entire effect is captured by that data.
# Methods that stamp Python-side metadata (e.g. block-FP8's format_ue8m0) or
# repack/transpose weights into shapes the meta-init client can't reproduce
# (per-tensor FP8, Marlin, AWQ/GPTQ) would serve silently-wrong numerics. Only
# methods verified to round-trip through pure tensor export are allowed; every
# other method hard-errors. Extend the registry below only after verifying a
# method end-to-end.
class UnsupportedQuantForIPCError(RuntimeError):
"""Raised when a quantization method is not on the verified allowlist for
CUDA IPC zero-copy weight sharing."""
def _get_quant_field(quant_config: Any, key: str) -> Any:
"""Read a field from a quant config that may be a dict or an object."""
if quant_config is None:
return None
if isinstance(quant_config, dict):
return quant_config.get(key)
return getattr(quant_config, key, None)
def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:
"""Only block-wise FP8 is verified.
Block-wise FP8 (weight_block_size set) preserves weight shape and the only
post-load metadata it stamps is accounted for. Per-tensor FP8 transposes
`layer.weight` during post-processing, a shape change the meta-init client
cannot reproduce, so it is not supported.
"""
return _get_quant_field(quant_config, "weight_block_size") is not None
# quant_method name -> predicate(quant_config) -> bool (True == verified safe).
# A method absent from this registry is unsupported and hard-errors.
IPC_QUANT_ALLOWLIST = {
"": lambda _quant_config: True, # unquantized
"fp8": _fp8_round_trips_via_ipc, # only block-wise FP8 verified
}
def is_ipc_quant_supported(quant_method: str, quant_config: Any) -> bool:
"""Return True if `quant_method` is verified safe for IPC zero-copy sharing."""
predicate = IPC_QUANT_ALLOWLIST.get(quant_method)
if predicate is None:
return False
return bool(predicate(quant_config))
def check_ipc_quant_support(
quant_method: str, quant_config: Any, *, where: str
) -> None:
"""Hard-error unless `quant_method` is verified safe for IPC zero-copy sharing.
`where` is a short tag (e.g. "daemon"/"client") used only in the error
message. Raises UnsupportedQuantForIPCError with an actionable message.
"""
if is_ipc_quant_supported(quant_method, quant_config):
return
verified = ", ".join(
(repr(m) if m else "'' (unquantized)") for m in IPC_QUANT_ALLOWLIST
)
raise UnsupportedQuantForIPCError(
f"[weight_cache:{where}] quantization method {quant_method!r} is not "
f"verified for CUDA IPC zero-copy weight sharing. Its "
f"process_weights_after_loading may stamp Python-side metadata "
f"(e.g. format_ue8m0) or repack/transpose weights into shapes the "
f"meta-initialized client cannot reproduce, which would silently serve "
f"wrong-numerics weights. Verified methods: {verified}. Note: FP8 is "
f"only verified for block-wise configs (weight_block_size set), not "
f"per-tensor FP8. Disable the weight cache (--weight-cache-mode off) "
f"for this model."
)
# ---------------------------------------------------------------------------
# Socket protocol helpers
# ---------------------------------------------------------------------------
MAX_MSG_SIZE = 256 * 1024 * 1024 # 256 MiB
def send_msg(sock, obj: Any) -> None:
"""Send a length-prefixed pickled message over a socket."""
data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
header = struct.pack("!I", len(data))
sock.sendall(header + data)
def recv_msg(sock) -> Any:
"""Receive a length-prefixed pickled message from a socket."""
header = _recv_exact(sock, 4)
if header is None:
raise ConnectionError("Connection closed while reading message header")
length = struct.unpack("!I", header)[0]
if length > MAX_MSG_SIZE:
raise ValueError(f"Message size {length} exceeds {MAX_MSG_SIZE} byte cap")
data = _recv_exact(sock, length)
if data is None:
raise ConnectionError("Connection closed while reading message body")
return safe_pickle_loads(data)
def _recv_exact(sock, n: int) -> Optional[bytes]:
"""Receive exactly n bytes from a socket."""
buf = bytearray()
while len(buf) < n:
chunk = sock.recv(n - len(buf))
if not chunk:
return None
buf.extend(chunk)
return bytes(buf)
def compute_env_stamp() -> Dict[str, str]:
"""Local environment fingerprint for the IPC weight cache.
Returns the device compute capability and torch version of the current
process. A daemon and a connecting client that differ on either may have run
different post-processing / kernel-selection branches, producing weights that
map cleanly over IPC yet serve garbage; stamping these into CacheConfig turns
that into a clean mismatch. Imported lazily so protocol.py stays cheap to
import and usable on CPU-only hosts (both fields degrade to "").
"""
device_capability = ""
torch_version = ""
try:
import torch
torch_version = str(torch.__version__)
except Exception:
pass
try:
from sglang.srt.platforms import current_platform
cap = current_platform.get_device_capability()
if cap is not None:
device_capability = f"{cap.major}.{cap.minor}"
except Exception:
pass
return {"device_capability": device_capability, "torch_version": torch_version}
def compute_global_rank(tp_size: int, pp_rank: int, tp_rank: int) -> int:
"""Single source of truth for the daemon rank formula.
global_rank = tp_size * pp_rank + tp_rank, so each daemon gets a unique
socket/ready path even across PP stages and nodes. Every call site (engine,
loader, model_runner, daemon) must go through this so the copies can't drift.
"""
return tp_size * pp_rank + tp_rank
def compute_local_gpu_id(
pp_rank: int,
tp_rank: int,
pp_size_per_node: int,
tp_size_per_node: int,
base_gpu_id: int = 0,
gpu_id_step: int = 1,
) -> int:
"""Single source of truth for the local GPU id a daemon rank runs on.
Mirrors the engine's device assignment so a daemon and the engine rank it
serves always land on the same physical GPU (a prerequisite for CUDA IPC).
``base_gpu_id``/``gpu_id_step`` default to the identity mapping used by the
standalone launcher; the engine passes its real ``--base-gpu-id`` /
``--gpu-id-step`` so every call site computes the id the same way instead of
keeping three drifting copies of the formula.
"""
return (
base_gpu_id
+ (pp_rank % pp_size_per_node) * tp_size_per_node
+ (tp_rank % tp_size_per_node) * gpu_id_step
)
def get_socket_path(global_rank: int) -> str:
"""Get the Unix socket path for a weight cache daemon.
global_rank = tp_size * pp_rank + tp_rank
"""
return WEIGHT_CACHE_SOCKET_TEMPLATE.format(global_rank=global_rank)
def get_ready_path(global_rank: int) -> str:
"""Get the ready-file path for a weight cache daemon.
global_rank = tp_size * pp_rank + tp_rank
"""
return WEIGHT_CACHE_READY_TEMPLATE.format(global_rank=global_rank)
def _read_ready_pid(ready_path: str) -> Optional[int]:
"""Read the daemon PID from a .ready file. Returns None if unreadable."""
try:
with open(ready_path) as f:
for line in f:
if line.startswith("pid="):
return int(line.strip().split("=", 1)[1])
except (OSError, ValueError):
pass
return None
def _is_pid_alive(pid: int) -> bool:
"""Check whether a process is still running."""
try:
os.kill(pid, 0)
return True
except ProcessLookupError:
return False
except PermissionError:
return True
def cleanup_stale_daemon_files(global_rank: int, *, force: bool = False) -> None:
"""Validate and clean up .ready/.sock files for a daemon rank.
If the .ready file exists and the recorded PID is still alive, the daemon
is still running — raise RuntimeError so the caller doesn't clobber it,
unless ``force`` is set, in which case the running daemon is killed and its
files are taken over (stale-takeover path for a wedged/orphaned daemon).
If the PID is dead (or unreadable), the files are stale leftovers from a
crashed/killed daemon and are safe to remove.
"""
ready_path = get_ready_path(global_rank)
socket_path = get_socket_path(global_rank)
if not os.path.exists(ready_path) and not os.path.exists(socket_path):
return
pid = _read_ready_pid(ready_path) if os.path.exists(ready_path) else None
if pid is not None and _is_pid_alive(pid):
if not force:
raise RuntimeError(
f"Weight cache daemon for rank {global_rank} is already running "
f"(pid={pid}, ready={ready_path}). Stop the existing daemon before "
f"launching a new one, or pass force=True (--force) to kill it and "
f"take over."
)
logger.warning(
f"[weight_cache] force takeover: killing existing daemon pid={pid} "
f"for rank {global_rank} and reclaiming its socket/ready files."
)
try:
os.kill(pid, signal.SIGKILL)
except ProcessLookupError:
pass
for path in (ready_path, socket_path):
if os.path.exists(path):
os.unlink(path)
logger.info(f"Removed stale daemon file: {path}")