[NPU] use standalone group for moe ep (#29030)

This commit is contained in:
Even Zhou
2026-07-10 08:49:49 +08:00
committed by GitHub
parent 295f85df08
commit b2f9a95867
2 changed files with 18 additions and 6 deletions
@@ -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 = []
@@ -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,