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.

docs/advanced_features/expert_parallelism.md

@@ -141,7 +141,7 @@ SGLang introduces a dispatcher-hook system for Single-Batch Overlap (SBO), enabl
 
 SGLang integrates the [Expert Parallelism Load Balancer (EPLB)](https://github.com/deepseek-ai/EPLB) from DeepSeek to address routing imbalances in MoE models. By analyzing expert activation statistics, EPLB computes an optimal expert arrangement, strategically placing or replicating experts to minimize GPU utilization variance, reduce idle cycles, and enhance scalability.
 
-To enable EPLB, use the flags `--enable-eplb`. For optimal performance, increase batch sizes to stabilize activation statistics and configure periodic rebalancing (e.g., every 1000 requests) to adapt to evolving workloads. Simulations demonstrate significant improvements in load balancedness (ratio of mean to max computation time), correlating strongly with throughput gains.
+To enable EPLB, use `--enable-eplb`. For optimal performance, increase batch sizes to stabilize activation statistics and configure periodic rebalancing (e.g., every 1000 requests) to adapt to evolving workloads. If you only want EPLB to act during its initial warmup rebalances and then freeze the resulting expert placement, add `--eplb-warmup-only`. Simulations demonstrate significant improvements in load balancedness (ratio of mean to max computation time), correlating strongly with throughput gains.
 
 For more details, refer to the [EPLB Section in the Large-Scale EP Blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/#expert-parallelism-load-balancer) and the [EPLB Repository](https://github.com/deepseek-ai/eplb).
 

docs/advanced_features/server_arguments.md

@@ -321,6 +321,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
 | `--ep-dispatch-algorithm` | The algorithm to choose ranks for redundant experts in expert parallel. | `None` | Type: str |
 | `--init-expert-location` | Initial location of EP experts. | `trivial` | Type: str |
 | `--enable-eplb` | Enable EPLB algorithm | `False` | bool flag (set to enable) |
+| `--eplb-warmup-only` | Run EPLB only during its warmup rebalances, then keep expert layout static. | `False` | bool flag (set to enable) |
 | `--eplb-algorithm` | Chosen EPLB algorithm | `auto` | Type: str |
 | `--eplb-rebalance-num-iterations` | Number of iterations to automatically trigger a EPLB re-balance. | `1000` | Type: int |
 | `--eplb-rebalance-layers-per-chunk` | Number of layers to rebalance per forward pass. | `None` | Type: int |

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,