diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index 0fdff48fe..7bd0254ba 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/python/sglang/srt/eplb/expert_location.py @@ -284,9 +284,19 @@ class ExpertLocationMetadata: # -------------------------------- usage ------------------------------------ def logical_to_all_physical( - self, layer_id: int, logical_expert_id: int + self, + layer_id: int, + logical_expert_id: int, + require_global_experts: bool = False, ) -> List[int]: # Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario + if require_global_experts: + num_physical_experts = self.logical_to_all_physical_map_cpu[layer_id].shape[ + -1 + ] + return list( + range(logical_expert_id, num_physical_experts, self.num_logical_experts) + ) return [ physical_expert_id for physical_expert_id in self.logical_to_all_physical_map_cpu[ @@ -355,14 +365,10 @@ def _compute_logical_to_all_physical_map( ) # Replace by the nearest physical expert - mapped_physical_experts = logical_to_all_physical_map[layer_id][ - logical_expert_id - ] - if ( - nearest_expert != -1 - and nearest_expert not in mapped_physical_experts - ): - mapped_physical_experts[0] = nearest_expert + if nearest_expert != -1: + logical_to_all_physical_map[layer_id][logical_expert_id] = [ + nearest_expert + ] logical_to_all_physical_map = _pad_nested_array( logical_to_all_physical_map, pad_value=-1 diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 47a9da781..d7fdb1079 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -548,9 +548,12 @@ class FusedMoE(torch.nn.Module): # This is a shared expert. physical_expert_ids = [expert_id] else: + require_global_experts = getattr( + param, "_sglang_require_global_experts", False + ) physical_expert_ids = ( global_expert_location_metadata.logical_to_all_physical( - self.layer_id, expert_id + self.layer_id, expert_id, require_global_experts ) )