[NPU] use standalone group for moe ep (#29030)
This commit is contained in:
@@ -2239,7 +2239,8 @@ def initialize_model_parallel(
|
|||||||
|
|
||||||
global _MOE_EP
|
global _MOE_EP
|
||||||
assert _MOE_EP is None, "expert model parallel group is already initialized"
|
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
|
_MOE_EP = _TP
|
||||||
else:
|
else:
|
||||||
group_ranks = []
|
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.single_batch_overlap import DownGemmOverlapArgs
|
||||||
from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
|
from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
|
||||||
from sglang.srt.distributed import (
|
from sglang.srt.distributed import (
|
||||||
|
get_moe_ep_group,
|
||||||
get_tp_group,
|
get_tp_group,
|
||||||
tensor_model_parallel_all_reduce,
|
tensor_model_parallel_all_reduce,
|
||||||
)
|
)
|
||||||
@@ -72,6 +73,7 @@ from sglang.srt.utils import (
|
|||||||
get_bool_env_var,
|
get_bool_env_var,
|
||||||
is_cpu,
|
is_cpu,
|
||||||
is_hip,
|
is_hip,
|
||||||
|
is_npu,
|
||||||
print_info_once,
|
print_info_once,
|
||||||
round_up,
|
round_up,
|
||||||
)
|
)
|
||||||
@@ -80,9 +82,22 @@ from sglang.srt.utils.custom_op import register_custom_op
|
|||||||
_is_hip = is_hip()
|
_is_hip = is_hip()
|
||||||
_is_cpu_amx_available = cpu_has_amx_support()
|
_is_cpu_amx_available = cpu_has_amx_support()
|
||||||
_is_cpu = is_cpu()
|
_is_cpu = is_cpu()
|
||||||
|
_is_npu = is_npu()
|
||||||
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
_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:
|
def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
|
||||||
a2a_backend = get_moe_a2a_backend()
|
a2a_backend = get_moe_a2a_backend()
|
||||||
if (
|
if (
|
||||||
@@ -101,11 +116,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
|
|||||||
or a2a_backend.is_nixl()
|
or a2a_backend.is_nixl()
|
||||||
):
|
):
|
||||||
return MaybeTboDeepEPDispatcher(
|
return MaybeTboDeepEPDispatcher(
|
||||||
group=(
|
group=_get_deepep_comm_group(a2a_backend),
|
||||||
get_tp_group().device_group
|
|
||||||
if not a2a_backend.is_mori()
|
|
||||||
else get_tp_group()
|
|
||||||
),
|
|
||||||
router_topk=moe_runner_config.top_k,
|
router_topk=moe_runner_config.top_k,
|
||||||
permute_fusion=True,
|
permute_fusion=True,
|
||||||
num_experts=moe_runner_config.num_experts,
|
num_experts=moe_runner_config.num_experts,
|
||||||
|
|||||||
Reference in New Issue
Block a user