sglang 0.5.3rc2__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/layers/moe/token_dispatcher/mooncake.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import NamedTuple, Optional, Tuple
+
+from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
+from sglang.srt.layers.dp_attention import get_is_extend_in_batch
+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.utils import DeepEPMode
+from sglang.srt.utils import get_int_env_var
+
+try:
+    from mooncake.mooncake_ep_buffer import Buffer
+
+    use_mooncake_ep = True
+except ImportError:
+    use_mooncake_ep = False
+
+from enum import Enum, auto
+
+import torch
+import torch.distributed as dist
+
+logger = logging.getLogger(__name__)
+
+
+class MooncakeDispatchOutput(NamedTuple):
+    """Mooncake EP dispatch output."""
+
+    hidden_states: torch.Tensor
+    hidden_states_scale: Optional[torch.Tensor]
+    topk_ids: torch.Tensor
+    topk_weights: torch.Tensor
+    masked_m: torch.Tensor
+    expected_m: int
+
+    @property
+    def format(self) -> DispatchOutputFormat:
+        return DispatchOutputFormat.DEEPEP_LL
+
+
+assert isinstance(MooncakeDispatchOutput, DispatchOutput)
+
+
+class MooncakeCombineInput(NamedTuple):
+    """Mooncake EP combine input."""
+
+    pass
+
+    @property
+    def format(self) -> CombineInputFormat:
+        return CombineInputFormat.DEEPEP_LL
+
+
+assert isinstance(MooncakeCombineInput, CombineInput)
+
+
+class EPBuffer:
+    _buffer = None
+    _hidden_size: Optional[int] = None
+    _num_max_dispatch_tokens_per_rank: Optional[int] = None
+    _num_experts: Optional[int] = None
+
+    @classmethod
+    def get_ep_buffer(
+        cls,
+        group: dist.ProcessGroup,
+        hidden_size: int,
+        param_bytes: int,
+        deepep_mode: DeepEPMode,
+        num_max_dispatch_tokens_per_rank: int = -1,
+        num_experts: int = -1,
+    ):
+        if cls._buffer is not None:
+            return cls._buffer
+
+        cls._hidden_size = hidden_size
+        cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
+        cls._num_experts = num_experts
+
+        num_ep_buffer_bytes = 0
+        if deepep_mode.enable_normal():
+            raise NotImplementedError(
+                "Normal mode is not supported for Mooncake EP yet."
+            )
+        if deepep_mode.enable_low_latency():
+            assert num_max_dispatch_tokens_per_rank != -1
+            assert num_experts != -1 and num_experts % group.size() == 0
+            num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint(
+                num_max_dispatch_tokens_per_rank,
+                hidden_size,
+                group.size(),
+                num_experts,
+            )
+
+        cls._buffer = Buffer(group, num_ep_buffer_bytes)
+        return cls._buffer
+
+
+class _MooncakeEPDispatcherImpl:
+    def __init__(
+        self,
+        group: torch.distributed.ProcessGroup,
+        router_topk: int,
+        permute_fusion: bool,
+        num_experts: int,
+        num_local_experts: int,
+        hidden_size: int,
+        params_dtype: torch.dtype,
+        return_recv_hook: bool,
+        deepep_mode: DeepEPMode,
+    ):
+        if not use_mooncake_ep:
+            raise ImportError(
+                "Mooncake EP is not installed. Please install Mooncake package at "
+                "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "
+                "with EP support to run SGLang with Mooncake EP."
+            )
+        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.return_recv_hook = return_recv_hook
+        self.deepep_mode = deepep_mode
+
+        self.params_bytes = 2
+        self.num_max_dispatch_tokens_per_rank = get_int_env_var(
+            "SGLANG_MOONCAKE_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 128
+        )
+        # Mooncake EP dispatch uses FINISHED_SUM_TAG=1024
+        # and the logic requires num-tokens-sent-from-one-rank-to-another-rank less than it
+        assert self.num_max_dispatch_tokens_per_rank <= 1024
+
+        self.first_execution = True
+        self.timeout_us = 10000000
+
+        self.active_ranks = ElasticEPStateManager.instance().active_ranks
+
+        self.handle = None
+
+    def dispatch_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_output: TopKOutput,
+    ):
+        topk_ids, topk_weights = topk_output.topk_ids, topk_output.topk_weights
+        buffer = self._get_buffer()
+        topk_ids = topk_ids.to(torch.int64)
+        expected_m = (
+            hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1]
+            + self.num_experts
+        ) // self.num_experts
+        hidden_states, masked_m, event, hook = self._dispatch_core(
+            hidden_states,
+            topk_ids,
+            use_fp8=True,
+        )
+        return (
+            hidden_states,
+            topk_ids,
+            topk_weights,
+            masked_m,
+            expected_m,
+            event,
+            hook,
+        )
+
+    def dispatch_b(
+        self,
+        hidden_states,
+        topk_ids,
+        topk_weights,
+        masked_m,
+        expected_m,
+        event,
+        hook,
+    ):
+        hook() if self.return_recv_hook else event.current_stream_wait()
+
+        get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency(
+            masked_m
+        )
+
+        if isinstance(hidden_states, tuple):
+            hidden_states, hidden_states_scale = hidden_states
+        else:
+            hidden_states_scale = None
+
+        return MooncakeDispatchOutput(
+            hidden_states,
+            hidden_states_scale,
+            topk_ids,
+            topk_weights,
+            masked_m,
+            expected_m,
+        )
+
+    def _dispatch_core(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        use_fp8: bool = False,
+    ):
+        buffer = self._get_buffer()
+        packed_recv_hidden, packed_recv_count, self.handle, event, hook = (
+            buffer.dispatch(
+                hidden_states,
+                topk_ids,
+                self.active_ranks,
+                self.num_max_dispatch_tokens_per_rank,
+                self.num_experts,
+                -1 if self.first_execution else self.timeout_us,
+                use_fp8=use_fp8,
+                async_finish=not self.return_recv_hook,
+                return_recv_hook=self.return_recv_hook,
+            )
+        )
+        return packed_recv_hidden, packed_recv_count, event, hook
+
+    def combine_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ):
+        hidden_states, event, hook = self._combine_core(
+            hidden_states,
+            topk_ids,
+            topk_weights,
+        )
+        return hidden_states, event, hook
+
+    def combine_b(self, hidden_states, event, hook):
+        hook() if self.return_recv_hook else event.current_stream_wait()
+        return hidden_states
+
+    def _combine_core(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ):
+        buffer = self._get_buffer()
+        combined_hidden_states, event, hook = buffer.combine(
+            hidden_states,
+            topk_ids,
+            topk_weights,
+            self.active_ranks,
+            -1 if self.first_execution else self.timeout_us,
+            self.handle,
+            async_finish=not self.return_recv_hook,
+            return_recv_hook=self.return_recv_hook,
+        )
+        self.first_execution = False
+        self.handle = None
+        return combined_hidden_states, event, hook
+
+    def _get_buffer(self):
+        return EPBuffer.get_ep_buffer(
+            self.group,
+            self.hidden_size,
+            self.params_bytes,
+            self.deepep_mode,
+            self.num_max_dispatch_tokens_per_rank,
+            self.num_experts,
+        )
+
+
+@dataclass
+class _Stage(Enum):
+    INITIAL = auto()
+    AFTER_DISPATCH_A = auto()
+    AFTER_DISPATCH_B = auto()
+    AFTER_COMBINE_A = auto()
+
+
+class MooncakeEPDispatcher(BaseDispatcher):
+    def __init__(
+        self,
+        group: torch.distributed.ProcessGroup,
+        router_topk: int,
+        permute_fusion: bool = False,
+        num_experts: int = None,
+        num_local_experts: int = None,
+        hidden_size: int = None,
+        params_dtype: torch.dtype = None,
+        deepep_mode: DeepEPMode = DeepEPMode.AUTO,
+        async_finish: bool = False,
+        return_recv_hook: bool = False,
+    ):
+        self.deepep_mode = deepep_mode
+
+        if self.deepep_mode.enable_low_latency():
+            self._low_latency_dispatcher = _MooncakeEPDispatcherImpl(
+                group=group,
+                router_topk=router_topk,
+                permute_fusion=permute_fusion,
+                num_experts=num_experts,
+                num_local_experts=num_local_experts,
+                hidden_size=hidden_size,
+                params_dtype=params_dtype,
+                return_recv_hook=return_recv_hook,
+                deepep_mode=deepep_mode,
+            )
+        if self.deepep_mode.enable_normal():
+            raise NotImplementedError
+
+        self._stage = _Stage.INITIAL
+
+    def dispatch(self, *args, **kwargs) -> DispatchOutput:
+        self.dispatch_a(*args, **kwargs)
+        ret = self.dispatch_b()
+        return ret
+
+    def dispatch_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_output: TopKOutput,
+    ):
+        self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A)
+        inner_state = self._get_impl().dispatch_a(
+            hidden_states=hidden_states,
+            topk_output=topk_output,
+        )
+        self._dispatch_intermediate_state = inner_state
+
+    def dispatch_b(self):
+        self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B)
+        inner_state = self._dispatch_intermediate_state
+        del self._dispatch_intermediate_state
+        return self._get_impl().dispatch_b(*inner_state)
+
+    def combine(self, *args, **kwargs) -> Tuple:
+        self.combine_a(*args, **kwargs)
+        ret = self.combine_b()
+        return ret
+
+    def combine_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        overlap_args: Optional = None,
+    ):
+        self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A)
+        inner_state = self._get_impl().combine_a(
+            hidden_states=hidden_states,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+        )
+        self._combine_intermediate_state = inner_state
+
+    def combine_b(self):
+        self._update_stage(_Stage.AFTER_COMBINE_A, _Stage.INITIAL)
+        inner_state = self._combine_intermediate_state
+        del self._combine_intermediate_state
+        return self._get_impl().combine_b(*inner_state)
+
+    def _get_impl(self) -> _MooncakeEPDispatcherImpl:
+        is_extend_in_batch = get_is_extend_in_batch()
+        resolved_deepep_mode = self.deepep_mode.resolve(is_extend_in_batch)
+        if resolved_deepep_mode == DeepEPMode.NORMAL:
+            raise NotImplementedError
+        elif resolved_deepep_mode == DeepEPMode.LOW_LATENCY:
+            return self._low_latency_dispatcher
+        else:
+            raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}")
+
+    def _update_stage(self, old_stage, new_stage):
+        assert self._stage == old_stage
+        self._stage = new_stage
+
+    def set_quant_config(self, quant_config: dict):
+        pass

sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -4,6 +4,11 @@ from typing import TYPE_CHECKING, NamedTuple
 
 import torch
 
+from sglang.srt.distributed import (
+    get_moe_expert_parallel_rank,
+    get_moe_expert_parallel_world_size,
+)
+from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
 from sglang.srt.layers.moe.token_dispatcher.base import (
     BaseDispatcher,
     CombineInput,
@@ -11,6 +16,8 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
     DispatchOutput,
     DispatchOutputFormat,
 )
+from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
+from sglang.srt.layers.moe.utils import get_moe_runner_backend
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.topk import TopKOutput
@@ -45,9 +52,45 @@ assert isinstance(StandardCombineInput, CombineInput)
 
 class StandardDispatcher(BaseDispatcher):
 
+    def __init__(self, moe_runner_config: MoeRunnerConfig):
+        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.enable_flashinfer_cutlass_moe = (
+            get_moe_runner_backend().is_flashinfer_cutlass()
+        )
+        self.num_experts = moe_runner_config.num_experts
+        self.num_local_experts = moe_runner_config.num_local_experts
+        self.moe_ep_rank = get_moe_expert_parallel_rank()
+        self.local_expert_mapping = None
+
     def dispatch(
         self, hidden_states: torch.Tensor, topk_output: TopKOutput
     ) -> DispatchOutput:
