[Fix] weight_cache: Keep the EPLB expert layout in the daemon

python/sglang/srt/eplb/eplb_manager.py

@@ -19,6 +19,7 @@
 )
 from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
 from sglang.srt.runtime_context import get_exec, get_model, get_parallel
+from sglang.srt.weight_cache.expert_layout import get_expert_layout_session
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -67,6 +68,13 @@ def __init__(
             f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations."
         )
 
+        # With the weight cache on, the expert weights are normally the
+        # daemon's copy, shared with every engine that maps it. A rebalance
+        # moves experts in place, so it has to keep the daemon's record of the
+        # expert layout in step (see weight_cache/expert_layout.py).
+        self._weight_cache_mode = _weight_cache_mode()
+        self._weight_cache_off_reason = None
+
         self._main_generator = self._entrypoint()
 
     def on_forward_pass_end(self):
@@ -112,6 +120,32 @@ def rebalance(self):
         ):
             return
 
+        # The daemon's expert layout record to move experts through, or None
+        # when the weights are this engine's own.
+        expert_layout = None
+        if self._weight_cache_mode != "off":
+            session = get_expert_layout_session()
+            mode, reason = (
+                ("off", "this engine has no expert layout session")
+                if session is None
+                else session.rebalance_mode()
+            )
+            if mode == "coordinated" and get_exec().moe.elastic_ep_backend is not None:
+                mode, reason = (
+                    "off",
+                    "elastic EP also moves experts outside rebalance()",
+                )
+            if mode == "off":
+                if reason != self._weight_cache_off_reason:
+                    self._weight_cache_off_reason = reason
+                    logger.warning(
+                        "[EPLBManager] rebalancing is off with the weight cache: "
+                        f"{reason}. The expert layout stays as it is."
+                    )
+                return
+            if mode == "coordinated":
+                expert_layout = session
+
         logger.info("[EPLBManager] rebalance start")
 
         enable_timing = self._rebalance_layers_per_chunk is None
@@ -152,6 +186,14 @@ def rebalance(self):
         for chunk_layer_ids in update_layer_ids_chunks:
             if len(update_layer_ids_chunks) > 1:
                 yield
