diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 5cd4f7266..1c03030ae 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -58,6 +58,11 @@ SGLang supports various environment variables that can be used to configure its The interval of passes to collect the metric of selected count of physical experts on each layer and GPU rank. 0 means disabled. 0 + + SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE + 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. + 32 + SGLANG_FORWARD_UNKNOWN_TOOLS Forward unknown tool calls to clients instead of dropping them diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 5e3adcedb..6472dca01 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/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" diff --git a/python/sglang/srt/eplb/expert_location_dispatch.py b/python/sglang/srt/eplb/expert_location_dispatch.py index 7ac03390a..533086e03 100644 --- a/python/sglang/srt/eplb/expert_location_dispatch.py +++ b/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 diff --git a/python/sglang/srt/eplb/expert_location_updater.py b/python/sglang/srt/eplb/expert_location_updater.py index d1d6387fc..52d80a15d 100644 --- a/python/sglang/srt/eplb/expert_location_updater.py +++ b/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 ( diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 2ea4e66ef..7c32cf0b8 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 5f102e267..d5fe1ef0b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index b104e40d7..f9d89828d 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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 diff --git a/test/manual/ep/test_eplb_mori.py b/test/manual/ep/test_eplb_mori.py new file mode 100644 index 000000000..1fef948af --- /dev/null +++ b/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() diff --git a/test/registered/amd/test_moriep_small.py b/test/registered/amd/test_moriep_small.py index d32cf3002..0e6780c5c 100644 --- a/test/registered/amd/test_moriep_small.py +++ b/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()