minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Support eplb for moriep (#22985)

Browse Source 
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
billishyahao
2026-06-10 10:23:51 -07:00
committed by GitHub
co-authored by HAI
parent 91ff7baa28
commit 0ae27405d0
9 changed files with 391 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/references/environment_variables.mdx
+5
View File
@@ -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
+3
View File
@@ -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
+9 -1
View File
@@ -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
+28 -4
View File
@@ -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
+6 -1
View File
@@ -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
+27 -5
View File
@@ -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
+3
View File
@@ -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
+242
View File
@@ -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
+68
View File
@@ -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()