[AMD] Skip eplb bookkeeping and topk remap when EPLB is not in use on mori-ep / HIP (#22985) (#28188)

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

@@ -54,6 +54,11 @@ if _use_aiter:
 logger = logging.getLogger(__name__)
 
 
+def _should_record_expert_distribution() -> bool:
+    recorder = get_global_expert_distribution_recorder()
+    return recorder.recording or torch.get_device_module().is_current_stream_capturing()
+
+
 class MoriEPPDispatchHooks(DeepEPPDispatchHooks):
 
     def __call__(self, dispatcher: BaseDispatcher):
@@ -634,6 +639,8 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
     ):
         done_event: Optional[torch.cuda.Event] = None
 
+        record = _should_record_expert_distribution()
+
         if self._comm_stream:
             compute_stream = torch.cuda.current_stream()
             comm_stream = self._comm_stream  # comm stream
@@ -668,7 +675,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                     topk_weights,
                     scale,
                     topk_ids,
-                    call_local_expert_count=True,
+                    call_local_expert_count=record,
                 )
                 if self.enable_sdma:
                     self.mori_op.dispatch_recv()
@@ -700,16 +707,15 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                 topk_weights,
                 scale,
                 topk_ids,
-                call_local_expert_count=True,
+                call_local_expert_count=record,
             )
 
-        # Use low_latency hook instead of normal since mori local_expert_count is
-        # a GPU tensor, while the normal hook expects a Python list (CPU).  The
-        # low_latency path accumulates counts directly on GPU via
-        # _DeepepLowLatencySinglePassGatherer, which is CUDA-graph safe.
-        get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
-            self.mori_op.local_expert_count
-        )
+        # mori local_expert_count is a GPU tensor; route it through the
+        # low_latency hook only when the recorder is actually active.
+        if record:
+            get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
+                self.mori_op.local_expert_count
+            )
 
         return (
             packed_recv_hidden,
@@ -888,11 +894,13 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
             is mori.ops.EpDispatchCombineKernelType.AsyncLL
         ), "mori asyncll mismatch"
 
-        self.mori_op.dispatch_recv(call_local_expert_count=True)
+        record = _should_record_expert_distribution()
+        self.mori_op.dispatch_recv(call_local_expert_count=record)
 
-        get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
-            self.mori_op.local_expert_count
-        )
+        if record:
+            get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
+                self.mori_op.local_expert_count
+            )
 
         return MoriEPLLDispatchOutput(
             hidden_states=hidden_states,

python/sglang/srt/layers/moe/topk.py

@@ -1111,18 +1111,34 @@ def is_power_of_two(n):
     return n > 0 and math.log2(n).is_integer()
 
 
+def _eplb_remap_enabled() -> bool:
+    # A real logical->physical mapping only exists when EPLB is enabled, the
+    # initial expert placement is non-trivial, or there are redundant physical
+    # experts. Otherwise the map is identity and the remap must be skipped (it is
+    # both unnecessary and not well-defined over the padded region of topk_ids).
+    from sglang.srt.server_args import get_global_server_args
+
+    server_args = get_global_server_args()
+    return (
+        server_args.enable_eplb
+        or server_args.init_expert_location != "trivial"
+        or server_args.ep_num_redundant_experts > 0
+    )
+
+
 def _mask_topk_ids_padded_region(
     topk_ids: torch.Tensor,
     num_token_non_padded: Optional[torch.Tensor] = None,
+    fill_value: int = -1,
 ) -> None:
     if num_token_non_padded is None:
         return
     # TODO: let the kernel support other dtypes
-    if _is_cuda and topk_ids.dtype == torch.int32:
+    if _is_cuda and topk_ids.dtype == torch.int32 and fill_value == -1:
         mask_topk_ids(topk_ids, num_token_non_padded)
     else:
         indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
-        topk_ids[indices >= num_token_non_padded, :] = -1
+        topk_ids[indices >= num_token_non_padded, :] = fill_value
 
 
 def _zero_topk_weights_padded_region(
@@ -1506,13 +1522,22 @@ def _post_process_topk_ids(
                 topk_ids, expert_location_dispatch_info, num_token_non_padded
             )
     elif _is_hip:
-        topk_ids = _biased_grouped_topk_postprocess(
-            topk_ids, expert_location_dispatch_info, num_token_non_padded
-        )
-        # On AMD HIP, the aiter MoE kernels do not handle topk_ids=-1 safely
-        # (negative indices cause illegal memory access). Instead, zero the
-        # routing weights for padded tokens so their MoE output contributes
-        # nothing to the hidden state after the weighted sum.
+        # The logical->physical remap is only meaningful when a real
+        # expert-location mapping exists. With a trivial placement and EPLB off
+        # the map is identity, and indexing it here is both unnecessary and not
+        # well-defined for the dp-attention padded region of topk_ids; skip it.
+        # ``--ep-dispatch-algorithm fake`` is a routing benchmark artifact and
+        # does not by itself require remapping.
+        if _eplb_remap_enabled():
+            # Mask the padded region to a valid in-range id (0) before the
+            # gather so it never indexes the dispatch map out of bounds. -1 is
+            # not used here: the aiter MoE kernels do not handle topk_ids=-1.
+            _mask_topk_ids_padded_region(topk_ids, num_token_non_padded, fill_value=0)
+            topk_ids = topk_ids_logical_to_physical(
+                topk_ids, expert_location_dispatch_info
+            )
+        # On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely, so
+        # padded tokens are neutralized by zeroing their routing weights.
         _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
 
     if recorder_topk_ids is None: