diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 07329b6cd993..652518387622 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -2504,6 +2504,7 @@ def _cutlass_moe_env_override(view: Any) -> dict: { "megamoe", "deepep", + "moonep", "mooncake", "nixl", "ascend_fuseep", diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 1e89fba14db7..d158fb23d0fe 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -27,6 +27,7 @@ ) from sglang.srt.layers.moe.token_dispatcher import ( DeepEPDispatcher, + MoonEPDispatcher, MooncakeEPDispatcher, MoriEPDispatcher, NixlEPDispatcher, @@ -1077,6 +1078,10 @@ def __init__(self, **kwargs): self._inners = [ DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] + elif get_moe_a2a_backend().is_moonep(): + self._inners = [ + MoonEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) + ] elif get_moe_a2a_backend().is_mooncake(): self._inners = [ MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) @@ -1120,6 +1125,9 @@ def combine_a(self, **kwargs): def combine_b(self, **kwargs): return self._execute("combine_b", **kwargs) + def prefetch_weight(self, **kwargs): + return self._execute("prefetch_weight", **kwargs) + def register_deepep_dispatch_hook(self, hook): handle_list = [] for inner in self._inners: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 09c406584254..c8a91366776e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -976,6 +976,11 @@ class Envs: SGLANG_DISABLE_STATIC_WATERFILL = EnvBool(False) SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False) SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # -1 uses MoonEP's training-safe default B = E / EP. + SGLANG_MOONEP_NUM_PREFETCH_SLOTS = EnvInt(-1) + SGLANG_MOONEP_TOKEN_PADDING = EnvInt(128) + SGLANG_MOONEP_NUM_SMS = EnvInt(32) SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True) # DeepSeek/GLM MoE (deepseek_v2.py): quantize the (dp-gathered) MoE input diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 513ba0aaa667..af20b529bd5e 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -25,6 +25,10 @@ DeepEPLLCombineInput, DeepEPNormalCombineInput, ) +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + get_moonep_expert_weight_layout, + run_moonep_bf16_expert, +) from sglang.srt.layers.moe.topk import ( StandardTopKOutput, TopKOutput, @@ -299,6 +303,26 @@ def run_moe_core( output = self.forward_cutlass_w4afp8_masked(dispatch_output) else: assert False, "forward_deepgemm_masked is deprecated" + elif DispatchOutputChecker.format_is_moonep(dispatch_output): + weight_layout = get_moonep_expert_weight_layout( + self, + num_prefetch_slots=( + int(dispatch_output.expert_ids.numel()) - self.num_experts + ), + ) + self.dispatcher.prefetch_weight( + plan=dispatch_output.plan, + weight_layout=weight_layout, + ) + return run_moonep_bf16_expert( + dispatch_output, + weight_layout, + activation=self.moe_runner_config.activation, + ) + else: + raise NotImplementedError( + f"Unsupported DeepEPMoE dispatch format: {dispatch_output.format}" + ) combine_input_wrapper = ( DeepEPNormalCombineInput @@ -354,6 +378,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): if ( get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_moonep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_pplx() 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 f9bcb949cee6..f781993652a9 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -149,6 +149,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: return StandardDispatcher(moe_runner_config) elif ( a2a_backend.is_deepep() + or a2a_backend.is_moonep() or a2a_backend.is_mooncake() or a2a_backend.is_mori() or a2a_backend.is_nixl() @@ -396,9 +397,24 @@ def __init__( f"quant_method={type(self.quant_method).__name__})." ) + moonep_global_weight_storage = get_moe_a2a_backend().is_moonep() + if moonep_global_weight_storage: + if quant_config is not None: + raise NotImplementedError( + "MoonEP PoC supports unquantized BF16 MoE weights only." + ) + if num_fused_shared_experts != 0: + raise NotImplementedError( + "MoonEP PoC does not support fused shared experts yet." + ) + self.quant_method.create_weights( layer=self, - num_experts=self.num_local_experts, + num_experts=( + self.num_experts + if moonep_global_weight_storage + else self.num_local_experts + ), hidden_size=hidden_size, intermediate_size_per_partition=self.intermediate_size_per_partition, params_dtype=params_dtype, @@ -410,6 +426,13 @@ def __init__( with_bias=with_bias, moe_intermediate_size=intermediate_size, ) + if moonep_global_weight_storage: + self.w13_weight._sglang_require_global_experts = True + self.w2_weight._sglang_require_global_experts = True + if getattr(self, "w13_weight_bias", None) is not None: + self.w13_weight_bias._sglang_require_global_experts = True + if getattr(self, "w2_weight_bias", None) is not None: + self.w2_weight_bias._sglang_require_global_experts = True self.quant_method.create_moe_runner(self, self.moe_runner_config) self.dispatcher = create_moe_dispatcher(self.moe_runner_config) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index 7f2c0942f95c..9c65901e685b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -30,6 +30,17 @@ MooncakeDispatchOutput, MooncakeEPDispatcher, ) +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPBufferKey, + MoonEPCombineInput, + MoonEPDispatcher, + MoonEPDispatchOutput, + MoonEPExpertWeightLayout, + get_moonep_expert_weight_layout, + get_moonep_num_prefetch_slots, + run_moonep_bf16_expert, +) from sglang.srt.layers.moe.token_dispatcher.moriep import ( MoriEPDispatcher, MoriEPLLCombineInput, @@ -67,6 +78,15 @@ "MooncakeCombineInput", "MooncakeDispatchOutput", "MooncakeEPDispatcher", + "MoonEPCombineInput", + "MoonEPBuffer", + "MoonEPBufferKey", + "MoonEPDispatcher", + "MoonEPDispatchOutput", + "MoonEPExpertWeightLayout", + "get_moonep_expert_weight_layout", + "get_moonep_num_prefetch_slots", + "run_moonep_bf16_expert", "MoriEPNormalDispatchOutput", "MoriEPNormalCombineInput", "MoriEPLLDispatchOutput", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 1ff2beb5bb33..cad75fba7f16 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -29,6 +29,8 @@ DeepEPNormalDispatchOutput, FlashinferCombineInput, FlashinferDispatchOutput, + MoonEPCombineInput, + MoonEPDispatchOutput, StandardCombineInput, StandardDispatchOutput, ) @@ -165,12 +167,19 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferDispatchOutput]: return dispatch_output.format.is_flashinfer() + @staticmethod + def format_is_moonep( + dispatch_output: DispatchOutput, + ) -> TypeGuard[MoonEPDispatchOutput]: + return dispatch_output.format.is_moonep() + class DispatchOutputFormat(Enum): STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" + MOONEP = "moonep" FLASHINFER = "flashinfer" ASCEND_TP = "ascend_tp" @@ -195,6 +204,9 @@ def is_deepep(self) -> bool: def is_flashinfer(self) -> bool: return self == DispatchOutputFormat.FLASHINFER + def is_moonep(self) -> bool: + return self == DispatchOutputFormat.MOONEP + @runtime_checkable class DispatchOutput(Protocol): @@ -249,11 +261,18 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferCombineInput]: return combine_input.format == CombineInputFormat.FLASHINFER + @staticmethod + def format_is_moonep( + combine_input: CombineInput, + ) -> TypeGuard[MoonEPCombineInput]: + return combine_input.format == CombineInputFormat.MOONEP + class CombineInputFormat(Enum): STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" + MOONEP = "moonep" FLASHINFER = "flashinfer" ASCEND_TP = "ascend_tp" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moonep.py b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py new file mode 100644 index 000000000000..55ab3be63268 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py @@ -0,0 +1,708 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, NamedTuple, NoReturn, Optional + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.topk import TopKOutputChecker +from sglang.srt.layers.moe.utils import DeepEPMode + + +_MOONEP_UNSUPPORTED_MESSAGE = ( + "MoonEP MoE A2A is recognized by SGLang, but the runtime dispatcher is not " + "implemented yet. MoonEP is not a drop-in DeepEP-compatible backend: it " + "returns a MoonEPCommPlan/cu_seqlens and requires MoonEP-compatible " + "contiguous symmetric-memory expert weights plus a VM-group expert GEMM." +) + + +class MoonEPDispatchOutput(NamedTuple): + """MoonEP dispatch output. + + ``plan`` is intentionally typed as ``Any`` so this module can define the + SGLang-side contract without importing the optional ``moonep`` package at + module import time. + """ + + hidden_states: torch.Tensor + route_weights_nvs: Optional[torch.Tensor] + cu_seqlens: torch.Tensor + plan: Any + expert_ids: torch.Tensor + num_tokens: int + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.MOONEP + + +assert isinstance(MoonEPDispatchOutput, DispatchOutput) + + +class MoonEPCombineInput(NamedTuple): + """MoonEP combine input.""" + + hidden_states: torch.Tensor + route_weights_nvs: Optional[torch.Tensor] + plan: Any + num_tokens: int + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.MOONEP + + +assert isinstance(MoonEPCombineInput, CombineInput) + + +class MoonEPExpertWeightLayout(NamedTuple): + """Contiguous BF16 expert weights in MoonEP prefetch layout.""" + + full_gate_weight: torch.Tensor + full_up_weight: torch.Tensor + full_down_weight: torch.Tensor + num_prefetch_slots: int + + +@dataclass(frozen=True) +class MoonEPBufferKey: + """Static MoonEP buffer dimensions. + + MoonEP allocates its communication buffers from static shape parameters, + unlike DeepEP's normal-dispatch path. Keep the dimensions explicit so the + process-wide facade never reuses a buffer with incompatible token capacity, + model shape, EP topology, or prefetch-slot layout. + """ + + num_max_dispatch_tokens_per_rank: int + hidden_size: int + router_topk: int + num_experts: int + num_ep_ranks: int + group_id: int + num_prefetch_slots: int + token_padding: int + num_sms: int + + +class MoonEPBuffer: + """Process-wide facade for MoonEP communication buffers. + + The underlying ``moonep.Buffer`` owns NVLink/VMM allocations and is keyed by + MoonEP's static allocation dimensions. The state lives on + ``ctx.resources.buffers`` so tests can reset it with ``reset_context()`` and + future runtime code has one lifecycle hook per process. + """ + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get("moonep_ep_state") + if state is None: + state = SimpleNamespace( + buffers={}, + active_key=None, + ) + buffers["moonep_ep_state"] = state + return state + + @staticmethod + def _require_positive_int(name: str, value: int) -> int: + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + return value + + @staticmethod + def _resolve_num_ep_ranks(group: dist.ProcessGroup) -> int: + try: + num_ep_ranks = dist.get_world_size(group=group) + except (AssertionError, RuntimeError, TypeError, ValueError): + group_size = getattr(group, "size", None) + if not callable(group_size): + raise + num_ep_ranks = group_size() + return MoonEPBuffer._require_positive_int("num_ep_ranks", int(num_ep_ranks)) + + @staticmethod + def _resolve_num_prefetch_slots( + num_prefetch_slots: int | None, + num_experts: int, + num_ep_ranks: int, + ) -> int: + if num_experts % num_ep_ranks != 0: + raise ValueError( + "MoonEP requires num_experts to be divisible by the EP group size: " + f"num_experts={num_experts}, num_ep_ranks={num_ep_ranks}" + ) + + if num_prefetch_slots is None: + num_prefetch_slots = envs.SGLANG_MOONEP_NUM_PREFETCH_SLOTS.get() + num_prefetch_slots = int(num_prefetch_slots) + if num_prefetch_slots <= 0: + return num_experts // num_ep_ranks + return MoonEPBuffer._require_positive_int( + "num_prefetch_slots", num_prefetch_slots + ) + + @classmethod + def build_key( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_experts: int, + num_max_dispatch_tokens_per_rank: int | None = None, + num_prefetch_slots: int | None = None, + token_padding: int | None = None, + num_sms: int | None = None, + ) -> MoonEPBufferKey: + if num_max_dispatch_tokens_per_rank is None: + num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + if token_padding is None: + token_padding = envs.SGLANG_MOONEP_TOKEN_PADDING.get() + if num_sms is None: + num_sms = envs.SGLANG_MOONEP_NUM_SMS.get() + + num_ep_ranks = cls._resolve_num_ep_ranks(group) + num_experts = cls._require_positive_int("num_experts", int(num_experts)) + num_prefetch_slots = cls._resolve_num_prefetch_slots( + num_prefetch_slots, + num_experts, + num_ep_ranks, + ) + + return MoonEPBufferKey( + num_max_dispatch_tokens_per_rank=cls._require_positive_int( + "num_max_dispatch_tokens_per_rank", + int(num_max_dispatch_tokens_per_rank), + ), + hidden_size=cls._require_positive_int("hidden_size", int(hidden_size)), + router_topk=cls._require_positive_int("router_topk", int(router_topk)), + num_experts=num_experts, + num_ep_ranks=num_ep_ranks, + group_id=id(group), + num_prefetch_slots=num_prefetch_slots, + token_padding=cls._require_positive_int( + "token_padding", int(token_padding) + ), + num_sms=cls._require_positive_int("num_sms", int(num_sms)), + ) + + @classmethod + def get_existing_buffer( + cls, + key: MoonEPBufferKey | None = None, + ): + """Return an already-created buffer, if any. + + Without a key this returns the most recently requested MoonEP buffer. + Future runtime code should pass an explicit key when multiple static + capacities are in play. + """ + + state = cls._state() + if key is None: + key = state.active_key + if key is None: + return None + return state.buffers.get(key) + + @classmethod + def get_moonep_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_experts: int, + num_max_dispatch_tokens_per_rank: int | None = None, + num_prefetch_slots: int | None = None, + token_padding: int | None = None, + num_sms: int | None = None, + ): + key = cls.build_key( + group=group, + hidden_size=hidden_size, + router_topk=router_topk, + num_experts=num_experts, + num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, + num_prefetch_slots=num_prefetch_slots, + token_padding=token_padding, + num_sms=num_sms, + ) + + state = cls._state() + buffer = state.buffers.get(key) + if buffer is not None: + state.active_key = key + return buffer + + try: + from moonep import Buffer + except ImportError as exc: + raise ImportError( + "MoonEP is not installed. Install MoonEP before running SGLang " + "with --moe-a2a-backend moonep." + ) from exc + + buffer = Buffer( + S=key.num_max_dispatch_tokens_per_rank, + H=key.hidden_size, + K=key.router_topk, + E=key.num_experts, + num_ep_ranks=key.num_ep_ranks, + num_sms=key.num_sms, + token_padding=key.token_padding, + B=key.num_prefetch_slots, + group=group, + ) + state.buffers[key] = buffer + state.active_key = key + return buffer + + @classmethod + def destroy_buffer(cls, key: MoonEPBufferKey | None = None) -> None: + state = cls._state() + if key is None: + key = state.active_key + if key is None: + return + + buffer = state.buffers.pop(key, None) + destroy = getattr(buffer, "destroy", None) + if callable(destroy): + destroy() + if state.active_key == key: + state.active_key = next(reversed(state.buffers), None) + + @classmethod + def destroy_all_buffers(cls) -> None: + state = cls._state() + for key in list(state.buffers): + cls.destroy_buffer(key) + state.active_key = None + + +def get_moonep_num_prefetch_slots(num_experts: int, num_ep_ranks: int) -> int: + return MoonEPBuffer._resolve_num_prefetch_slots( + num_prefetch_slots=None, + num_experts=num_experts, + num_ep_ranks=num_ep_ranks, + ) + + +def get_moonep_expert_weight_layout( + layer: torch.nn.Module, + num_prefetch_slots: int, +) -> MoonEPExpertWeightLayout: + """Return cached contiguous BF16 gate/up/down tensors for MoonEP. + + The first executable PoC path supports only unquantized BF16 weights stored + in global expert-id order. Rows ``[E, E+B)`` are mutable prefetch slots and + are intentionally preserved across calls so ``buffer.prefetch_weight`` can + fill them before the MoonEP expert runner consumes the layout. + """ + + if num_prefetch_slots <= 0: + raise ValueError( + f"num_prefetch_slots must be positive, got {num_prefetch_slots}" + ) + + quant_config = getattr(layer, "quant_config", None) + if quant_config is not None: + raise NotImplementedError( + "MoonEP PoC expert weight layout supports unquantized BF16 only." + ) + + if getattr(layer.moe_runner_config, "num_fused_shared_experts", 0) != 0: + raise NotImplementedError( + "MoonEP PoC does not support fused shared experts yet." + ) + if not getattr(layer.moe_runner_config, "is_gated", True): + raise NotImplementedError("MoonEP PoC requires gated w13 experts.") + if getattr(layer, "use_triton_kernels", False): + raise NotImplementedError( + "MoonEP PoC expects canonical [E, 2I, H] w13 layout, not " + "triton-kernels transposed weight layout." + ) + + w13_weight = layer.w13_weight + w2_weight = layer.w2_weight + if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16: + raise NotImplementedError( + "MoonEP PoC expert runner supports BF16 weights only." + ) + if not w13_weight.is_contiguous() or not w2_weight.is_contiguous(): + raise ValueError("MoonEP expert source weights must be contiguous.") + + num_experts = int(layer.num_experts) + intermediate_size = int(layer.intermediate_size_per_partition) + hidden_size = int(layer.hidden_size) + expected_w13_shape = (num_experts, 2 * intermediate_size, hidden_size) + expected_w2_shape = (num_experts, hidden_size, intermediate_size) + if tuple(w13_weight.shape) != expected_w13_shape: + raise ValueError( + "MoonEP PoC requires global w13_weight shape " + f"{expected_w13_shape}, got {tuple(w13_weight.shape)}." + ) + if tuple(w2_weight.shape) != expected_w2_shape: + raise ValueError( + "MoonEP PoC requires global w2_weight shape " + f"{expected_w2_shape}, got {tuple(w2_weight.shape)}." + ) + + cache_key = ( + num_prefetch_slots, + w13_weight.data_ptr(), + w2_weight.data_ptr(), + tuple(w13_weight.shape), + tuple(w2_weight.shape), + ) + cache = getattr(layer, "_moonep_weight_layout_cache", None) + if cache is not None and cache[0] == cache_key: + return cache[1] + + full_gate_weight = torch.empty( + num_experts + num_prefetch_slots, + intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=w13_weight.device, + ) + full_up_weight = torch.empty_like(full_gate_weight) + full_down_weight = torch.empty( + num_experts + num_prefetch_slots, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=w2_weight.device, + ) + + full_gate_weight[:num_experts].copy_(w13_weight[:, :intermediate_size, :]) + full_up_weight[:num_experts].copy_( + w13_weight[:, intermediate_size : 2 * intermediate_size, :] + ) + full_down_weight[:num_experts].copy_(w2_weight) + full_gate_weight[num_experts:].zero_() + full_up_weight[num_experts:].zero_() + full_down_weight[num_experts:].zero_() + + layout = MoonEPExpertWeightLayout( + full_gate_weight=full_gate_weight.contiguous(), + full_up_weight=full_up_weight.contiguous(), + full_down_weight=full_down_weight.contiguous(), + num_prefetch_slots=num_prefetch_slots, + ) + layer._moonep_weight_layout_cache = (cache_key, layout) + return layout + + +def run_moonep_bf16_expert( + dispatch_output: MoonEPDispatchOutput, + weight_layout: MoonEPExpertWeightLayout, + *, + activation: str = "silu", +) -> MoonEPCombineInput: + """Run a simple BF16 MoonEP expert core over ``cu_seqlens`` segments. + + This is a correctness-first PoC runner. It consumes MoonEP's already + expert-grouped ``[NvS, H]`` token layout, applies gate/up/down expert + weights for each non-empty `cu_seqlens` segment, multiplies each dispatched + row by its route weight, and returns a `MoonEPCombineInput` for + `MoonEPDispatcher.combine`. + """ + + if activation != "silu": + raise NotImplementedError( + f"MoonEP BF16 PoC runner only supports silu, got {activation!r}" + ) + + hidden_states = dispatch_output.hidden_states + route_weights_nvs = dispatch_output.route_weights_nvs + cu_seqlens = dispatch_output.cu_seqlens + expert_ids = dispatch_output.expert_ids + + if hidden_states.ndim != 2: + raise ValueError( + f"MoonEP hidden states must be [NvS, H], got {hidden_states.shape}" + ) + if cu_seqlens.ndim != 1: + raise ValueError(f"cu_seqlens must be 1D, got {cu_seqlens.shape}") + if expert_ids.shape != cu_seqlens.shape: + raise ValueError( + f"expert_ids shape {expert_ids.shape} must match cu_seqlens " + f"shape {cu_seqlens.shape}" + ) + if route_weights_nvs is not None and route_weights_nvs.ndim != 1: + raise ValueError( + f"route_weights_nvs must be 1D, got {route_weights_nvs.shape}" + ) + + output = torch.empty_like(hidden_states) + prev = 0 + for group_id in range(cu_seqlens.numel()): + cur = int(cu_seqlens[group_id].item()) + if cur < prev: + raise ValueError("MoonEP cu_seqlens must be non-decreasing") + if cur == prev: + continue + + expert_id = int(expert_ids[group_id].item()) + if expert_id < 0: + output[prev:cur].zero_() + prev = cur + continue + if expert_id >= weight_layout.full_gate_weight.shape[0]: + raise ValueError( + f"expert_id {expert_id} exceeds MoonEP weight rows " + f"{weight_layout.full_gate_weight.shape[0]}" + ) + + x = hidden_states[prev:cur] + gate = F.linear(x, weight_layout.full_gate_weight[expert_id]) + up = F.linear(x, weight_layout.full_up_weight[expert_id]) + activated = F.silu(gate) * up + y = F.linear(activated, weight_layout.full_down_weight[expert_id]) + if route_weights_nvs is not None: + y = y * route_weights_nvs[prev:cur].to(dtype=y.dtype).unsqueeze(-1) + output[prev:cur].copy_(y) + prev = cur + + if prev < hidden_states.shape[0]: + output[prev:].zero_() + + return MoonEPCombineInput( + hidden_states=output, + route_weights_nvs=route_weights_nvs, + plan=dispatch_output.plan, + num_tokens=dispatch_output.num_tokens, + ) + + +class MoonEPDispatcher(BaseDispatcher): + """MoonEP dispatcher for the initial BF16 inference PoC.""" + + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool = False, + num_experts: int | None = None, + num_local_experts: int | None = None, + hidden_size: int | None = None, + params_dtype: torch.dtype | None = None, + deepep_mode: DeepEPMode = DeepEPMode.AUTO, + async_finish: bool = False, + return_recv_hook: bool = False, + ): + super().__init__() + self.group = group + self.router_topk = router_topk + self.permute_fusion = permute_fusion + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.params_dtype = params_dtype + self.deepep_mode = deepep_mode + self.async_finish = async_finish + self.return_recv_hook = return_recv_hook + self.expert_mask_gpu = None + self.num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + self.num_prefetch_slots = None + + @staticmethod + def _raise_unimplemented() -> NoReturn: + raise NotImplementedError(_MOONEP_UNSUPPORTED_MESSAGE) + + def _get_buffer(self): + if self.hidden_size is None or self.num_experts is None: + raise ValueError( + "MoonEPDispatcher requires hidden_size and num_experts to " + "create a MoonEP buffer." + ) + return MoonEPBuffer.get_moonep_buffer( + group=self.group, + hidden_size=self.hidden_size, + router_topk=self.router_topk, + num_experts=self.num_experts, + num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, + num_prefetch_slots=self.num_prefetch_slots, + ) + + def _get_rank(self) -> int: + try: + return dist.get_rank(group=self.group) + except (AssertionError, RuntimeError, TypeError, ValueError): + return 0 + + def _pad_to_capacity( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + num_tokens = int(hidden_states.shape[0]) + capacity = int(self.num_max_dispatch_tokens_per_rank) + if num_tokens > capacity: + raise ValueError( + "MoonEP runtime batch has more tokens than its static buffer " + f"capacity: num_tokens={num_tokens}, capacity={capacity}. " + "Increase SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK." + ) + + hidden_states = hidden_states.contiguous() + topk_ids = topk_ids.to(dtype=torch.int32).contiguous() + topk_weights = topk_weights.to(dtype=torch.float32).contiguous() + if num_tokens == capacity: + return hidden_states, topk_ids, topk_weights, num_tokens + + pad_tokens = capacity - num_tokens + hidden_pad = hidden_states.new_zeros(pad_tokens, hidden_states.shape[1]) + id_pad = topk_ids.new_zeros(pad_tokens, topk_ids.shape[1]) + weight_pad = topk_weights.new_zeros(pad_tokens, topk_weights.shape[1]) + return ( + torch.cat((hidden_states, hidden_pad), dim=0).contiguous(), + torch.cat((topk_ids, id_pad), dim=0).contiguous(), + torch.cat((topk_weights, weight_pad), dim=0).contiguous(), + num_tokens, + ) + + def _tokens_per_expert(self, topk_ids: torch.Tensor) -> torch.Tensor: + assert self.num_experts is not None + return torch.bincount( + topk_ids.reshape(-1).to(dtype=torch.int64), + minlength=self.num_experts, + ).to(dtype=torch.int32) + + def _expert_ids_from_plan( + self, + cu_seqlens: torch.Tensor, + plan: Any, + ) -> torch.Tensor: + assert self.num_experts is not None + num_groups = int(cu_seqlens.numel()) + expert_ids = torch.full_like(cu_seqlens, -1) + experts_to_copy = plan.experts_to_copy + if experts_to_copy.ndim == 2: + experts_to_copy = experts_to_copy[self._get_rank()] + + prev = 0 + for group_id in range(num_groups): + cur = int(cu_seqlens[group_id].item()) + if cur > prev: + if group_id < self.num_experts: + expert_ids[group_id] = group_id + else: + expert_ids[group_id] = experts_to_copy[group_id - self.num_experts] + prev = cur + return expert_ids + + def dispatch( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ) -> DispatchOutput: + if not TopKOutputChecker.format_is_standard(topk_output): + raise NotImplementedError( + "MoonEP PoC requires standard top-k output before dispatch." + ) + if hidden_states.dtype != torch.bfloat16: + raise NotImplementedError( + "MoonEP PoC dispatcher supports BF16 hidden states only." + ) + if self.num_experts is None: + raise ValueError("MoonEPDispatcher requires num_experts.") + + hidden_states, topk_ids, topk_weights, num_tokens = self._pad_to_capacity( + hidden_states, + topk_output.topk_ids, + topk_output.topk_weights, + ) + tokens_per_expert = self._tokens_per_expert(topk_ids) + buffer = self._get_buffer() + hidden_nvsh, route_weights_nvs, cu_seqlens, plan = buffer.dispatch( + hidden_states, + topk_weights, + topk_ids, + tokens_per_expert, + async_finish=False, + ) + return MoonEPDispatchOutput( + hidden_states=hidden_nvsh, + route_weights_nvs=route_weights_nvs, + cu_seqlens=cu_seqlens, + plan=plan, + expert_ids=self._expert_ids_from_plan(cu_seqlens, plan), + num_tokens=num_tokens, + ) + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): + self._raise_unimplemented() + + def dispatch_b(self): + self._raise_unimplemented() + + def combine( + self, + combine_input: CombineInput, + ) -> torch.Tensor: + if combine_input.format != CombineInputFormat.MOONEP: + raise TypeError( + f"MoonEPDispatcher.combine expected MOONEP input, got " + f"{combine_input.format}" + ) + hidden_states, _route_weights_sk, _event = self._get_buffer().combine( + plan=combine_input.plan, + hidden_nvsh=combine_input.hidden_states, + route_weights_nvs=None, + async_finish=False, + ) + return hidden_states[: combine_input.num_tokens].contiguous() + + def combine_a( + self, + combine_input: CombineInput, + ): + self._raise_unimplemented() + + def combine_b(self): + self._raise_unimplemented() + + def prefetch_weight( + self, + plan: Any, + weight_layout: MoonEPExpertWeightLayout, + ) -> None: + self._get_buffer().prefetch_weight( + plan=plan, + async_finish=False, + full_gate_weight=weight_layout.full_gate_weight, + full_up_weight=weight_layout.full_up_weight, + full_down_weight=weight_layout.full_down_weight, + ) + + def register_deepep_dispatch_hook(self, hook): + self._raise_unimplemented() diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 593c65372204..ce2bdb9cd51e 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -30,6 +30,7 @@ class MoeA2ABackend(Enum): NONE = "none" DEEPEP = "deepep" + MOONEP = "moonep" MOONCAKE = "mooncake" NIXL = "nixl" MORI = "mori" @@ -55,6 +56,9 @@ def is_none(self): def is_deepep(self): return self == MoeA2ABackend.DEEPEP + def is_moonep(self): + return self == MoeA2ABackend.MOONEP + def is_mooncake(self): return self == MoeA2ABackend.MOONCAKE @@ -489,9 +493,15 @@ def is_sbo_enabled() -> bool: def is_deepep_class_backend() -> bool: - """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, Mori, or PPLX).""" + """Check if the MoE backend is DeepEP-family.""" b = get_moe_a2a_backend() - return b.is_deepep() or b.is_mooncake() or b.is_mori() or b.is_pplx() + return ( + b.is_deepep() + or b.is_moonep() + or b.is_mooncake() + or b.is_mori() + or b.is_pplx() + ) def uses_per_rank_fused_shared_slots() -> bool: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index afad543b1dba..6f3414c54d10 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -272,6 +272,7 @@ MOE_A2A_BACKEND_CHOICES = [ "none", "deepep", + "moonep", "mooncake", "nixl", "mori", @@ -2324,6 +2325,7 @@ class ServerArgs: Literal[ "none", "deepep", + "moonep", "mooncake", "nixl", "mori", @@ -6859,6 +6861,15 @@ def _handle_a2a_moe(self): self.cuda_graph_config.decode.backend = Backend.DISABLED self.cuda_graph_config.prefill.backend = Backend.DISABLED + if a2a_backend == "moonep": + logger.warning( + "MoonEP MoE is enabled in experimental BF16 PoC mode. " + "Cuda graph is disabled while the eager MoonEP dispatch/" + "prefetch/compute/combine path is validated." + ) + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED + if ( self.moe_a2a_backend == "none" and is_npu() ) or self.moe_a2a_backend == "ascend_tp": diff --git a/scripts/moonep/validate_moonep_bf16_poc.py b/scripts/moonep/validate_moonep_bf16_poc.py new file mode 100755 index 000000000000..35f054f851a1 --- /dev/null +++ b/scripts/moonep/validate_moonep_bf16_poc.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +"""Distributed BF16 MoonEP PoC validation. + +Run with torchrun on a single NVLink/NVSwitch node, for example: + + SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128 \ + torchrun --standalone --nproc-per-node=4 \ + scripts/moonep/validate_moonep_bf16_poc.py --tokens 128 --hidden-size 1024 + +The script validates SGLang's MoonEP BF16 path: +MoonEPDispatcher.dispatch -> MoonEPBuffer.prefetch_weight -> BF16 segment runner +-> MoonEPDispatcher.combine. It compares the final token-major output against a +rank-local PyTorch reference using the same top-k and expert weights. +""" + +from __future__ import annotations + +import argparse +from enum import IntEnum, auto +import json +import os +import sys +from pathlib import Path +from types import ModuleType +from typing import NamedTuple, Protocol, TypeGuard, runtime_checkable + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=128) + parser.add_argument("--hidden-size", type=int, default=1024) + parser.add_argument("--intermediate-size", type=int, default=128) + parser.add_argument("--top-k", type=int, default=2) + parser.add_argument("--experts-per-rank", type=int, default=2) + parser.add_argument("--prefetch-slots", type=int, default=-1) + parser.add_argument("--seed", type=int, default=20260802) + parser.add_argument("--atol", type=float, default=5e-2) + parser.add_argument("--rtol", type=float, default=5e-2) + return parser.parse_args() + + +def setup_dist() -> tuple[int, int, int]: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + return rank, world_size, local_rank + + +def make_topk(tokens: int, top_k: int, num_experts: int, device: torch.device): + # Deterministic but rank-local routing. Keep weights normalized so the + # reference magnitude remains bounded. + topk_ids = torch.randint( + 0, + num_experts, + (tokens, top_k), + device=device, + dtype=torch.int64, + ) + raw_weights = torch.rand(tokens, top_k, device=device, dtype=torch.float32) + topk_weights = raw_weights / raw_weights.sum(dim=-1, keepdim=True) + return topk_ids, topk_weights + + +class LightweightTopKOutputFormat(IntEnum): + STANDARD = auto() + + +@runtime_checkable +class LightweightTopKOutput(Protocol): + @property + def format(self) -> LightweightTopKOutputFormat: ... + + +class LightweightStandardTopKOutput(NamedTuple): + topk_weights: torch.Tensor + topk_ids: torch.Tensor + router_logits: torch.Tensor + + @property + def format(self) -> LightweightTopKOutputFormat: + return LightweightTopKOutputFormat.STANDARD + + +class LightweightTopKOutputChecker: + @staticmethod + def format_is_standard( + topk_output: LightweightTopKOutput, + ) -> TypeGuard[LightweightStandardTopKOutput]: + return isinstance(topk_output, LightweightStandardTopKOutput) + + +class LightweightDeepEPMode(IntEnum): + NORMAL = 1 + LOW_LATENCY = 2 + AUTO = 3 + + +def install_lightweight_sglang_imports() -> None: + """Install validation-only SGLang package stubs. + + Full SGLang environments should use the normal imports. Minimal remote + validation images often do not carry every frontend/model dependency that + SGLang's package ``__init__`` imports, though this script only needs the + MoonEP dispatcher, its base protocol, envs, and runtime buffer registry. + These stubs keep the package search path pointed at the checked-out source + tree while bypassing heavyweight parent ``__init__`` files. + """ + + repo_root = Path(__file__).resolve().parents[2] + sglang_root = repo_root / "python" / "sglang" + + for name in list(sys.modules): + if name == "sglang" or name.startswith("sglang."): + del sys.modules[name] + + def add_package(name: str, path: Path) -> None: + module = ModuleType(name) + module.__path__ = [str(path)] + sys.modules[name] = module + + add_package("sglang", sglang_root) + add_package("sglang.srt", sglang_root / "srt") + add_package("sglang.srt.layers", sglang_root / "srt" / "layers") + add_package("sglang.srt.layers.moe", sglang_root / "srt" / "layers" / "moe") + add_package( + "sglang.srt.layers.moe.token_dispatcher", + sglang_root / "srt" / "layers" / "moe" / "token_dispatcher", + ) + + topk_module = ModuleType("sglang.srt.layers.moe.topk") + topk_module.TopKOutputFormat = LightweightTopKOutputFormat + topk_module.TopKOutput = LightweightTopKOutput + topk_module.StandardTopKOutput = LightweightStandardTopKOutput + topk_module.TopKOutputChecker = LightweightTopKOutputChecker + sys.modules[topk_module.__name__] = topk_module + + utils_module = ModuleType("sglang.srt.layers.moe.utils") + utils_module.DeepEPMode = LightweightDeepEPMode + sys.modules[utils_module.__name__] = utils_module + + +def import_moonep_validation_symbols(): + try: + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + + return ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) + except ModuleNotFoundError as exc: + print( + "Normal SGLang import failed; retrying with lightweight validation " + f"imports: {exc}", + file=sys.stderr, + flush=True, + ) + install_lightweight_sglang_imports() + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + + return ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) + + +def expert_mlp(x, expert_id: int, gate, up, down): + return F.linear( + F.silu(F.linear(x, gate[expert_id])) * F.linear(x, up[expert_id]), + down[expert_id], + ) + + +def reference_output(hidden, topk_ids, topk_weights, gate, up, down): + out = torch.zeros_like(hidden) + tokens, top_k = topk_ids.shape + for token_idx in range(tokens): + x = hidden[token_idx : token_idx + 1] + acc = torch.zeros_like(x) + for k in range(top_k): + expert_id = int(topk_ids[token_idx, k].item()) + acc += expert_mlp(x, expert_id, gate, up, down) * topk_weights[ + token_idx, k + ].to(hidden.dtype) + out[token_idx] = acc[0] + return out + + +def main() -> None: + args = parse_args() + os.environ["SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = str(args.tokens) + if args.prefetch_slots > 0: + os.environ["SGLANG_MOONEP_NUM_PREFETCH_SLOTS"] = str(args.prefetch_slots) + + rank, world_size, local_rank = setup_dist() + device = torch.device(f"cuda:{local_rank}") + torch.manual_seed(args.seed + rank) + + ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) = import_moonep_validation_symbols() + + num_experts = world_size * args.experts_per_rank + hidden = torch.randn( + args.tokens, + args.hidden_size, + device=device, + dtype=torch.bfloat16, + ) + topk_ids, topk_weights = make_topk(args.tokens, args.top_k, num_experts, device) + topk_output = StandardTopKOutput( + topk_weights=topk_weights, + topk_ids=topk_ids, + router_logits=torch.empty(0, device=device), + ) + + dispatcher = MoonEPDispatcher( + group=dist.group.WORLD, + router_topk=args.top_k, + num_experts=num_experts, + num_local_experts=args.experts_per_rank, + hidden_size=args.hidden_size, + params_dtype=torch.bfloat16, + ) + + dispatch_output = dispatcher.dispatch(hidden, topk_output) + num_prefetch_slots = int(dispatch_output.expert_ids.numel()) - num_experts + + # Full global rows are deliberately replicated for this correctness PoC. + # The production path should replace this with true symmetric expert-row + # mappings whose physical storage is owned by each expert's home rank. + torch.manual_seed(args.seed) + gate = torch.randn( + num_experts + num_prefetch_slots, + args.intermediate_size, + args.hidden_size, + device=device, + dtype=torch.bfloat16, + ) / 8 + up = torch.randn_like(gate) / 8 + down = torch.randn( + num_experts + num_prefetch_slots, + args.hidden_size, + args.intermediate_size, + device=device, + dtype=torch.bfloat16, + ) / 8 + gate[num_experts:].zero_() + up[num_experts:].zero_() + down[num_experts:].zero_() + layout = MoonEPExpertWeightLayout( + gate.contiguous(), + up.contiguous(), + down.contiguous(), + num_prefetch_slots, + ) + + dispatcher.prefetch_weight(dispatch_output.plan, layout) + combine_input = run_moonep_bf16_expert(dispatch_output, layout) + output = dispatcher.combine(combine_input) + + expected = reference_output(hidden, topk_ids, topk_weights, gate, up, down) + max_abs_err = (output.float() - expected.float()).abs().max() + rel_err = max_abs_err / expected.float().abs().max().clamp_min(1e-6) + local_ok = bool( + torch.allclose(output.float(), expected.float(), atol=args.atol, rtol=args.rtol) + ) + ok_tensor = torch.tensor([1 if local_ok else 0], device=device, dtype=torch.int32) + dist.all_reduce(ok_tensor, op=dist.ReduceOp.MIN) + + result = { + "rank": rank, + "world_size": world_size, + "tokens": args.tokens, + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "top_k": args.top_k, + "num_experts": num_experts, + "num_prefetch_slots": num_prefetch_slots, + "max_abs_err": float(max_abs_err.item()), + "relative_err": float(rel_err.item()), + "local_ok": local_ok, + "global_ok": bool(ok_tensor.item()), + } + print(json.dumps(result, sort_keys=True), flush=True) + dist.barrier(device_ids=[local_rank]) + MoonEPBuffer.destroy_all_buffers() + dist.destroy_process_group() + if rank == 0 and not bool(ok_tensor.item()): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/test/registered/unit/layers/moe/test_moonep_buffer.py b/test/registered/unit/layers/moe/test_moonep_buffer.py new file mode 100644 index 000000000000..dd06f5f9fbc9 --- /dev/null +++ b/test/registered/unit/layers/moe/test_moonep_buffer.py @@ -0,0 +1,306 @@ +import sys +import types +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatchOutput, + MoonEPExpertWeightLayout, + get_moonep_expert_weight_layout, + run_moonep_bf16_expert, +) +from sglang.srt.runtime_context import reset_context +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _FakeMoonEPBuffer: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.destroy_calls = 0 + self.__class__.instances.append(self) + + def destroy(self): + self.destroy_calls += 1 + + +def _fake_moonep_module(): + module = types.ModuleType("moonep") + module.Buffer = _FakeMoonEPBuffer + return module + + +class TestMoonEPBuffer(unittest.TestCase): + def setUp(self): + reset_context() + _FakeMoonEPBuffer.instances.clear() + + def tearDown(self): + try: + MoonEPBuffer.destroy_all_buffers() + finally: + reset_context() + _FakeMoonEPBuffer.instances.clear() + + def test_lazily_constructs_and_reuses_buffer_for_static_key(self): + group = object() + + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=4, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + same_buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + larger_buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=512, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + + self.assertIs(buffer, same_buffer) + self.assertIsNot(buffer, larger_buffer) + self.assertEqual(len(_FakeMoonEPBuffer.instances), 2) + self.assertEqual( + buffer.kwargs, + { + "S": 256, + "H": 1024, + "K": 8, + "E": 64, + "num_ep_ranks": 4, + "num_sms": 20, + "token_padding": 64, + "B": 16, + "group": group, + }, + ) + self.assertIs(MoonEPBuffer.get_existing_buffer(), larger_buffer) + + def test_resolves_env_defaults_and_training_safe_prefetch_slots(self): + group = object() + + with ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(384), + envs.SGLANG_MOONEP_NUM_PREFETCH_SLOTS.override(-1), + envs.SGLANG_MOONEP_TOKEN_PADDING.override(32), + envs.SGLANG_MOONEP_NUM_SMS.override(18), + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=8, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=2048, + router_topk=6, + num_experts=128, + ) + + self.assertEqual(buffer.kwargs["S"], 384) + self.assertEqual(buffer.kwargs["token_padding"], 32) + self.assertEqual(buffer.kwargs["num_sms"], 18) + self.assertEqual(buffer.kwargs["B"], 16) + + def test_rejects_non_divisible_experts_before_allocating(self): + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=6, + ), + ): + with self.assertRaisesRegex(ValueError, "divisible"): + MoonEPBuffer.get_moonep_buffer( + group=object(), + hidden_size=1024, + router_topk=8, + num_experts=64, + ) + + self.assertEqual(_FakeMoonEPBuffer.instances, []) + + def test_destroy_all_buffers_releases_cached_buffers(self): + group = object() + + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=4, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + ) + + MoonEPBuffer.destroy_all_buffers() + + self.assertEqual(buffer.destroy_calls, 1) + self.assertIsNone(MoonEPBuffer.get_existing_buffer()) + + +class TestMoonEPExpertWeightLayout(unittest.TestCase): + def _fake_layer(self): + num_experts, hidden_size, intermediate_size = 3, 4, 5 + w13_weight = torch.arange( + num_experts * 2 * intermediate_size * hidden_size, + dtype=torch.bfloat16, + ).reshape(num_experts, 2 * intermediate_size, hidden_size) + w2_weight = torch.arange( + num_experts * hidden_size * intermediate_size, + dtype=torch.bfloat16, + ).reshape(num_experts, hidden_size, intermediate_size) + + return SimpleNamespace( + quant_config=None, + moe_runner_config=SimpleNamespace( + num_fused_shared_experts=0, + is_gated=True, + ), + use_triton_kernels=False, + w13_weight=w13_weight, + w2_weight=w2_weight, + num_experts=num_experts, + intermediate_size_per_partition=intermediate_size, + hidden_size=hidden_size, + ) + + def test_layout_splits_gate_up_down_and_adds_prefetch_slots(self): + layer = self._fake_layer() + + layout = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + self.assertEqual(tuple(layout.full_gate_weight.shape), (5, 5, 4)) + self.assertEqual(tuple(layout.full_up_weight.shape), (5, 5, 4)) + self.assertEqual(tuple(layout.full_down_weight.shape), (5, 4, 5)) + torch.testing.assert_close( + layout.full_gate_weight[:3], + layer.w13_weight[:, :5, :], + ) + torch.testing.assert_close( + layout.full_up_weight[:3], + layer.w13_weight[:, 5:10, :], + ) + torch.testing.assert_close(layout.full_down_weight[:3], layer.w2_weight) + self.assertTrue(torch.all(layout.full_gate_weight[3:] == 0)) + self.assertTrue(torch.all(layout.full_up_weight[3:] == 0)) + self.assertTrue(torch.all(layout.full_down_weight[3:] == 0)) + + def test_layout_is_cached_for_same_weight_storage(self): + layer = self._fake_layer() + + first = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + second = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + self.assertIs(first, second) + + def test_layout_rejects_local_expert_storage(self): + layer = self._fake_layer() + layer.w13_weight = layer.w13_weight[:2].contiguous() + + with self.assertRaisesRegex(ValueError, "global w13_weight"): + get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + +class TestMoonEPBf16ExpertRunner(unittest.TestCase): + def test_segment_runner_applies_expert_weights_and_route_weights(self): + hidden_states = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + dtype=torch.float32, + ) + route_weights = torch.tensor([1.0, 0.5, 2.0], dtype=torch.float32) + cu_seqlens = torch.tensor([2, 3], dtype=torch.int32) + expert_ids = torch.tensor([0, 1], dtype=torch.int32) + gate = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.5, 0.0], [0.0, 0.5]], + ] + ) + up = torch.tensor( + [ + [[2.0, 0.0], [0.0, 2.0]], + [[1.5, 0.0], [0.0, 1.5]], + ] + ) + down = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[2.0, 0.0], [0.0, 2.0]], + ] + ) + layout = MoonEPExpertWeightLayout( + full_gate_weight=gate, + full_up_weight=up, + full_down_weight=down, + num_prefetch_slots=0, + ) + dispatch_output = MoonEPDispatchOutput( + hidden_states=hidden_states, + route_weights_nvs=route_weights, + cu_seqlens=cu_seqlens, + plan=object(), + expert_ids=expert_ids, + num_tokens=2, + ) + + combine_input = run_moonep_bf16_expert(dispatch_output, layout) + + expected = torch.empty_like(hidden_states) + for start, end, expert in [(0, 2, 0), (2, 3, 1)]: + x = hidden_states[start:end] + y = torch.nn.functional.linear( + torch.nn.functional.silu( + torch.nn.functional.linear(x, gate[expert]) + ) + * torch.nn.functional.linear(x, up[expert]), + down[expert], + ) + expected[start:end] = y * route_weights[start:end, None] + + torch.testing.assert_close(combine_input.hidden_states, expected) + self.assertIs(combine_input.plan, dispatch_output.plan) + self.assertEqual(combine_input.num_tokens, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 2acef3594a03..aed81383f025 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1500,6 +1500,23 @@ def test_waterfill_supports_deepep_low_latency_mode(self): self.assertTrue(server_args.enforce_shared_experts_fusion) +class TestMoonEPArgs(CustomTestCase): + def test_moonep_backend_is_recognized_and_disables_cuda_graph(self): + server_args = ServerArgs( + model_path="dummy", + moe_a2a_backend="moonep", + ) + + server_args._handle_a2a_moe() + + from sglang.srt.arg_groups.overrides import resolved_view + + self.assertEqual(resolved_view(server_args).moe_a2a_backend, "moonep") + self.assertEqual(resolved_view(server_args).ep_size, server_args.tp_size) + self.assertEqual(server_args.cuda_graph_config.decode.backend, Backend.DISABLED) + self.assertEqual(server_args.cuda_graph_config.prefill.backend, Backend.DISABLED) + + class TestPrefillOnlyDisableKvCache(unittest.TestCase): """Validation for --prefill-only-disable-kv-cache. diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 3d519b5ccd1c..de445840caed 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -440,6 +440,39 @@ def test_initialize_materializes_group(self): self.assertTrue(is_tbo_enabled()) self.assertEqual(get_flags().moe.deepep_config, "") + def test_initialize_materializes_moonep_backend(self): + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + + self._init(moe_a2a_backend="moonep") + self.assertTrue(get_moe_a2a_backend().is_moonep()) + + def test_moonep_dispatch_contract_formats(self): + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInputChecker, + DispatchOutputChecker, + MoonEPCombineInput, + MoonEPDispatchOutput, + ) + + dispatch_output = MoonEPDispatchOutput( + hidden_states=object(), + route_weights_nvs=None, + cu_seqlens=object(), + plan=object(), + expert_ids=object(), + num_tokens=1, + ) + combine_input = MoonEPCombineInput( + hidden_states=object(), + route_weights_nvs=None, + plan=dispatch_output.plan, + num_tokens=1, + ) + + self.assertTrue(dispatch_output.format.is_moonep()) + self.assertTrue(DispatchOutputChecker.format_is_moonep(dispatch_output)) + self.assertTrue(CombineInputChecker.format_is_moonep(combine_input)) + def test_speculative_swap_and_restore(self): from sglang.srt.layers.moe.utils import ( get_moe_a2a_backend,