+            if expert_layout is not None:
+                reason = expert_layout.begin(chunk_layer_ids)
+                if reason is not None:
+                    logger.info(
+                        f"[EPLBManager] rebalance stopped before layers "
+                        f"{chunk_layer_ids}: {reason}"
+                    )
+                    return
             update_expert_location_with_recovery(
                 expert_location_updater=self._get_expert_location_updater(),
                 model=self._get_model(),
@@ -174,6 +216,22 @@ def rebalance(self):
                 # P2P waits only synchronize participating peers. Ranks without
                 # moves must also install this chunk before NIXL resumes.
                 dist.barrier()
+            if expert_layout is not None:
+                # Once committed, the daemon may serve these weights to the
+                # next engine, so the moves must have landed.
+                torch.get_device_module().synchronize()
+                reason = expert_layout.commit(
+                    chunk_layer_ids,
+                    get_global_expert_location_metadata().physical_to_logical_map_cpu,
+                )
+                if reason is not None:
+                    logger.error(
+                        f"[EPLBManager] the weight cache daemons did not record the "
+                        f"move of layers {chunk_layer_ids}: {reason}. Rebalancing "
+                        f"is off; restart the daemons before another engine uses "
+                        f"them."
+                    )
+                    return
 
         self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
 
@@ -351,3 +409,9 @@ def update_expert_location_with_recovery(
 def _chunk_list(items: List, chunk_size):
     for start_index in range(0, len(items), chunk_size):
         yield items[start_index : start_index + chunk_size]
+
+
+def _weight_cache_mode() -> str:
+    # The model config's field. Inside EPLBManager.__init__, ``get_model`` names
+    # the model getter argument instead, and the model is not loaded yet there.
+    return get_model().weight_cache_mode

python/sglang/srt/model_executor/model_runner.py

@@ -255,6 +255,7 @@
 from sglang.srt.utils.profile_utils import build_step_span_name
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.srt.utils.weight_checker import WeightChecker
+from sglang.srt.weight_cache.expert_layout import adopt_weight_cache_expert_layout
 
 _is_npu = is_npu()
 _is_cpu_amx_available = cpu_has_amx_support()
@@ -751,9 +752,16 @@ def maybe_init_expert_location_metadata(self):
             else 0
         )
         set_global_expert_location_metadata(
-            compute_initial_expert_location_metadata(
+            # With the weight cache, route with the layout the daemon's
+            # weights are in, which an EPLB rebalance may have changed.
+            adopt_weight_cache_expert_layout(
+                compute_initial_expert_location_metadata(
+                    model_config=self.model_config,
+                    moe_ep_rank=expert_rank,
+                ),
                 model_config=self.model_config,
                 moe_ep_rank=expert_rank,
+                gpu_id=self.gpu_id,
             )
         )
         if (

python/sglang/srt/runtime_context.py

@@ -621,6 +621,9 @@ class Resources(_FlagGroupBase):
     # (owning accessors live in sglang.srt.eplb).
     expert_distribution_recorder: Any = None
     expert_location_metadata: Any = None
+    # Weight cache: this engine's link to its daemon's record of the expert
+    # layout (owning accessors live in sglang.srt.weight_cache.expert_layout).
+    expert_layout_session: Any = None
     # LPLB: layer_id -> solver.
     lplb_solvers: dict = msgspec.field(default_factory=dict)
     # Named side streams (see RuntimeContext.get_stream): name -> stream.

python/sglang/srt/weight_cache/daemon.py

@@ -59,6 +59,15 @@
     spawn_world_rank,
 )
 
+from .expert_layout import (
+    ATTACH,
+    LAYOUT_REFUSED,
+    OK,
+    REQUEST_TYPES,
+    ExpertLayoutPresence,
+    ExpertLayoutRegistry,
+    handle_expert_layout_request,
+)
 from .protocol import (
     CacheConfig,
     check_ipc_quant_support,
@@ -188,6 +197,10 @@ def __init__(
         self.state_entries: Dict[str, Dict[str, Any]] = {}
         self.preloaded_weights_bytes = 0
         self.transport_backend = None
+        # Where the routed experts of the served weights are (see
+        # expert_layout.py); tracked once load() builds the EPLB layout.
+        self.expert_layout = ExpertLayoutRegistry(None)
+        self.expert_layout_presence: Optional[ExpertLayoutPresence] = None
 
     def _init_distributed(self, server_args, model_config):
         """Initialize the distributed backend required for model loading.
@@ -324,6 +337,7 @@ def load(self):
         # avoids maintaining a second copy of the model-parallel hierarchy.
         self._init_distributed(server_args, model_config)
         self._initialize_eplb_expert_location_metadata(model_config)
+        self._track_expert_layout()
         moe_dp_rank = get_parallel().moe_dp_rank
         moe_ep_rank = get_parallel().moe_ep_rank
         self.config = CacheConfig(
@@ -488,6 +502,18 @@ def _initialize_eplb_expert_location_metadata(self, model_config) -> None:
             )
         )
 
+    def _track_expert_layout(self) -> None:
+        """Keep the record of the EPLB layout built above. An EPLB rebalance in
+        an engine moves experts in these weights and reports it here."""
+        if not get_exec().moe.enable_eplb:
+            return
+
+        from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
+
+        self.expert_layout = ExpertLayoutRegistry.from_metadata(
+            get_global_expert_location_metadata()
+        )
+
     def serve(self):
         """Block and serve IPC handles over Unix socket."""
         # Do NOT unlink an existing socket here: stale-file cleanup is the launch
@@ -523,6 +549,9 @@ def _signal_handler(signum, frame):
         signal.signal(signal.SIGTERM, _signal_handler)
         signal.signal(signal.SIGINT, _signal_handler)
 
+        self.expert_layout_presence = ExpertLayoutPresence(self.expert_layout)
+        self.expert_layout_presence.start()
+
         try:
             while self._running:
                 try:
@@ -533,28 +562,36 @@ def _signal_handler(signum, frame):
                     # (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()
+                    self._serve_connection(conn)
                 except socket.timeout:
                     continue
         finally:
+            self.expert_layout_presence.stop()
             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."""
+    def _serve_connection(self, conn: socket.socket) -> None:
+        """Handle one accepted connection, then close it unless it became an
+        engine's expert layout attach connection."""
+        keep_open = False
+        try:
+            keep_open = self._handle_connection(conn)
+        except Exception as e:
+            logger.error(
+                f"[WeightCacheDaemon gpu={self.gpu_id}] Error handling connection: {e}",
+                exc_info=True,
+            )
+        finally:
+            if not keep_open:
+                conn.close()
+
+    def _handle_connection(self, conn: socket.socket) -> bool:
+        """Handle a single client connection. Returns True when the connection
+        must stay open: an engine's expert layout attach connection."""
         req = recv_msg(conn)
 
         if req.get("type") == "query_config":
@@ -580,7 +617,17 @@ def _handle_connection(self, conn: socket.socket):
                 send_msg(
                     conn, {"status": "mismatch", "daemon_config": self.config.to_dict()}
                 )
-                return
+                return False
+
+            # The engine must route with the expert layout these weights are in.
+            refusal = self.expert_layout.check_fetch(req.get("expert_layout"))
+            if refusal is not None:
+                logger.warning(
+                    f"[WeightCacheDaemon gpu={self.gpu_id}] "
+                    f"Refusing to serve the weights: {refusal}"
+                )
+                send_msg(conn, {"status": LAYOUT_REFUSED, "message": refusal})
+                return False
 
             logger.info(
                 f"[WeightCacheDaemon gpu={self.gpu_id}] "
@@ -601,6 +648,9 @@ def _handle_connection(self, conn: socket.socket):
         elif req.get("type") == "ping":
             send_msg(conn, {"status": "ok"})
 
+        elif req.get("type") in REQUEST_TYPES:
+            return self._handle_expert_layout_request(conn, req)
+
         else:
             send_msg(
                 conn,
@@ -609,6 +659,21 @@ def _handle_connection(self, conn: socket.socket):
                     "message": f"Unknown request type: {req.get('type')}",
                 },
             )
+        return False
+
+    def _handle_expert_layout_request(self, conn: socket.socket, req) -> bool:
+        reply = handle_expert_layout_request(self.expert_layout, req)
+        if req["type"] != ATTACH or reply.get("status") != OK:
+            send_msg(conn, reply)
+            return False
+        # Attached: the connection now tells whether the engine is alive.
+        try:
+            send_msg(conn, reply)
+        except BaseException:
+            self.expert_layout.detach(req["engine_id"])
+            raise
+        self.expert_layout_presence.watch(conn, req["engine_id"])
+        return True
 
     def shutdown(self):
         """Release GPU memory and clean up."""