[MoonEP] Harden BF16 reference adapter contracts and validation

python/sglang/srt/arg_groups/overrides.py

@@ -2503,6 +2503,7 @@ def _cutlass_moe_env_override(view: Any) -> dict:
     {
         "megamoe",
         "deepep",
+        "moonep",
         "mooncake",
         "nixl",
         "ascend_fuseep",

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -28,6 +28,7 @@
 from sglang.srt.layers.moe.token_dispatcher import (
     DeepEPDispatcher,
     MooncakeEPDispatcher,
+    MoonEPDispatcher,
     MoriEPDispatcher,
     NixlEPDispatcher,
     PplxDispatcher,
@@ -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:

python/sglang/srt/distributed/parallel_state.py

@@ -51,6 +51,7 @@
 )
 from sglang.srt.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND
 from sglang.srt.runtime_context import (
+    cleanup_distributed_resources,
     get_global_dwdp_manager,
     set_global_dwdp_manager,
 )
@@ -2856,6 +2857,7 @@ def get_moe_tensor_parallel_rank():
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
+    cleanup_distributed_resources()
     dwdp_mgr = get_global_dwdp_manager()
     if dwdp_mgr is not None:
         dwdp_mgr.cleanup()
@@ -2909,6 +2911,7 @@ def destroy_model_parallel():
 
 
 def destroy_distributed_environment():
+    cleanup_distributed_resources()
     global _WORLD, _MODEL_PARALLEL_GROUP_TIMEOUT
     if _WORLD:
         _WORLD.destroy()

python/sglang/srt/environ.py

@@ -807,6 +807,13 @@ class Envs:
     SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False)
     SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
 
+    # MoonEP
+    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)
+
     # PPLX-EP (Perplexity pplx-kernels NVSHMEM all-to-all)
     SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
 

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,
@@ -60,6 +64,16 @@
 logger = logging.getLogger(__name__)
 
 
+def _should_use_deepep_bf16_dispatch_fallback() -> bool:
+    """Return whether the deprecated DeepEP BF16 dispatch path applies."""
+
+    return (
+        deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+        and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
+        and get_moe_a2a_backend().is_deepep()
+    )
+
+
 class DeepEPMoE(FusedMoE):
     """
     MoE Expert Parallel Impl based on DeepEP (https://github.com/deepseek-ai/DeepEP/tree/main)
@@ -113,10 +127,7 @@ def __init__(
             quant_config, Fp8Config
         ):
             self.deprecate_flag = True
-        elif (
-            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
-            and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
-        ):
+        elif _should_use_deepep_bf16_dispatch_fallback():
             self.deprecate_flag = True
         elif (
             get_moe_runner_backend().is_flashinfer_cutedsl()
@@ -299,6 +310,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.cu_seqlens.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 +385,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()

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -39,6 +39,9 @@
 )
 from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher
 from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
+from sglang.srt.layers.moe.token_dispatcher.moonep import (
+    validate_moonep_reference_bf16_config,
+)
 from sglang.srt.layers.moe.token_dispatcher.standard import (
     StandardDispatcher,
 )
@@ -149,6 +152,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 +400,25 @@ 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:
+            validate_moonep_reference_bf16_config(
+                quant_config=quant_config,
+                params_dtype=params_dtype,
+                num_fused_shared_experts=num_fused_shared_experts,
+                with_bias=with_bias,
+                activation=activation,
+                quant_method=self.quant_method,
+                layer=self,
+            )
+
         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 +430,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)

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",

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"