Extract MoE/EP setup into a moe_ep_setup module (#31159)

2026-07-14 16:00:24 +08:00
parent 440aebdfe0
commit 08798dba0d
3 changed files with 175 additions and 124 deletions
@@ -86,6 +86,10 @@ class EPLBManager:
            self._server_args, self._model_runner.model_config, logical_count
        )

+        from sglang.srt.model_executor.model_runner_components.moe_ep_setup import (
+            init_lplb_solvers,
+        )
+
        update_layer_ids_chunks = self._compute_update_layer_ids_chunks()
        all_update_layer_ids = [
            layer_id for chunk in update_layer_ids_chunks for layer_id in chunk
@@ -107,7 +111,9 @@ class EPLBManager:
                expert_backup_client=self._model_runner.expert_backup_client,
                update_weights_from_disk_callable=self._model_runner.weight_updater.update_weights_from_disk,
                ep_dispatch_algorithm=self._model_runner.server_args.ep_dispatch_algorithm,
-                init_lplb_solvers_callable=self._model_runner._init_lplb_solvers,
+                init_lplb_solvers_callable=lambda: init_lplb_solvers(
+                    model_config=self._model_runner.model_config
+                ),
            )

        self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)