+
+        if (
+            self.moe_ep_size > 1
+            and not self.enable_flashinfer_cutlass_moe
+            and TopKOutputChecker.format_is_standard(topk_output)
+        ):
+            if self.local_expert_mapping is None:
+                self.local_expert_mapping = torch.full(
+                    (self.num_experts,), -1, dtype=torch.int32, device="cuda"
+                )
+                self.local_expert_mapping[
+                    self.moe_ep_rank
+                    * self.num_local_experts : (self.moe_ep_rank + 1)
+                    * self.num_local_experts
+                ] = torch.arange(
+                    0, self.num_local_experts, dtype=torch.int32, device="cuda"
+                )
+
+        if self.local_expert_mapping is not None:
+            if TopKOutputChecker.format_is_standard(topk_output):
+                topk_output = topk_output._replace(
+                    topk_ids=self.local_expert_mapping[topk_output.topk_ids]
+                )
+            elif TopKOutputChecker.format_is_triton_kernel(topk_output):
+                raise NotImplementedError()
+
         return StandardDispatchOutput(
             hidden_states=hidden_states, topk_output=topk_output
         )
@@ -59,3 +102,6 @@ class StandardDispatcher(BaseDispatcher):
             # TODO: this branch should be removed in the future
             assert isinstance(combine_input, torch.Tensor)
             return combine_input
