[Fix] eplb: Load all local replicas of each expert

python/sglang/srt/eplb/expert_location.py

@@ -346,8 +346,8 @@ def logical_to_all_physical(
         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
-        cpu_map = self.logical_to_all_physical_map_cpu
+        # Use CPU copy to avoid GPU-to-CPU sync on every call, which is expensive in update weights scenario
+        cpu_map = self.physical_to_logical_map_cpu
         # Draft workers can query MoE layers whose layer_id lies beyond the
         # target-sized expert map; fall back to the identity mapping (no EPLB
         # rebalancing for those layers) instead of indexing out of range.
@@ -367,11 +367,14 @@ def logical_to_all_physical(
             return list(
                 range(logical_expert_id, num_physical_experts, self.num_logical_experts)
             )
-        return [
-            physical_expert_id
-            for physical_expert_id in cpu_map[layer_id, logical_expert_id].tolist()
-            if physical_expert_id != -1
-        ]
+        if not 0 <= logical_expert_id < self.num_logical_experts:
+            raise IndexError(
+                f"{logical_expert_id=} is not one of the "
+                f"{self.num_logical_experts} logical experts"
+            )
+        # Every replica on every rank: logical_to_all_physical_map may narrow the
+        # dispatch candidates to one replica, but the loaders must fill every slot.
+        return (cpu_map[layer_id] == logical_expert_id).nonzero().flatten().tolist()
 
 
 def format_expert_location_layout(

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -1160,6 +1160,11 @@ def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
         else:
             return -1
 
+    def _local_routed_expert_ids(self, physical_expert_ids: List[int]) -> List[int]:
+        start_idx = self._expert_storage_rank * self._num_local_routed
+        end_idx = start_idx + self._num_local_routed
+        return [i for i in physical_expert_ids if start_idx <= i < end_idx]
+
     def weight_loader(
         self,
         param: torch.nn.Parameter,
@@ -1227,6 +1232,10 @@ def weight_loader(
                     self.layer_id, expert_id, require_global_experts
                 )
             )
+            if not require_global_experts:
+                # Before the local mapping: after an elastic EP scale-up, another rank's
+                # routed slot can lie past _num_global_routed and map to the shared slot.
+                physical_expert_ids = self._local_routed_expert_ids(physical_expert_ids)
 
         for physical_expert_id in physical_expert_ids:
             self._weight_loader_physical(

python/sglang/srt/layers/quantization/moe_wna16.py

@@ -235,11 +235,12 @@ def _local_expert_ids(layer, expert_id: int) -> List[int]:
     expert-parallel slice this rank stores.
     """
     metadata = get_global_expert_location_metadata()
-    physical = (
-        [expert_id]
-        if metadata is None or not layer.uses_expert_location
-        else metadata.logical_to_all_physical(layer.layer_id, expert_id)
-    )
+    if metadata is None or not layer.uses_expert_location:
+        physical = [expert_id]
+    else:
+        physical = metadata.logical_to_all_physical(layer.layer_id, expert_id)
+        if expert_id < metadata.num_logical_experts:
+            physical = layer._local_routed_expert_ids(physical)
     local = (layer._map_global_expert_id_to_local_expert_id(p) for p in physical)
     return [i for i in local if 0 <= i < layer.num_local_experts]