Support layerwise rebalancing experts (#6851)

test/srt/test_eplb.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import sglang as sgl
-from sglang.srt.managers.expert_distribution_storage import ExpertDistributionStorage
 from sglang.srt.utils import kill_process_tree
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -17,7 +16,9 @@ from sglang.test.test_utils import (
 )
 
 
-class TestDynamicEPLB(CustomTestCase):
+class _BaseTestDynamicEPLB(CustomTestCase):
+    extra_args = []
+
     @classmethod
     def setUpClass(cls):
         cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
@@ -51,8 +52,13 @@ class TestDynamicEPLB(CustomTestCase):
                 "stat",
                 "--ep-dispatch-algorithm",
                 "static",
+                *cls.extra_args,
             ],
-            env={"SGL_ENABLE_JIT_DEEPGEMM": "0", **os.environ},
+            env={
+                "SGL_ENABLE_JIT_DEEPGEMM": "0",
+                "SGLANG_EXPERT_LOCATION_UPDATER_CANARY": "1",
+                **os.environ,
+            },
         )
 
     @classmethod
@@ -72,6 +78,14 @@ class TestDynamicEPLB(CustomTestCase):
         self.assertGreater(metrics["score"], 0.5)
 
 
+class TestDynamicEPLBSimple(_BaseTestDynamicEPLB):
+    pass
+
+
+class TestDynamicEPLBMultiChunk(_BaseTestDynamicEPLB):
+    extra_args = ["--eplb-rebalance-layers-per-chunk", "1"]
+
+
 class TestStaticEPLB(CustomTestCase):
     def test_save_expert_distribution_and_init_expert_location(self):
         os.environ["SGL_ENABLE_JIT_DEEPGEMM"] = "0"