[AMD] Support eplb for moriep (#22985)

Co-authored-by: HAI <hixiao@gmail.com>

python/sglang/srt/eplb/expert_distribution.py

@@ -311,6 +311,9 @@ class _SinglePassGatherer(ABC):
                 server_args, expert_location_metadata, rank
             )
 
+        if server_args.moe_a2a_backend == "mori":
+            return _DeepepLowLatencySinglePassGatherer(expert_location_metadata, rank)
+
         if server_args.expert_distribution_recorder_mode == "stat_approx":
             if server_args.moe_a2a_backend != "none" and (
                 server_args.deepep_mode == "normal"

python/sglang/srt/eplb/expert_location_dispatch.py

@@ -19,6 +19,9 @@ import torch
 
 from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
 from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
 
 
 @dataclass
@@ -89,7 +92,10 @@ def topk_ids_logical_to_physical(
 def _topk_ids_logical_to_physical_static(
     topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
 ) -> torch.Tensor:
-    return info.partial_logical_to_rank_dispatch_physical_map[topk_ids]
+    physical_topk_ids = info.partial_logical_to_rank_dispatch_physical_map[topk_ids]
+    if _is_hip:
+        physical_topk_ids = physical_topk_ids.to(topk_ids.dtype)
+    return physical_topk_ids
 
 
 def _topk_ids_logical_to_physical_dynamic(
@@ -104,6 +110,8 @@ def _topk_ids_logical_to_physical_dynamic(
         % info.partial_logical_to_all_physical_map_num_valid[topk_ids]
     )
     topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index]
+    if _is_hip:
+        topk_ids = topk_ids.to(topk_ids.dtype)
 
     topk_ids = topk_ids.view(topk_ids_original_shape)
     return topk_ids

python/sglang/srt/eplb/expert_location_updater.py

@@ -26,13 +26,15 @@ from sglang.srt.eplb.expert_location import (
     get_global_expert_location_metadata,
 )
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils import get_bool_env_var, get_int_env_var, is_hip
 
 logger = logging.getLogger(__name__)
 
 
 _LOG_INPUT = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_INPUT")
 
+_is_hip = is_hip()
+
 
 class ExpertLocationUpdater:
     def __init__(self):
@@ -483,9 +485,31 @@ def update_expert_weights_single_layer(
         if len(p2p_ops) == 0:
             return
 
-        reqs = torch.distributed.batch_isend_irecv(p2p_ops)
-        for req in reqs:
-            req.wait()
+        if _is_hip:
+            # Submit P2P ops in batches to prevent RCCL GPU-side
+            # accumulation hangs. All ranks use the same expert_id ranges
+            # (based on num_physical_experts) to ensure matching send/recv
+            # pairs land in the same batch. Setting batch_chunk_size >=
+            # num_physical_experts disables batching behavior.
+            batch_chunk_size = get_int_env_var(
+                "SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE", 32
+            )
+            ops_by_expert = {eid: ops for eid, ops in sorted_infos}
+            for start in range(0, num_physical_experts, batch_chunk_size):
+                batch_ops = []
+                for eid in range(
+                    start, min(start + batch_chunk_size, num_physical_experts)
+                ):
+                    if eid in ops_by_expert:
+                        batch_ops.extend(ops_by_expert[eid])
+                if batch_ops:
+                    reqs = torch.distributed.batch_isend_irecv(batch_ops)
+                    for req in reqs:
+                        req.wait()
+        else:
+            reqs = torch.distributed.batch_isend_irecv(p2p_ops)
+            for req in reqs:
+                req.wait()
 
     def _execute_buffer2weight_copies(buffer2weight_copy_infos):
         for (

python/sglang/srt/layers/moe/moe_runner/aiter.py

@@ -146,7 +146,6 @@ class AiterRunnerCore(MoeRunnerCore):
             return AiterRunnerOutput(hidden_states=runner_input.hidden_states)
 
         from aiter.fused_moe import fused_moe
-        from aiter.ops.flydsl.moe_common import GateMode
 
         from sglang.srt.environ import envs
 
@@ -164,6 +163,12 @@ class AiterRunnerCore(MoeRunnerCore):
         if runner_input.output_dtype is not None:
             extra["dtype"] = runner_input.output_dtype
         if quant_info.swiglu_limit > 0:
+            # GateMode is only needed for the gpt-oss MXFP4 swiglu_limit path.
+            # Import lazily so models that don't use it (e.g. DeepSeek-V3 fp8,
+            # swiglu_limit==0) still run on aiter builds where this module
+            # lives elsewhere / is absent.
+            from aiter.ops.flydsl.moe_common import GateMode
+
             # Default (INTERLEAVE) preserves the pre-fix behavior for paths
             # that prepare weights in the gate/up-interleaved layout. Set
             # `SGLANG_USE_AITER_MOE_GU_ITLV=0` to switch to SEPARATED, which

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

@@ -5,6 +5,7 @@ import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple
 
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.dp_attention import get_is_extend_in_batch
 from sglang.srt.layers.moe.token_dispatcher.base import (
     BaseDispatcher,
@@ -662,7 +663,13 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                     recv_scales,
                     recv_topk_ids,
                     packed_recv_count,
-                ) = dispatch_fn(hidden_states, topk_weights, scale, topk_ids)
+                ) = dispatch_fn(
+                    hidden_states,
+                    topk_weights,
+                    scale,
+                    topk_ids,
+                    call_local_expert_count=True,
+                )
                 if self.enable_sdma:
                     self.mori_op.dispatch_recv()
 
@@ -688,10 +695,21 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                 recv_scales,
                 recv_topk_ids,
                 packed_recv_count,
-            ) = self.mori_op.dispatch(hidden_states, topk_weights, scale, topk_ids)
+            ) = self.mori_op.dispatch(
+                hidden_states,
+                topk_weights,
+                scale,
+                topk_ids,
+                call_local_expert_count=True,
+            )
 
-        # TODO(billishyahao): EPLB
-        # get_global_expert_distribution_recorder().on_deepep_dispatch_normal(
+        # 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
+        )
 
         return (
             packed_recv_hidden,
@@ -870,7 +888,11 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
             is mori.ops.EpDispatchCombineKernelType.AsyncLL
         ), "mori asyncll mismatch"
 
-        self.mori_op.dispatch_recv()
+        self.mori_op.dispatch_recv(call_local_expert_count=True)
+
+        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

@@ -1506,6 +1506,9 @@ 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