minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Weight Cache] Support static DP/EP layouts (#33684)

Browse Source 
Co-authored-by: siyu <liusy58@linux.alibaba.com>
This commit is contained in:
Xun Sun
2026-08-23 19:57:52 -07:00
committed by GitHub
co-authored by siyu
parent 514b997e6c
commit a90d770c40
9 changed files with 354 additions and 343 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/engine.py
+14 -56
View File
@@ -28,8 +28,6 @@ import multiprocessing as mp
import os
import random
import signal
import subprocess
import sys
import tempfile
import threading
import time
@@ -672,12 +670,6 @@ class Engine(EngineScoreMixin, EngineBase):
(``python -m sglang.srt.weight_cache.daemon``) plus
``--weight-cache-mode client``, where the daemon outlives the engine.
"""
if get_parallel().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.
@@ -721,6 +713,7 @@ class Engine(EngineScoreMixin, EngineBase):
# 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.daemon import spawn_weight_cache_daemon
from sglang.srt.weight_cache.protocol import (
cleanup_stale_daemon_files,
compute_global_rank,
@@ -743,49 +736,14 @@ class Engine(EngineScoreMixin, EngineBase):
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",
get_model().model_path,
"--gpu-id",
str(gpu_id),
"--tp-size",
str(tp_size),
"--tp-rank",
str(tp_rank),
"--pp-size",
str(configured_pp_size()),
"--pp-rank",
str(pp_rank),
"--dp-size",
"1",
"--ep-size",
str(get_parallel().ep_size),
"--load-format",
get_model().load_format,
"--dtype",
get_model().dtype,
"--dist-init-method",
dist_init_method,
]
if get_model().quantization:
cmd += ["--quantization", get_model().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 = spawn_weight_cache_daemon(
server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=pp_rank,
dist_init_method=dist_init_method,
)
proc = subprocess.Popen(cmd)
daemon_procs.append(proc)
# Wait for all daemons to be ready (ready file exists). On any failure
@@ -810,10 +768,10 @@ class Engine(EngineScoreMixin, EngineBase):
)
# Check if daemon process is still alive
for p in daemon_procs:
if p.poll() is not None:
if not p.is_alive():
raise RuntimeError(
f"Weight cache daemon (pid={p.pid}) exited prematurely "
f"with code {p.returncode}"
f"with code {p.exitcode}"
)
logger.info(
f"Weight cache daemon for pp_rank={pp_rank} "
@@ -844,17 +802,17 @@ class Engine(EngineScoreMixin, EngineBase):
if not procs:
return
for p in procs:
if p.poll() is None:
if p.is_alive():
p.terminate() # SIGTERM -> daemon cleanup handler runs
for p in procs:
try:
p.wait(timeout=timeout)
except subprocess.TimeoutExpired:
p.join(timeout=timeout)
if p.is_alive():
logger.warning(
f"Weight cache daemon (pid={p.pid}) did not exit within "
f"{timeout}s of SIGTERM; sending SIGKILL."
)
p.kill()
p.join()
@classmethod
def _launch_scheduler_processes(
python/sglang/srt/layers/moe/token_dispatcher/mooncake.py
+8 -2
View File
@@ -231,7 +231,7 @@ class _MooncakeEPDispatcherImpl:
use_fp8: bool = False,
):
buffer = self._get_buffer()
active_ranks = ElasticEPStateManager.instance().active_ranks
active_ranks = self._get_active_ranks()
packed_recv_hidden, packed_recv_count, self.handle, event, hook = (
buffer.dispatch(
hidden_states,
@@ -271,7 +271,7 @@ class _MooncakeEPDispatcherImpl:
topk_weights: torch.Tensor,
):
buffer = self._get_buffer()
active_ranks = ElasticEPStateManager.instance().active_ranks
active_ranks = self._get_active_ranks()
combined_hidden_states, event, hook = buffer.combine(
hidden_states,
topk_ids,
@@ -286,6 +286,12 @@ class _MooncakeEPDispatcherImpl:
self.handle = None
return combined_hidden_states, event, hook
def _get_active_ranks(self) -> torch.Tensor:
elastic_state = ElasticEPStateManager.instance()
if elastic_state is not None:
return elastic_state.active_ranks
return torch.ones(self.group.size(), dtype=torch.int32, device="cuda")
def _get_buffer(self):
return EPBuffer.get_ep_buffer(
self.group,
python/sglang/srt/server_args.py
+5
View File
@@ -8179,6 +8179,11 @@ class ServerArgs:
"(--weight-cache-mode off) for this configuration."
)
if self.weight_cache_mode != "off" and self.enable_eplb:
raise ValueError(
"--weight-cache-mode is not supported together with --enable-eplb."
)
def _is_mistral_native_format(self) -> bool:
"""True iff the checkpoint requires load_format=mistral.
python/sglang/srt/weight_cache/daemon.py
+200 -278
View File
@@ -34,19 +34,22 @@ Usage:
--dist-init-method tcp://127.0.0.1:29500
"""
import argparse
import dataclasses
import logging
import multiprocessing
import os
import signal
import socket
import time
from typing import Any, Dict, Optional, Tuple
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
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.runtime_context import publish
from sglang.srt.runtime_context import get_parallel, publish
from .protocol import (
CacheConfig,
@@ -66,6 +69,9 @@ from .transport import choose_daemon_transport_backend
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
# 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
@@ -73,6 +79,59 @@ logger = logging.getLogger(__name__)
CLIENT_CONNECTION_TIMEOUT = 30.0
@dataclasses.dataclass
class WeightCacheDaemonArgs:
"""Daemon-private worker identity and standalone launcher controls.
``gpu_id``, ``tp_rank``, ``pp_rank``, and ``dist_init_method`` identify a
worker after the shared server configuration has been resolved. ``timeout``
and ``force`` apply only when this entrypoint launches and monitors a local
daemon group. None has a corresponding ServerArgs field with these
per-worker semantics.
"""
gpu_id: Optional[int] = None
tp_rank: Optional[int] = None
pp_rank: int = 0
dist_init_method: Optional[str] = None
timeout: int = 1800
force: bool = False
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--gpu-id",
type=int,
default=None,
help="GPU device ID for a single daemon. If omitted, launches all local ranks.",
)
parser.add_argument(
"--tp-rank",
type=int,
default=None,
help="TP rank for a single daemon. If omitted, launches all local ranks.",
)
parser.add_argument("--pp-rank", type=int, default=0)
parser.add_argument(
"--dist-init-method",
default=None,
help="Daemon distributed init method (for example tcp://host:29500).",
)
parser.add_argument("--timeout", type=int, default=1800)
parser.add_argument("--force", action="store_true")
@classmethod
def from_cli_args(cls, args: argparse.Namespace) -> "WeightCacheDaemonArgs":
return cls(
gpu_id=args.gpu_id,
tp_rank=args.tp_rank,
pp_rank=args.pp_rank,
dist_init_method=args.dist_init_method,
timeout=args.timeout,
force=args.force,
)
class WeightCacheDaemon:
"""Persistent GPU weight cache for a single TP rank.
@@ -82,42 +141,42 @@ class WeightCacheDaemon:
def __init__(
self,
model_path: str,
server_args: "ServerArgs",
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,
tp_rank: int,
pp_rank: int,
dist_init_method: Optional[str] = None,
):
self.model_path = model_path
self.server_args = server_args
self.model_path = server_args.model_path
self.gpu_id = gpu_id
self.tp_size = tp_size
self.tp_size = server_args.tp_size
self.tp_rank = tp_rank
self.pp_size = pp_size
self.pp_size = server_args.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.dp_size = server_args.dp_size
self.ep_size = server_args.ep_size
self.moe_dp_size = server_args.moe_dp_size
self.enable_dp_attention = server_args.enable_dp_attention
self.enable_dp_lm_head = server_args.enable_dp_lm_head
self.attn_cp_size = server_args.attn_cp_size
self.moe_dense_tp_size = server_args.moe_dense_tp_size
self.moe_a2a_backend = server_args.moe_a2a_backend
self.deepep_mode = server_args.deepep_mode
self.load_format = server_args.load_format
self.dtype = server_args.dtype
self.quantization = server_args.quantization
self.model_loader_extra_config = server_args.model_loader_extra_config
self.trust_remote_code = server_args.trust_remote_code
self.revision = server_args.revision
self.dist_init_method = dist_init_method
self.socket_path = get_socket_path(
compute_global_rank(tp_size, pp_rank, tp_rank)
compute_global_rank(self.tp_size, pp_rank, tp_rank)
)
self.ready_path = get_ready_path(
compute_global_rank(self.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
@@ -164,12 +223,18 @@ class WeightCacheDaemon:
distributed_init_method=self.dist_init_method,
local_rank=self.gpu_id,
backend=current_platform.get_torch_distributed_backend_str(),
moe_a2a_backend=server_args.moe_a2a_backend,
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
pipeline_model_parallel_size=self.pp_size,
expert_model_parallel_size=self.ep_size,
attention_data_parallel_size=(
self.dp_size if self.enable_dp_attention else 1
),
attention_context_model_parallel_size=self.attn_cp_size,
moe_data_model_parallel_size=self.moe_dp_size,
)
# Initialize DP attention state (required by some models like Qwen3 MoE)
@@ -210,22 +275,14 @@ class WeightCacheDaemon:
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.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,
)
server_args = self.server_args
publish(server_args, role="weight_cache_daemon")
from sglang.srt.layers.moe import initialize_moe_config
initialize_moe_config(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
@@ -237,7 +294,7 @@ class WeightCacheDaemon:
quantization=self.quantization,
)
# Build cache config fingerprint BEFORE loading the model.
# 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.
@@ -250,6 +307,15 @@ class WeightCacheDaemon:
if not quant_method and quant_config is not None:
quant_method = get_quant_method_name(quant_config)
# Refuse unsupported quant methods before creating distributed groups
# or touching model weights.
check_ipc_quant_support(quant_method, quant_config, where="daemon")
# The initialized groups are the authority for rank identity. This
# avoids maintaining a second copy of the model-parallel hierarchy.
self._init_distributed(server_args, model_config)
moe_dp_rank = get_parallel().moe_dp_rank
moe_ep_rank = get_parallel().moe_ep_rank
self.config = CacheConfig(
model_path=self.model_path,
model_arch=(
@@ -263,6 +329,14 @@ class WeightCacheDaemon:
pp_rank=self.pp_rank,
dp_size=self.dp_size,
ep_size=self.ep_size,
moe_dp_size=self.moe_dp_size,
moe_dp_rank=moe_dp_rank,
moe_ep_rank=moe_ep_rank,
enable_dp_attention=self.enable_dp_attention,
enable_dp_lm_head=self.enable_dp_lm_head,
attn_cp_size=self.attn_cp_size,
moe_dense_tp_size=self.moe_dense_tp_size,
moe_a2a_backend=self.moe_a2a_backend,
quant_method=quant_method,
quant_config_hash=hash_quant_config(quant_config),
dtype=str(model_config.dtype),
@@ -270,14 +344,6 @@ class WeightCacheDaemon:
**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,
@@ -523,20 +589,10 @@ class WeightCacheDaemon:
def run_weight_cache_daemon(
model_path: str,
server_args: "ServerArgs",
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,
tp_rank: int,
pp_rank: int,
dist_init_method: Optional[str] = None,
):
"""Entry point for running a weight cache daemon process."""
@@ -554,20 +610,10 @@ def run_weight_cache_daemon(
kill_itself_when_parent_died()
daemon = WeightCacheDaemon(
model_path=model_path,
server_args=server_args,
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,
)
@@ -575,22 +621,26 @@ def run_weight_cache_daemon(
daemon.serve()
def spawn_weight_cache_daemon(
server_args: "ServerArgs",
*,
gpu_id: int,
tp_rank: int,
pp_rank: int,
dist_init_method: str,
):
"""Start one daemon from the complete resolved server configuration."""
ctx = multiprocessing.get_context("spawn")
proc = ctx.Process(
target=run_weight_cache_daemon,
args=(server_args, gpu_id, tp_rank, pp_rank, dist_init_method),
)
proc.start()
return proc
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,
server_args: "ServerArgs",
dist_init_method: Optional[str] = None,
timeout: int = 1800,
force: bool = False,
@@ -601,9 +651,9 @@ def launch_weight_cache_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.
Uses ``multiprocessing`` with the ``spawn`` start method. The child starts
in a clean interpreter and receives the complete resolved ``ServerArgs``
through pickle, avoiding a second hand-maintained CLI configuration path.
Usage (single-node):
python -m sglang.srt.weight_cache.daemon \\
@@ -623,24 +673,22 @@ def launch_weight_cache_daemons(
--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_size_per_node = max(server_args.pp_size // server_args.nnodes, 1)
nnodes_per_pp_rank = max(server_args.nnodes // server_args.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),
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
pp_size_per_node * (server_args.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_size_per_node = server_args.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),
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group),
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
)
if nnodes > 1 and dist_init_method is None:
if server_args.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 "
@@ -654,13 +702,10 @@ def launch_weight_cache_daemons(
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)
global_rank = compute_global_rank(server_args.tp_size, pp_rank, tp_rank)
cleanup_stale_daemon_files(global_rank, force=force)
procs = []
@@ -671,46 +716,16 @@ def launch_weight_cache_daemons(
tp_rank,
pp_size_per_node,
tp_size_per_node,
base_gpu_id=base_gpu_id,
gpu_id_step=gpu_id_step,
base_gpu_id=server_args.base_gpu_id,
gpu_id_step=server_args.gpu_id_step,
)
proc = spawn_weight_cache_daemon(
server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=pp_rank,
dist_init_method=dist_init_method,
)
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} "
@@ -723,7 +738,7 @@ def launch_weight_cache_daemons(
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)
global_rank = compute_global_rank(server_args.tp_size, pp_rank, tp_rank)
ready_path = get_ready_path(global_rank)
while not os.path.exists(ready_path):
time.sleep(check_interval)
@@ -740,39 +755,38 @@ def launch_weight_cache_daemons(
)
# Check if any daemon exited prematurely
for p in procs:
retcode = p.poll()
if retcode is not None:
if not p.is_alive():
logger.error(
f"Weight cache daemon exited prematurely "
f"with code {retcode}"
f"with code {p.exitcode}"
)
for other in procs:
if other.poll() is None:
if other.is_alive():
other.terminate()
raise RuntimeError(
f"Weight cache daemon exited prematurely "
f"with code {retcode}"
f"with code {p.exitcode}"
)
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"All {num_daemons} weight cache daemons on node {server_args.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
# Monitor daemons and, the moment any one exits, terminate the rest and
# raise. A serial proc.join() 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:
if not proc.is_alive():
exited = proc
break
else:
@@ -782,149 +796,57 @@ def launch_weight_cache_daemons(
logger.info("Received KeyboardInterrupt, shutting down daemons")
finally:
for proc in procs:
if proc.poll() is None:
if proc.is_alive():
proc.terminate()
for proc in procs:
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
proc.join(timeout=5)
if proc.is_alive():
proc.kill()
proc.join()
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."
f"{exited.exitcode}; terminated the remaining daemons."
)
if __name__ == "__main__":
import argparse
worker_parser = argparse.ArgumentParser(add_help=False)
WeightCacheDaemonArgs.add_cli_args(worker_parser)
worker_ns, server_argv = worker_parser.parse_known_args()
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).",
)
from sglang.srt.server_args import prepare_server_args
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.
server_args = prepare_server_args(server_argv)
daemon_args = WeightCacheDaemonArgs.from_cli_args(worker_ns)
if daemon_args.gpu_id is not None or daemon_args.tp_rank is not None:
gpu_id = (
daemon_args.gpu_id
if daemon_args.gpu_id is not None
else daemon_args.tp_rank
)
tp_rank = (
daemon_args.tp_rank
if daemon_args.tp_rank is not None
else daemon_args.gpu_id
)
cleanup_stale_daemon_files(
compute_global_rank(args.tp_size, args.pp_rank, tp_rank),
force=args.force,
compute_global_rank(server_args.tp_size, daemon_args.pp_rank, tp_rank),
force=daemon_args.force,
)
run_weight_cache_daemon(
model_path=args.model_path,
server_args,
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,
pp_rank=daemon_args.pp_rank,
dist_init_method=daemon_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,
server_args,
dist_init_method=daemon_args.dist_init_method,
timeout=daemon_args.timeout,
force=daemon_args.force,
)
python/sglang/srt/weight_cache/ipc_loader.py
+13 -2
View File
@@ -478,7 +478,7 @@ class IpcModelLoader(BaseModelLoader):
try:
# Build engine's config fingerprint
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import get_exec, get_parallel
ps = get_parallel()
tp_size = ps.tp_size
@@ -488,8 +488,11 @@ class IpcModelLoader(BaseModelLoader):
pp_rank = ps.pp_rank
ep_size = ps.moe_ep_size
moe_dp_size = ps.moe_dp_size
moe_dp_rank = ps.moe_dp_rank
moe_ep_rank = ps.moe_ep_rank
dp_size = get_parallel().dp_size
dp_size = ps.dp_size
quant_method, quant_config = self._resolve_engine_quant(model_config)
@@ -506,6 +509,14 @@ class IpcModelLoader(BaseModelLoader):
pp_rank=pp_rank,
dp_size=dp_size,
ep_size=ep_size,
moe_dp_size=moe_dp_size,
moe_dp_rank=moe_dp_rank,
moe_ep_rank=moe_ep_rank,
enable_dp_attention=ps.enable_dp_attention,
enable_dp_lm_head=ps.enable_dp_lm_head,
attn_cp_size=ps.attn_cp_size,
moe_dense_tp_size=ps.moe_dense_tp_size,
moe_a2a_backend=get_exec().moe.moe_a2a_backend,
quant_method=quant_method,
quant_config_hash=hash_quant_config(quant_config),
dtype=str(model_config.dtype),
python/sglang/srt/weight_cache/protocol.py
+8
View File
@@ -42,6 +42,14 @@ class CacheConfig(msgspec.Struct):
pp_rank: int
dp_size: int
ep_size: int
moe_dp_size: int
moe_dp_rank: int
moe_ep_rank: int
enable_dp_attention: bool
enable_dp_lm_head: bool
attn_cp_size: int
moe_dense_tp_size: Optional[int]
moe_a2a_backend: str
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"
test/manual/test_weight_cache_e2e.py
+8 -4
View File
@@ -54,20 +54,24 @@ def run_single_daemon(
"""Run a single daemon process for one (pp_rank, tp_rank)."""
import traceback
from sglang.srt.server_args import ServerArgs
from sglang.srt.weight_cache.daemon import WeightCacheDaemon
daemon = WeightCacheDaemon(
server_args = ServerArgs(
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=1,
load_format=load_format,
dtype=dtype,
quantization=quantization,
trust_remote_code=trust_remote_code,
)
daemon = WeightCacheDaemon(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=pp_rank,
dist_init_method=dist_init_method,
)
daemon.socket_path = socket_path
test/registered/unit/model_loader/test_weight_cache_protocol.py
+77 -1
View File
@@ -7,7 +7,7 @@ These cover the pure-Python logic that the GPU end-to-end test
- length-prefixed socket framing (send_msg/recv_msg) over socketpair()
- CacheConfig fingerprint matching / (de)serialization
- quant-config hashing and method-name extraction
- the daemon rank formula and socket/ready path derivation
- daemon spawn configuration and socket/ready path derivation
- the IPC quantization allowlist (the gate that keeps silently-wrong
quant methods off the zero-copy path)
- stale-vs-live daemon file cleanup
@@ -21,6 +21,7 @@ import os
import socket
import struct
import unittest
from types import SimpleNamespace
import torch
@@ -61,6 +62,14 @@ def _make_cache_config(**overrides) -> CacheConfig:
pp_rank=0,
dp_size=1,
ep_size=1,
moe_dp_size=1,
moe_dp_rank=0,
moe_ep_rank=0,
enable_dp_attention=False,
enable_dp_lm_head=False,
attn_cp_size=1,
moe_dense_tp_size=None,
moe_a2a_backend="none",
quant_method="",
quant_config_hash="",
dtype="torch.float16",
@@ -159,6 +168,11 @@ class TestCacheConfig(CustomTestCase):
base = _make_cache_config()
for field, value in (
("tp_rank", 1),
("moe_dp_rank", 1),
("moe_ep_rank", 1),
("enable_dp_attention", True),
("moe_dense_tp_size", 1),
("moe_a2a_backend", "mooncake"),
("dtype", "torch.bfloat16"),
("quant_method", "fp8"),
("model_path", "/models/other"),
@@ -250,6 +264,68 @@ class TestGlobalRankAndPaths(CustomTestCase):
)
class TestDaemonLaunchConfiguration(CustomTestCase):
def test_spawn_forwards_complete_server_args_without_projection(self):
from sglang.srt.weight_cache import daemon
# The spawn helper receives Engine's already-resolved ServerArgs. A
# minimal namespace keeps this projection test CPU-only and
# model-independent; importantly, no EPLB configuration is involved.
server_args = SimpleNamespace(
model_path="/models/demo",
tp_size=8,
pp_size=1,
dp_size=8,
ep_size=8,
moe_dp_size=2,
enable_dp_attention=True,
enable_dp_lm_head=True,
attn_cp_size=2,
moe_dense_tp_size=1,
moe_a2a_backend="mooncake",
deepep_mode="low_latency",
load_format="safetensors",
dtype="bfloat16",
quantization="fp8",
model_loader_extra_config='{"key": "value"}',
trust_remote_code=True,
revision="test-revision",
)
class FakeProcess:
pid = 1234
def start(self):
pass
class FakeContext:
def Process(self, **kwargs):
self.kwargs = kwargs
return FakeProcess()
fake_context = FakeContext()
from unittest import mock
with mock.patch.object(
daemon.multiprocessing, "get_context", return_value=fake_context
) as get_context:
result = daemon.spawn_weight_cache_daemon(
server_args,
gpu_id=3,
tp_rank=3,
pp_rank=0,
dist_init_method="tcp://127.0.0.1:29500",
)
get_context.assert_called_once_with("spawn")
self.assertIsInstance(result, FakeProcess)
self.assertIs(fake_context.kwargs["target"], daemon.run_weight_cache_daemon)
self.assertEqual(
fake_context.kwargs["args"],
(server_args, 3, 3, 0, "tcp://127.0.0.1:29500"),
)
class TestIpcQuantAllowlist(CustomTestCase):
def test_unquantized_is_supported(self):
self.assertTrue(is_ipc_quant_supported("", None))
test/registered/unit/test_supplied_instance_exposure_ratchet.py
+21
View File
@@ -179,6 +179,22 @@ _EXPOSED = {
("utils/common.py", "speculative_num_draft_tokens"),
("utils/common.py", "speculative_num_steps"),
("utils/hf_transformers/processor.py", "image_processor_backend"),
# The daemon command and constructor snapshot the resolved startup layout
# before the daemon's loading lifecycle can apply any runtime overrides.
("weight_cache/daemon.py", "attn_cp_size"),
("weight_cache/daemon.py", "deepep_mode"),
("weight_cache/daemon.py", "dp_size"),
("weight_cache/daemon.py", "dtype"),
("weight_cache/daemon.py", "enable_dp_attention"),
("weight_cache/daemon.py", "enable_dp_lm_head"),
("weight_cache/daemon.py", "ep_size"),
("weight_cache/daemon.py", "load_format"),
("weight_cache/daemon.py", "model_path"),
("weight_cache/daemon.py", "moe_a2a_backend"),
("weight_cache/daemon.py", "moe_dense_tp_size"),
("weight_cache/daemon.py", "moe_dp_size"),
("weight_cache/daemon.py", "pp_size"),
("weight_cache/daemon.py", "quantization"),
}
# Pairs whose resolution write only happens on a CUDA host (capability or
@@ -196,6 +212,11 @@ _OVERRIDDEN_AND_READ = {
("dllm/config.py", "model_path"),
("entrypoints/engine.py", "reasoning_parser"),
("entrypoints/engine.py", "tool_call_parser"),
("weight_cache/daemon.py", "dp_size"),
("weight_cache/daemon.py", "dtype"),
("weight_cache/daemon.py", "ep_size"),
("weight_cache/daemon.py", "load_format"),
("weight_cache/daemon.py", "model_path"),
("configs/model_config.py", "dtype"),
("configs/model_config.py", "model_path"),
("mem_cache/kv_cache_builder.py", "hicache_storage_backend"),