feat(eplb): add eplb warmup-only mode

Freeze expert layout after the initial EPLB warmup rebalances so heavy workloads avoid recurring rebalance memory spikes and OOMs.

python/sglang/srt/eplb/eplb_manager.py

@@ -18,6 +18,7 @@ class EPLBManager:
         super().__init__()
         self._model_runner = model_runner
         self._server_args = model_runner.server_args
+        self._warmup_only = self._server_args.eplb_warmup_only
         self._rebalance_layers_per_chunk = (
             self._server_args.eplb_rebalance_layers_per_chunk
         )
@@ -27,7 +28,9 @@ class EPLBManager:
         assert (
             self._server_args.eplb_rebalance_num_iterations
             >= self._server_args.expert_distribution_recorder_buffer_size
-        ), "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size"
+        ), (
+            "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size"
+        )
 
         if not get_global_expert_distribution_recorder().recording:
             get_global_expert_distribution_recorder().start_record()
@@ -35,6 +38,10 @@ class EPLBManager:
         logger.info(
             f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations."
         )
+        if self._warmup_only:
+            logger.info(
+                "[EPLBManager] warmup-only mode enabled; expert layout will become static after warmup rebalances."
+            )
 
         self._main_generator = self._entrypoint()
 
@@ -55,12 +62,26 @@ class EPLBManager:
                     yield
                 yield from self.rebalance()
 
+        if self._warmup_only:
+            self._freeze_after_warmup()
+            while True:
+                yield
+
         while True:
             for _ in range(self._rebalance_num_iterations):
                 yield
 
             yield from self.rebalance()
 
+    def _freeze_after_warmup(self):
+        recorder = get_global_expert_distribution_recorder()
+        if recorder.recording:
+            recorder.stop_record()
+        torch.get_device_module().empty_cache()
+        logger.info(
+            "[EPLBManager] warmup-only mode completed; expert layout is now static and EPLB rebalancing is disabled."
+        )
+
     def rebalance(self):
         logger.info("[EPLBManager] rebalance start")
 

python/sglang/srt/server_args.py

@@ -525,6 +525,7 @@ class ServerArgs:
     ep_dispatch_algorithm: Optional[Literal["static", "dynamic", "fake"]] = None
     init_expert_location: str = "trivial"
     enable_eplb: bool = False
+    eplb_warmup_only: bool = False
     eplb_algorithm: str = "auto"
     eplb_rebalance_num_iterations: int = 1000
     eplb_rebalance_layers_per_chunk: Optional[int] = None
@@ -2754,6 +2755,9 @@ class ServerArgs:
                 ), "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK (default 4096) must be larger or equal to chunked_prefill_size"
 
     def _handle_eplb_and_dispatch(self):
+        if self.eplb_warmup_only:
+            assert self.enable_eplb, "--eplb-warmup-only requires --enable-eplb."
+
         if self.enable_eplb and (self.expert_distribution_recorder_mode is None):
             self.expert_distribution_recorder_mode = "stat"
             logger.warning(
@@ -4862,6 +4866,11 @@ class ServerArgs:
             action="store_true",
             help="Enable EPLB algorithm",
         )
+        parser.add_argument(
+            "--eplb-warmup-only",
+            action="store_true",
+            help="Run EPLB only during its warmup rebalances, then keep expert layout static.",
+        )
         parser.add_argument(
             "--eplb-algorithm",
             type=str,