diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index d5fe1ef0b..c43c95042 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 997823db4..10280df0c 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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: