weight cache: key daemon paths by GPU UUID (#36101)

Co-authored-by: siyu <liusy58@linux.alibaba.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/srt/entrypoints/engine.py

@@ -101,6 +101,7 @@ from sglang.srt.observability.startup_time import build_engine_startup_time
 from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
 from sglang.srt.parser.template_detection import resolve_auto_parsers
 from sglang.srt.parser.template_manager import TemplateManager
+from sglang.srt.platforms import current_platform
 from sglang.srt.plugins import load_plugins
 from sglang.srt.runtime_context import (
     get_disagg,
@@ -138,6 +139,12 @@ from sglang.srt.utils.network import (
 )
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.srt.utils.watchdog import SubprocessWatchdog
+from sglang.srt.weight_cache.daemon import spawn_weight_cache_daemon
+from sglang.srt.weight_cache.protocol import (
+    cleanup_stale_daemon_files,
+    compute_local_gpu_id,
+    get_ready_path,
+)
 from sglang.version import __version__
 
 logger = logging.getLogger(__name__)
@@ -715,19 +722,18 @@ 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,
-            compute_local_gpu_id,
-            get_ready_path,
-        )
-
+        # If a daemon is still alive at this GPU, raise instead of clobbering.
         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)
+                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,
+                )
+                cleanup_stale_daemon_files(current_platform.get_device_uuid(gpu_id))
 
         for pp_rank in pp_rank_range:
             for tp_rank in tp_rank_range:
@@ -759,8 +765,17 @@ class Engine(EngineScoreMixin, EngineBase):
         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)
+                    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,
+                    )
+                    ready_path = get_ready_path(
+                        current_platform.get_device_uuid(gpu_id)
+                    )
                     while not os.path.exists(ready_path):
                         time.sleep(check_interval)
                         if time.time() - start_time > timeout:

python/sglang/srt/environ.py

@@ -1612,14 +1612,14 @@ class Envs:
     # Weight Cache Daemon
     # ===================================================================
     # Paths the daemon and the engine ranks it serves must agree on. Both are
-    # format templates and must keep the {global_rank} placeholder: each rank
-    # talks to the daemon on its own GPU, so a rank-independent path would point
-    # every rank at one daemon and map another rank's shard.
+    # format templates and must keep the {device_uuid} placeholder: each daemon
+    # is keyed by the physical GPU it runs on, so a GPU-independent path would
+    # let one job's client discover another job's daemon.
     SGLANG_WEIGHT_CACHE_SOCKET_TEMPLATE = EnvStr(
-        "/tmp/sglang_weight_cache_rank{global_rank}.sock"
+        "/tmp/sglang_weight_cache_{device_uuid}.sock"
     )
     SGLANG_WEIGHT_CACHE_READY_TEMPLATE = EnvStr(
-        "/tmp/sglang_weight_cache_rank{global_rank}.ready"
+        "/tmp/sglang_weight_cache_{device_uuid}.ready"
     )
 
 

python/sglang/srt/model_executor/model_runner.py

@@ -1145,13 +1145,8 @@ class ModelRunner:
             weight_cache_socket=get_model().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,
-            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(

python/sglang/srt/model_executor/model_runner_components/load_model_utils.py

@@ -236,16 +236,13 @@ def build_load_config(
 def maybe_enable_ipc_weight_cache(
     *,
     load_config: LoadConfig,
-    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.
+    disk fallback). 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 get_model().weight_cache_mode == "off":
         return
@@ -254,17 +251,6 @@ def maybe_enable_ipc_weight_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(
     *,

python/sglang/srt/model_loader/loader.py

@@ -4379,21 +4379,10 @@ def get_model_loader(
 
     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:
-
-            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,
+            socket_path=load_config.weight_cache_socket,
             weight_cache_mode=load_config.weight_cache_mode,
             fallback_load_format=load_config.fallback_load_format,
         )

python/sglang/srt/server_args.py

@@ -3619,7 +3619,8 @@ class ServerArgs:
         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",
+            "If not set, derives the path from SGLANG_WEIGHT_CACHE_SOCKET_TEMPLATE "
+            "using the caller's physical GPU UUID.",
         ),
         NS("model"),
     ] = None

python/sglang/srt/weight_cache/daemon.py

@@ -173,12 +173,9 @@ class WeightCacheDaemon:
         self.revision = cfg.revision
         self.dist_init_method = dist_init_method
 
-        self.socket_path = get_socket_path(
-            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)
-        )
+        device_uuid = current_platform.get_device_uuid(gpu_id)
+        self.socket_path = get_socket_path(device_uuid)
+        self.ready_path = get_ready_path(device_uuid)
 
         self.model = None
         self.config: Optional[CacheConfig] = None
@@ -735,8 +732,17 @@ def launch_weight_cache_daemons(
     # 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(cfg.tp_size, pp_rank, tp_rank)
-            cleanup_stale_daemon_files(global_rank, force=force)
+            gpu_id = compute_local_gpu_id(
+                pp_rank,
+                tp_rank,
+                pp_size_per_node,
+                tp_size_per_node,
+                base_gpu_id=cfg.base_gpu_id,
+                gpu_id_step=cfg.gpu_id_step,
+            )
+            cleanup_stale_daemon_files(
+                current_platform.get_device_uuid(gpu_id), force=force
+            )
 
     procs = []
     for pp_rank in pp_rank_range:
@@ -768,8 +774,15 @@ 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(cfg.tp_size, pp_rank, tp_rank)
-            ready_path = get_ready_path(global_rank)
+            gpu_id = compute_local_gpu_id(
+                pp_rank,
+                tp_rank,
+                pp_size_per_node,
+                tp_size_per_node,
+                base_gpu_id=cfg.base_gpu_id,
+                gpu_id_step=cfg.gpu_id_step,
+            )
+            ready_path = get_ready_path(current_platform.get_device_uuid(gpu_id))
             while not os.path.exists(ready_path):
                 time.sleep(check_interval)
                 if time.time() - start_time > timeout:
@@ -863,7 +876,7 @@ if __name__ == "__main__":
             else daemon_args.gpu_id
         )
         cleanup_stale_daemon_files(
-            compute_global_rank(server_args.tp_size, daemon_args.pp_rank, tp_rank),
+            current_platform.get_device_uuid(gpu_id),
             force=daemon_args.force,
         )
         run_weight_cache_daemon(

python/sglang/srt/weight_cache/ipc_loader.py

@@ -22,6 +22,7 @@ from sglang.srt.model_loader.loader import (
     BaseModelLoader,
     _initialize_model,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_exec, get_parallel
 
 from .protocol import (
@@ -29,6 +30,7 @@ from .protocol import (
     check_ipc_quant_support,
     compute_env_stamp,
     get_quant_method_name,
+    get_socket_path,
     hash_quant_config,
     recv_msg,
     send_msg,
@@ -67,7 +69,7 @@ class IpcModelLoader(BaseModelLoader):
     def __init__(
         self,
         load_config: LoadConfig,
-        socket_path: str,
+        socket_path: Optional[str] = None,
         fallback_loader_cls=None,
         weight_cache_mode: str = "client",
         fallback_load_format: str = "auto",
@@ -103,7 +105,7 @@ class IpcModelLoader(BaseModelLoader):
         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)
+        cache_data = self._fetch_from_cache(model_config, device_config)
 
         if cache_data is None:
             if self.weight_cache_mode == "daemon":
@@ -446,7 +448,7 @@ class IpcModelLoader(BaseModelLoader):
 
         return model
 
-    def _fetch_from_cache(self, model_config) -> Optional[dict]:
+    def _fetch_from_cache(self, model_config, device_config) -> Optional[dict]:
         """Connect to daemon, validate config, fetch IPC handles.
 
         Returns the daemon response dict on success, None if the daemon is
@@ -455,6 +457,10 @@ class IpcModelLoader(BaseModelLoader):
         """
         import socket as socket_mod
 
+        if self.socket_path is None:
+            device_uuid = current_platform.get_device_uuid(int(device_config.gpu_id))
+            self.socket_path = get_socket_path(device_uuid)
+
         # 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).

python/sglang/srt/weight_cache/protocol.py

@@ -269,12 +269,7 @@ def compute_env_stamp() -> Dict[str, str]:
 
 
 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.
-    """
+    """Global rank for ``init_distributed_environment`` (tp_size * pp_rank + tp_rank)."""
     return tp_size * pp_rank + tp_rank
 
 
@@ -302,38 +297,32 @@ def compute_local_gpu_id(
     )
 
 
-def _format_daemon_path(env_field, global_rank: int) -> str:
-    """Fill in a daemon path template, rejecting one that drops the rank.
+def _format_daemon_path(env_field, device_uuid: str) -> str:
+    """Fill in a daemon path template, rejecting one that drops the GPU identity.
 
     The template is user-overridable, and ``str.format`` silently ignores a
-    missing placeholder. Every rank would then derive the same path and map the
-    shard belonging to whichever daemon got there first, so refuse up front
-    rather than serve wrong weights.
+    missing placeholder. Every physical GPU would then derive the same path,
+    letting one job's client discover another job's daemon, so refuse up
+    front rather than serve wrong weights.
     """
     template = env_field.get()
-    if "{global_rank}" not in template:
+    if "{device_uuid}" not in template:
         raise ValueError(
-            f"{env_field.name}={template!r} must contain '{{global_rank}}': each "
-            f"rank needs its own path, and a rank-independent one would point "
-            f"every rank at a single daemon."
+            f"{env_field.name}={template!r} must contain '{{device_uuid}}': "
+            f"each physical GPU needs its own path, and a GPU-independent one "
+            f"would point every caller at a single daemon."
         )
-    return template.format(global_rank=global_rank)
+    return template.format(device_uuid=device_uuid)
 
 
-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 _format_daemon_path(envs.SGLANG_WEIGHT_CACHE_SOCKET_TEMPLATE, global_rank)
+def get_socket_path(device_uuid: str) -> str:
+    """Get the Unix socket path for a weight cache daemon's physical GPU."""
+    return _format_daemon_path(envs.SGLANG_WEIGHT_CACHE_SOCKET_TEMPLATE, device_uuid)
 
 
-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 _format_daemon_path(envs.SGLANG_WEIGHT_CACHE_READY_TEMPLATE, global_rank)
+def get_ready_path(device_uuid: str) -> str:
+    """Get the ready-file path for a weight cache daemon's physical GPU."""
+    return _format_daemon_path(envs.SGLANG_WEIGHT_CACHE_READY_TEMPLATE, device_uuid)
 
 
 def _read_ready_pid(ready_path: str) -> Optional[int]:
@@ -359,8 +348,8 @@ def _is_pid_alive(pid: int) -> bool:
         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.
+def cleanup_stale_daemon_files(device_uuid: str, *, force: bool = False) -> None:
+    """Validate and clean up .ready/.sock files for a daemon's physical GPU.
 
     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,
@@ -369,8 +358,8 @@ def cleanup_stale_daemon_files(global_rank: int, *, force: bool = False) -> None
     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)
+    ready_path = get_ready_path(device_uuid)
+    socket_path = get_socket_path(device_uuid)
 
     if not os.path.exists(ready_path) and not os.path.exists(socket_path):
         return
@@ -380,14 +369,14 @@ def cleanup_stale_daemon_files(global_rank: int, *, force: bool = False) -> 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"Weight cache daemon for GPU {device_uuid} 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."
+            f"for GPU {device_uuid} and reclaiming its socket/ready files."
         )
         try:
             os.kill(pid, signal.SIGKILL)