diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 5c7279b80..46cfa4dfe 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -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( diff --git a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py index 5fbf494e3..0a8d9ec55 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7b584d561..bd850b330 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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. diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 5c6c0e1f5..d14d79075 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -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://: 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, ) diff --git a/python/sglang/srt/weight_cache/ipc_loader.py b/python/sglang/srt/weight_cache/ipc_loader.py index a649b7fdd..5e6c942bf 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -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), diff --git a/python/sglang/srt/weight_cache/protocol.py b/python/sglang/srt/weight_cache/protocol.py index 8b3b4b341..befdf8e1e 100644 --- a/python/sglang/srt/weight_cache/protocol.py +++ b/python/sglang/srt/weight_cache/protocol.py @@ -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" diff --git a/test/manual/test_weight_cache_e2e.py b/test/manual/test_weight_cache_e2e.py index 562820497..5cd757f11 100644 --- a/test/manual/test_weight_cache_e2e.py +++ b/test/manual/test_weight_cache_e2e.py @@ -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 diff --git a/test/registered/unit/model_loader/test_weight_cache_protocol.py b/test/registered/unit/model_loader/test_weight_cache_protocol.py index ae4b4354c..d130f98f0 100644 --- a/test/registered/unit/model_loader/test_weight_cache_protocol.py +++ b/test/registered/unit/model_loader/test_weight_cache_protocol.py @@ -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)) diff --git a/test/registered/unit/test_supplied_instance_exposure_ratchet.py b/test/registered/unit/test_supplied_instance_exposure_ratchet.py index bc94d70ec..aa0dc8fa1 100644 --- a/test/registered/unit/test_supplied_instance_exposure_ratchet.py +++ b/test/registered/unit/test_supplied_instance_exposure_ratchet.py @@ -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"),