+
+    def set_quant_config(self, quant_config: dict):
+        pass

sglang/srt/layers/moe/topk.py

@@ -365,9 +365,10 @@ class TopK(CustomOp):
     def empty_topk_output(self, device: torch.device) -> TopKOutput:
         topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
         topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
-        topk_idx = torch.full((0, topk), -1, dtype=torch.int32, device=device)
+        topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device)
+        # FIXME: router_logits should be of size (0, num_experts)
         router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
-        return StandardTopKOutput(topk_weights, topk_idx, router_logits)
+        return StandardTopKOutput(topk_weights, topk_ids, router_logits)
 
 
 # ------------------------------- TopK implementation -------------------------------------

sglang/srt/layers/moe/utils.py

@@ -13,6 +13,7 @@ from sglang.srt.layers.dp_attention import (
     get_attention_dp_size,
     is_dp_attention_enabled,
 )
+from sglang.srt.utils import log_info_on_rank0
 
 if TYPE_CHECKING:
     from sglang.srt.server_args import ServerArgs
@@ -24,6 +25,7 @@ class MoeA2ABackend(Enum):
 
     NONE = "none"
     DEEPEP = "deepep"
+    MOONCAKE = "mooncake"
 
     @classmethod
     def _missing_(cls, value):
@@ -40,20 +42,28 @@ class MoeA2ABackend(Enum):
     def is_deepep(self):
         return self == MoeA2ABackend.DEEPEP
 
+    def is_mooncake(self):
+        return self == MoeA2ABackend.MOONCAKE
+
 
 class MoeRunnerBackend(Enum):
 
     AUTO = "auto"
+    DEEP_GEMM = "deep_gemm"
     TRITON = "triton"
     TRITON_KERNEL = "triton_kernel"
     FLASHINFER_TRTLLM = "flashinfer_trtllm"
     FLASHINFER_CUTLASS = "flashinfer_cutlass"
     FLASHINFER_MXFP4 = "flashinfer_mxfp4"
     FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
+    CUTLASS = "cutlass"
 
     def is_auto(self):
         return self == MoeRunnerBackend.AUTO
 
