[AMD] Support eplb for moriep (#22985)

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

docs_new/docs/references/environment_variables.mdx

@@ -58,6 +58,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The interval of passes to collect the metric of selected count of physical experts on each layer and GPU rank. 0 means disabled.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>0</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of logical expert IDs per batch when submitting P2P ops during EPLB rebalance on ROCm. Smaller values prevent RCCL GPU-side accumulation hangs but increase overhead.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>32</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_FORWARD_UNKNOWN_TOOLS</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Forward unknown tool calls to clients instead of dropping them</td>

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,6 +485,28 @@ def update_expert_weights_single_layer(
         if len(p2p_ops) == 0:
             return
 
+        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()

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
+        # Use low_latency hook instead of normal since mori local_expert_count is
-        # get_global_expert_distribution_recorder().on_deepep_dispatch_normal(
+        # 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

test/manual/ep/test_eplb_mori.py

@@ -0,0 +1,242 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.server_args import ZMQ_TCP_PORT_DELTA
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.network import is_port_available
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.test_utils import (
+    DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+def wait_all_ports_release(base_url, timeout_s=60):
+    import time
+
+    port = int(base_url.split(":")[-1])
+    offsets = [
+        0,
+        ZMQ_TCP_PORT_DELTA,
+        ZMQ_TCP_PORT_DELTA + 1,
+        ZMQ_TCP_PORT_DELTA + 2,
+        ZMQ_TCP_PORT_DELTA + 3,
+        ZMQ_TCP_PORT_DELTA + 4,
+    ]
+    for _ in range(timeout_s):
+        if all(is_port_available(port + off) for off in offsets):
+            return
+        time.sleep(1)
+    print(f"Warning: some ports still occupied after {timeout_s}s")
+
+
+mori_env = {
+    **os.environ,
+    "SGLANG_USE_AITER": "1",
+    "SGLANG_MORI_DISPATCH_DTYPE": "bf16",
+    "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "4096",
+    "SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE": "32",
+    "MORI_SHMEM_MODE": "ISOLATION",
+}
+
+common_args = [
+    "--tp-size",
+    "8",
+    "--ep-size",
+    "8",
+    "--dp-size",
+    "8",
+    "--enable-dp-attention",
+    "--moe-a2a-backend",
+    "mori",
+    "--trust-remote-code",
+    "--load-balance-method",
+    "round_robin",
+    "--moe-dense-tp-size",
+    "1",
+    "--enable-dp-lm-head",
+    "--mem-fraction-static",
+    "0.6",
+    "--chunked-prefill-size",
+    "32768",
+    "--max-running-requests",
+    "128",
+    "--context-length",
+    "12288",
+    "--attention-backend",
+    "aiter",
+    "--cuda-graph-max-bs",
+    "32",
+]
+
+eplb_args = [
+    "--enable-eplb",
+    "--ep-num-redundant-experts",
+    "32",
+    "--eplb-rebalance-num-iterations",
+    "50",
+    "--expert-distribution-recorder-buffer-size",
+    "50",
+    "--ep-dispatch-algorithm",
+    "static",
+]
+
+
+mtp_args = [
+    "--speculative-algo",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+]
+
+
+class TestEPLBMoriStat(CustomTestCase):
+    """EPLB with mori backend, stat mode (on_select_experts path)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+class TestEPLBMoriStatApprox(CustomTestCase):
+    """EPLB with mori backend, stat_approx mode (local_expert_count kernel)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat_approx",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+class TestEPLBMoriMultiChunk(CustomTestCase):
+    """EPLB with mori backend, chunked layer updates."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat",
+                "--eplb-rebalance-layers-per-chunk",
+                "1",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/amd/test_moriep_small.py

@@ -76,6 +76,18 @@ common_args = [
     "32",
 ]
 
+eplb_args = [
+    "--enable-eplb",
+    "--ep-num-redundant-experts",
+    "32",
+    "--eplb-rebalance-num-iterations",
+    "50",
+    "--expert-distribution-recorder-buffer-size",
+    "50",
+    "--ep-dispatch-algorithm",
+    "static",
+]
+
 mtp_args = [
     "--speculative-algo",
     "EAGLE",
@@ -507,5 +519,61 @@ class TestMTPwithTBOLowLatency(CustomTestCase):
         self.assertGreaterEqual(avg_spec_accept_length, 2.8)
 
 
+class TestEPLBMoriStat(CustomTestCase):
+    """EPLB with mori backend, stat mode (on_select_experts path)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat",
+            ]
+        )
+
+        env = dict(os.environ)
+        env["SGLANG_USE_AITER"] = "1"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "128"
+        env["SGLANG_ENABLE_SPEC_V2"] = "false"
+        env["SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE"] = "32"
+        env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
+        # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
+        # env["MORI_DISABLE_P2P"] = "1"
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
 if __name__ == "__main__":
     unittest.main()