diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 73a64ef9b..72a87476d 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -2239,7 +2239,8 @@ def initialize_model_parallel( global _MOE_EP assert _MOE_EP is None, "expert model parallel group is already initialized" - if moe_ep_size == tensor_model_parallel_size: + # NPU requires a standalone group for MOE expert parallelism + if moe_ep_size == tensor_model_parallel_size and not _is_npu: _MOE_EP = _TP else: group_ranks = [] diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 289907dd3..e4f9c88fc 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -13,6 +13,7 @@ from torch.nn.parameter import UninitializedParameter from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher from sglang.srt.distributed import ( + get_moe_ep_group, get_tp_group, tensor_model_parallel_all_reduce, ) @@ -72,6 +73,7 @@ from sglang.srt.utils import ( get_bool_env_var, is_cpu, is_hip, + is_npu, print_info_once, round_up, ) @@ -80,9 +82,22 @@ from sglang.srt.utils.custom_op import register_custom_op _is_hip = is_hip() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() +_is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +def _get_deepep_comm_group(a2a_backend): + group = get_tp_group().device_group + + if a2a_backend.is_mori(): + group = get_tp_group() + + elif _is_npu: + group = get_moe_ep_group().device_group + + return group + + def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: a2a_backend = get_moe_a2a_backend() if ( @@ -101,11 +116,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: or a2a_backend.is_nixl() ): return MaybeTboDeepEPDispatcher( - group=( - get_tp_group().device_group - if not a2a_backend.is_mori() - else get_tp_group() - ), + group=_get_deepep_comm_group(a2a_backend), router_topk=moe_runner_config.top_k, permute_fusion=True, num_experts=moe_runner_config.num_experts,