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

Extract MoE/EP setup into a moe_ep_setup module (#31159)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 16:00:24 +08:00
committed by GitHub
parent 440aebdfe0
commit 08798dba0d
3 changed files with 175 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/eplb_manager.py
+7 -1
View File
@@ -86,6 +86,10 @@ class EPLBManager:
self._server_args, self._model_runner.model_config, logical_count
)
from sglang.srt.model_executor.model_runner_components.moe_ep_setup import (
init_lplb_solvers,
)
update_layer_ids_chunks = self._compute_update_layer_ids_chunks()
all_update_layer_ids = [
layer_id for chunk in update_layer_ids_chunks for layer_id in chunk
@@ -107,7 +111,9 @@ class EPLBManager:
expert_backup_client=self._model_runner.expert_backup_client,
update_weights_from_disk_callable=self._model_runner.weight_updater.update_weights_from_disk,
ep_dispatch_algorithm=self._model_runner.server_args.ep_dispatch_algorithm,
init_lplb_solvers_callable=self._model_runner._init_lplb_solvers,
init_lplb_solvers_callable=lambda: init_lplb_solvers(
model_config=self._model_runner.model_config
),
)
self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
python/sglang/srt/model_executor/model_runner.py
+19 -123
View File
@@ -65,12 +65,6 @@ from sglang.srt.eplb.expert_location import (
set_global_expert_location_metadata,
)
from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
from sglang.srt.eplb.lplb_solver import (
LPLBSolver,
assert_lplb_supported_model,
clear_global_lplb_solvers,
set_global_lplb_solver,
)
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.hardware_backend.xpu.graph_runner.xpu_graph_runner import XPUGraphRunner
from sglang.srt.kv_canary.api import install_canary
@@ -87,8 +81,6 @@ from sglang.srt.layers.cp.utils import (
get_cp_strategy,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.hash_topk import HashTopK
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.sampler import create_sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
@@ -134,6 +126,11 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import (
report_online_quantization,
resolve_sliding_window_size,
)
from sglang.srt.model_executor.model_runner_components.moe_ep_setup import (
check_quantized_moe_compatibility,
init_lplb_solvers,
prepare_moe_topk,
)
from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager import (
NgramEmbeddingManager,
)
@@ -188,7 +185,6 @@ from sglang.srt.utils import (
cpu_has_amx_support,
enable_show_time_cost,
get_available_gpu_memory,
get_bool_env_var,
init_cublas,
is_host_cpu_arm64,
is_npu,
@@ -214,7 +210,6 @@ from sglang.srt.utils.weight_checker import WeightChecker
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu_arm64 = is_host_cpu_arm64()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
@@ -512,7 +507,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
if self.server_args.ep_dispatch_algorithm == "lp" and not self.is_draft_worker:
self._init_lplb_solvers()
init_lplb_solvers(model_config=self.model_config)
# Expert parallelism
self.eplb_manager = (
@@ -531,7 +526,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Load the model
self.sampler = create_sampler()
self.load_model()
self._prepare_moe_topk()
prepare_moe_topk(
model=self.model,
model_config=self.model_config,
server_args=self.server_args,
moe_ep_size=self.moe_ep_size,
moe_ep_rank=self.moe_ep_rank,
)
# Must run before backend/graph init so no draft graph records a
# routed-experts capture-write kernel.
@@ -824,42 +825,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
def check_quantized_moe_compatibility(self):
if (
quantization_config := getattr(
self.model_config.hf_config, "quantization_config", None
)
) is not None and (
weight_block_size := quantization_config.get("weight_block_size", None)
) is not None:
weight_block_size_n = weight_block_size[0]
if self.tp_size % self.moe_ep_size != 0:
raise ValueError(
f"tp_size {self.tp_size} must be divisible by ep_size {self.moe_ep_size}"
)
moe_tp_size = self.tp_size // self.moe_ep_size // self.moe_dp_size
moe_intermediate_size = getattr(
self.model_config.hf_text_config, "moe_intermediate_size", None
)
if moe_intermediate_size is None:
return
if moe_intermediate_size % moe_tp_size != 0:
raise ValueError(
f"moe_intermediate_size {moe_intermediate_size} must be divisible by moe_tp_size ({moe_tp_size}) which is tp_size ({self.tp_size}) divided by moe_ep_size ({self.moe_ep_size})."
)
if (
not envs.SGLANG_SHARED_EXPERT_TP1.get()
and (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0
and not _use_aiter
):
raise ValueError(
f"For quantized MoE models, please make sure ({moe_intermediate_size=} / {moe_tp_size=}) % {weight_block_size_n=} == 0 "
f"where moe_tp_size is equal to tp_size ({self.tp_size}) divided by ep_size ({self.moe_ep_size}). "
f"You can fix this by setting arguments `--tp` and `--ep` correctly."
)
check_quantized_moe_compatibility(
model_config=self.model_config,
tp_size=self.tp_size,
moe_ep_size=self.moe_ep_size,
moe_dp_size=self.moe_dp_size,
)
def init_torch_distributed(self):
result = bootstrap.init_torch_distributed(
@@ -1044,81 +1015,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
tp_rank=self.tp_rank,
)
def _prepare_moe_topk(self):
balancer_cls = None
num_prepared = 0
num_routed_experts = None
for module in self.model.modules():
if not isinstance(module, (TopK, HashTopK)):
continue
if not module.enable_waterfill or module.waterfill_balancer is not None:
continue
if num_routed_experts is None:
num_routed_experts = getattr(
self.model_config.hf_config, "n_routed_experts", None
)
if num_routed_experts is None:
raise ValueError(
"Waterfill requires model config n_routed_experts."
)
if balancer_cls is None:
from sglang.srt.layers.moe.waterfill import WaterfillBalancer
balancer_cls = WaterfillBalancer
# Static EPLB remaps TopK ids to physical expert ids before Waterfill.
# Redundant experts therefore need to be included in the per-rank
# expert count used for Waterfill's shared-expert slot remapping.
num_physical_routed_experts = (
num_routed_experts + self.server_args.ep_num_redundant_experts
)
if isinstance(module, TopK):
routed_scaling_factor = module.topk_config.routed_scaling_factor
else:
routed_scaling_factor = module.routed_scaling_factor
module.waterfill_balancer = balancer_cls(
num_routed_experts=num_physical_routed_experts,
world_size=self.moe_ep_size,
rank=self.moe_ep_rank,
layer_id=module.layer_id,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
)
num_prepared += 1
if num_prepared:
log_info_on_rank0(
logger, f"Prepared {num_prepared} Waterfill TopK modules."
)
def _init_lplb_solvers(self):
"""Initialize per-layer LPLB solvers from current expert location metadata."""
from sglang.srt.distributed import get_moe_ep_group
# Gate: refuse LP for non-DeepSeek MoE families whose empty-token paths
# don't participate in the EP all-reduce (would deadlock under DP-
# attention). Failure here happens before any forward pass.
architectures = getattr(self.model_config.hf_config, "architectures", None)
if architectures:
assert_lplb_supported_model(architectures[0])
metadata = get_global_expert_location_metadata()
if metadata is None:
return
clear_global_lplb_solvers()
ep_group = get_moe_ep_group()
for lid in range(metadata.num_layers):
solver = LPLBSolver(
phy2log=metadata.physical_to_logical_map[lid],
log2phy=metadata.logical_to_all_physical_map[lid],
num_gpus=metadata.ep_size,
ep_group=ep_group,
logical_to_all_physical_map_num_valid=(
metadata.logical_to_all_physical_map_num_valid[lid]
),
)
set_global_lplb_solver(lid, solver)
logger.info(f"Initialized LPLB solvers for {metadata.num_layers} layers")
def maybe_recover_ep_ranks(self):
# TODO(perf): `active_ranks.all()` on a CUDA tensor triggers host-device
# synchronization, and this function is on the forward-path.
python/sglang/srt/model_executor/model_runner_components/moe_ep_setup.py
+149
View File
@@ -0,0 +1,149 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
from sglang.srt.eplb.lplb_solver import (
LPLBSolver,
assert_lplb_supported_model,
clear_global_lplb_solvers,
set_global_lplb_solver,
)
from sglang.srt.layers.moe.hash_topk import HashTopK
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.utils import get_bool_env_var, is_hip, log_info_on_rank0
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
def prepare_moe_topk(
*,
model,
model_config: ModelConfig,
server_args: ServerArgs,
moe_ep_size: int,
moe_ep_rank: int,
) -> None:
balancer_cls = None
num_prepared = 0
num_routed_experts = None
for module in model.modules():
if not isinstance(module, (TopK, HashTopK)):
continue
if not module.enable_waterfill or module.waterfill_balancer is not None:
continue
if num_routed_experts is None:
num_routed_experts = getattr(
model_config.hf_config, "n_routed_experts", None
)
if num_routed_experts is None:
raise ValueError("Waterfill requires model config n_routed_experts.")
if balancer_cls is None:
from sglang.srt.layers.moe.waterfill import WaterfillBalancer
balancer_cls = WaterfillBalancer
# Static EPLB remaps TopK ids to physical expert ids before Waterfill.
# Redundant experts therefore need to be included in the per-rank
# expert count used for Waterfill's shared-expert slot remapping.
num_physical_routed_experts = (
num_routed_experts + server_args.ep_num_redundant_experts
)
if isinstance(module, TopK):
routed_scaling_factor = module.topk_config.routed_scaling_factor
else:
routed_scaling_factor = module.routed_scaling_factor
module.waterfill_balancer = balancer_cls(
num_routed_experts=num_physical_routed_experts,
world_size=moe_ep_size,
rank=moe_ep_rank,
layer_id=module.layer_id,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
)
num_prepared += 1
if num_prepared:
log_info_on_rank0(logger, f"Prepared {num_prepared} Waterfill TopK modules.")
def init_lplb_solvers(*, model_config: ModelConfig) -> None:
"""Initialize per-layer LPLB solvers from current expert location metadata."""
from sglang.srt.distributed import get_moe_ep_group
# Gate: refuse LP for non-DeepSeek MoE families whose empty-token paths
# don't participate in the EP all-reduce (would deadlock under DP-
# attention). Failure here happens before any forward pass.
architectures = getattr(model_config.hf_config, "architectures", None)
if architectures:
assert_lplb_supported_model(architectures[0])
metadata = get_global_expert_location_metadata()
if metadata is None:
return
clear_global_lplb_solvers()
ep_group = get_moe_ep_group()
for lid in range(metadata.num_layers):
solver = LPLBSolver(
phy2log=metadata.physical_to_logical_map[lid],
log2phy=metadata.logical_to_all_physical_map[lid],
num_gpus=metadata.ep_size,
ep_group=ep_group,
logical_to_all_physical_map_num_valid=(
metadata.logical_to_all_physical_map_num_valid[lid]
),
)
set_global_lplb_solver(lid, solver)
logger.info(f"Initialized LPLB solvers for {metadata.num_layers} layers")
def check_quantized_moe_compatibility(
*,
model_config: ModelConfig,
tp_size: int,
moe_ep_size: int,
moe_dp_size: int,
) -> None:
if (
quantization_config := getattr(
model_config.hf_config, "quantization_config", None
)
) is not None and (
weight_block_size := quantization_config.get("weight_block_size", None)
) is not None:
weight_block_size_n = weight_block_size[0]
if tp_size % moe_ep_size != 0:
raise ValueError(
f"tp_size {tp_size} must be divisible by ep_size {moe_ep_size}"
)
moe_tp_size = tp_size // moe_ep_size // moe_dp_size
moe_intermediate_size = getattr(
model_config.hf_text_config, "moe_intermediate_size", None
)
if moe_intermediate_size is None:
return
if moe_intermediate_size % moe_tp_size != 0:
raise ValueError(
f"moe_intermediate_size {moe_intermediate_size} must be divisible by moe_tp_size ({moe_tp_size}) which is tp_size ({tp_size}) divided by moe_ep_size ({moe_ep_size})."
)
if (
not envs.SGLANG_SHARED_EXPERT_TP1.get()
and (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0
and not _use_aiter
):
raise ValueError(
f"For quantized MoE models, please make sure ({moe_intermediate_size=} / {moe_tp_size=}) % {weight_block_size_n=} == 0 "
f"where moe_tp_size is equal to tp_size ({tp_size}) divided by ep_size ({moe_ep_size}). "
f"You can fix this by setting arguments `--tp` and `--ep` correctly."
)