[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 (