Support EPLB in FusedMoE (#8448)

python/sglang/srt/eplb/expert_distribution.py

@@ -47,6 +47,11 @@ class ExpertDistributionRecorder(ABC):
         rank: int,
     ):
         if server_args.expert_distribution_recorder_mode is not None:
+            assert (
+                expert_location_metadata is not None
+            ), "ExpertLocationMetadata is required for expert distribution recording. One possible"
+            "reason is that you are using a model that does not support expert distribution"
+            "recording. Try setting `get_model_config_for_expert_location` in your model."
             return _ExpertDistributionRecorderReal(
                 server_args, expert_location_metadata, rank
             )

python/sglang/srt/eplb/expert_location.py

@@ -82,6 +82,10 @@ class ExpertLocationMetadata:
     def init_trivial(server_args: ServerArgs, model_config: ModelConfig):
         """Trivial location - logical expert i corresponds to physical expert i"""
         common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
         num_physical_experts = common["num_physical_experts"]
         model_config_for_expert_location = common["model_config_for_expert_location"]
         num_layers = model_config_for_expert_location.num_layers
@@ -109,6 +113,10 @@ class ExpertLocationMetadata:
         physical_to_logical_map = physical_to_logical_map.to(server_args.device)
 
         common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
         model_config_for_expert_location = common["model_config_for_expert_location"]
         logical_to_all_physical_map = _compute_logical_to_all_physical_map(
             physical_to_logical_map,
@@ -133,6 +141,10 @@ class ExpertLocationMetadata:
         logical_count = logical_count.to(server_args.device)
 
         common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
         model_config_for_expert_location = common["model_config_for_expert_location"]
         num_physical_experts = common["num_physical_experts"]
         num_groups = model_config_for_expert_location.num_groups
@@ -168,6 +180,9 @@ class ExpertLocationMetadata:
             ModelConfigForExpertLocation.from_model_config(model_config)
         )
 
+        if model_config_for_expert_location is None:
+            return None
+
         num_physical_experts = (
             model_config_for_expert_location.num_logical_experts
             + server_args.ep_num_redundant_experts
@@ -398,10 +413,6 @@ class ModelConfigForExpertLocation:
     num_logical_experts: int
     num_groups: Optional[int] = None
 
-    @staticmethod
-    def init_dummy():
-        return ModelConfigForExpertLocation(num_layers=1, num_logical_experts=1)
-
     @staticmethod
     def from_model_config(model_config: ModelConfig):
         model_class, _ = get_model_architecture(model_config)
@@ -410,12 +421,12 @@ class ModelConfigForExpertLocation:
                 model_config.hf_config
             )
         else:
-            return ModelConfigForExpertLocation.init_dummy()
+            return None
 
 
 def compute_initial_expert_location_metadata(
     server_args: ServerArgs, model_config: ModelConfig
-) -> ExpertLocationMetadata:
+) -> Optional[ExpertLocationMetadata]:
     data = server_args.init_expert_location
     if data == "trivial":
         return ExpertLocationMetadata.init_trivial(server_args, model_config)

python/sglang/srt/eplb/expert_location_dispatch.py

@@ -36,6 +36,7 @@ class ExpertLocationDispatchInfo:
     def init_new(cls, layer_id: int):
         ep_dispatch_algorithm = global_server_args_dict["ep_dispatch_algorithm"]
         expert_location_metadata = get_global_expert_location_metadata()
+        assert expert_location_metadata is not None
 
         if ep_dispatch_algorithm is None:
             return None

python/sglang/srt/eplb/expert_location_updater.py

@@ -50,6 +50,8 @@ class ExpertLocationUpdater:
             torch.cuda.empty_cache()
 
         old_expert_location_metadata = get_global_expert_location_metadata()
+        assert old_expert_location_metadata is not None
+
         _update_expert_weights(
             routed_experts_weights_of_layer=routed_experts_weights_of_layer,
             old_expert_location_metadata=old_expert_location_metadata,