+    def is_deep_gemm(self):
+        return self == MoeRunnerBackend.DEEP_GEMM
+
     def is_triton(self):
         return self == MoeRunnerBackend.TRITON
 
@@ -72,6 +82,9 @@ class MoeRunnerBackend(Enum):
     def is_flashinfer_mxfp4(self):
         return self == MoeRunnerBackend.FLASHINFER_MXFP4
 
+    def is_cutlass(self):
+        return self == MoeRunnerBackend.CUTLASS
+
 
 class DeepEPMode(Enum):
 
@@ -147,7 +160,10 @@ def get_moe_a2a_backend() -> MoeA2ABackend:
 def get_moe_runner_backend() -> MoeRunnerBackend:
     global MOE_RUNNER_BACKEND
     if MOE_RUNNER_BACKEND is None:
-        logger.warning("MOE_RUNNER_BACKEND is not initialized, using triton backend")
+        log_info_on_rank0(
+            logger,
+            "MOE_RUNNER_BACKEND is not initialized, the backend will be automatically selected",
+        )
         MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO
     return MOE_RUNNER_BACKEND
 

sglang/srt/layers/quantization/__init__.py

@@ -10,10 +10,6 @@ import torch
 try:
     from vllm.model_executor.layers.quantization.aqlm import AQLMConfig
     from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig
-    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (
-        CompressedTensorsW8A8Fp8MoEMethod,
-        CompressedTensorsWNA16MoEMethod,
-    )
     from vllm.model_executor.layers.quantization.deepspeedfp import DeepSpeedFPConfig
     from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config
     from vllm.model_executor.layers.quantization.gguf import GGUFConfig
@@ -72,7 +68,8 @@ if TYPE_CHECKING:
 BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
     "fp8": Fp8Config,
     "blockwise_int8": BlockInt8Config,
-    "modelopt": ModelOptFp8Config,
+    "modelopt": ModelOptFp8Config,  # Auto-detect, defaults to FP8
+    "modelopt_fp8": ModelOptFp8Config,
     "modelopt_fp4": ModelOptFp4Config,
     "w8a8_int8": W8A8Int8Config,
     "w8a8_fp8": W8A8Fp8Config,
@@ -174,51 +171,3 @@ def monkey_patch_isinstance_for_vllm_base_layer(reverse: bool = False):
         return original_isinstance(obj, classinfo)
 
     builtins.isinstance = patched_isinstance
-
-
-def monkey_patch_moe_apply(class_obj: "FusedMoEMethodBase"):
-    """
-    Monkey patch the apply function of vllm's FusedMoEMethodBase.
-    Convert sglang arguments to vllm arguments.
-    """
-    original_apply = class_obj.apply
-    sig = inspect.signature(original_apply)
-    param_names = list(sig.parameters.keys())
-    has_correction_bias = "e_score_correction_bias" in param_names
-
-    def new_apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        topk_output: TopKOutput,
-        *,
-        activation: str = "silu",
-        apply_router_weight_on_input: bool = False,
-        inplace: bool = True,
-        no_combine: bool = False,
-        routed_scaling_factor: Optional[float] = None,
-    ):
-        assert activation == "silu"
-        assert inplace and not no_combine
-
-        kwargs = {
-            "self": self,
-            "layer": layer,
-            "x": x,
-            "topk_output": topk_output,
-        }
-        return original_apply(**kwargs)
-
-    setattr(class_obj, "apply", new_apply)
-
-
-def monkey_patch_quant_configs():
-    """Apply all monkey patches in one place."""
-
-    monkey_patch_moe_apply(CompressedTensorsW8A8Fp8MoEMethod)
-    monkey_patch_moe_apply(CompressedTensorsWNA16MoEMethod)
-
-
-# Only apply monkey patches if vllm is available
-if VLLM_AVAILABLE:
-    monkey_patch_quant_configs()