Support PPLX kernels for Deepseek models

python/sglang/srt/layers/communicator.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
+import logging
 from dataclasses import dataclass
 from enum import Enum, auto
 from functools import partial
@@ -28,13 +28,16 @@
     attn_tp_reduce_scatter,
     dp_gather_partial,
     dp_scatter,
+    get_attention_tp_group,
     get_attention_tp_rank,
     get_attention_tp_size,
     get_local_attention_dp_size,
 )
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
+logger = logging.getLogger(__name__)
+
 
 class ScatterMode(Enum):
     """
@@ -101,11 +104,12 @@ def _compute_layer_input_mode(cls, context: _LayerModeComputationContext):
     @classmethod
     def _compute_mlp_mode(cls, context: _LayerModeComputationContext):
         if context.is_layer_sparse:
-            return (
-                ScatterMode.SCATTERED
-                if global_server_args_dict["enable_deepep_moe"]
-                else ScatterMode.FULL
-            )
+            if global_server_args_dict["enable_deepep_moe"]:
+                return ScatterMode.SCATTERED
+            elif global_server_args_dict["enable_pplx_moe"]:
+                return ScatterMode.TP_ATTN_FULL
+            else:
+                return ScatterMode.FULL
         else:
             return (
                 ScatterMode.SCATTERED
@@ -118,7 +122,7 @@ def _compute_middle_residual_mode(cls, context: _LayerModeComputationContext):
         mlp_mode = cls._compute_mlp_mode(context)
         if mlp_mode == ScatterMode.SCATTERED:
             return ScatterMode.SCATTERED
-        if mlp_mode == ScatterMode.FULL:
+        if mlp_mode == ScatterMode.FULL or mlp_mode == ScatterMode.TP_ATTN_FULL:
             return ScatterMode.TP_ATTN_FULL
         raise NotImplementedError
 
@@ -129,7 +133,7 @@ def _compute_layer_output_mode(cls, context: _LayerModeComputationContext):
             return ScatterMode.model_input_output()
         if mlp_mode == ScatterMode.SCATTERED:
             return ScatterMode.SCATTERED
-        if mlp_mode == ScatterMode.FULL:
+        if mlp_mode == ScatterMode.FULL or mlp_mode == ScatterMode.TP_ATTN_FULL:
             return ScatterMode.TP_ATTN_FULL
         raise NotImplementedError
 
@@ -342,8 +346,21 @@ def get_fn(
                 residual_input_mode=residual_input_mode,
             )
 
+        if (
+            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
+            and (
+                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
+            )
+            and (hidden_states_output_mode == ScatterMode.TP_ATTN_FULL)
+            and (residual_output_mode == ScatterMode.TP_ATTN_FULL)
+        ):
+            return partial(
+                CommunicateWithAllReduceAndLayerNormFn._all_reduce_hidden_states_and_residual,
+                residual_input_mode=residual_input_mode,
+            )
+
         raise NotImplementedError(
-            f"{hidden_states_input_mode=} {residual_input_mode=} {residual_output_mode=} {residual_output_mode=}"
+            f"{hidden_states_input_mode=} {residual_input_mode=} {hidden_states_output_mode=} {residual_output_mode=}"
         )
 
     @staticmethod
@@ -402,6 +419,28 @@ def _scatter_hidden_states_and_residual(
             hidden_states, residual = layernorm(hidden_states, residual)
         return hidden_states, residual
 
+    @staticmethod
+    def _all_reduce_hidden_states_and_residual(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        residual_input_mode,
+    ):
+        if hidden_states.shape[0] != 0:
+            if residual_input_mode == ScatterMode.SCATTERED:
+                tensor_list = list(hidden_states.tensor_split(context.attn_tp_size))
+                tensor_list[context.attn_tp_rank] += residual
+            else:
+                if context.attn_tp_rank == 0:
+                    hidden_states += residual
+            hidden_states = get_attention_tp_group().all_reduce(hidden_states)
+            residual = hidden_states
+            hidden_states = layernorm(hidden_states)
+        return hidden_states, residual
+
 
 class CommunicateSummableTensorPairFn:
     """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed."""

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

@@ -4,6 +4,7 @@
 import torch
 from torch.nn import Module
 
+from sglang.srt.layers.moe.ep_moe.modular_kernel import _moe_problem_size
 from sglang.srt.layers.quantization.deep_gemm import _ENABLE_JIT_DEEPGEMM
 from sglang.srt.managers.expert_location import get_global_expert_location_metadata
 from sglang.srt.managers.expert_location_dispatch import ExpertLocationDispatchInfo
@@ -42,6 +43,7 @@
     silu_and_mul_triton_kernel,
     tma_align_input_scale,
 )
+from sglang.srt.layers.moe.ep_moe.modular_kernel import BatchedTritonExperts
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE, FusedMoEMethodBase
 from sglang.srt.layers.moe.topk import select_experts
@@ -1257,4 +1259,136 @@ def get_moe_impl_class():
         return DeepEPMoE
     if global_server_args_dict["enable_ep_moe"]:
         return EPMoE
+    if global_server_args_dict["enable_pplx_moe"]:
+        return PPlxMoE
     return FusedMoE
+
+
+# Some of the codes are adapted from https://github.com/vllm-project/vllm/blob/f9c069c85e029830094ff9abb926ffbf37b7c7e7/vllm/model_executor/layers/fused_moe/layer.py#L722
+class PPlxMoE(EPMoE):
+    """
+    MoE Expert based on pplx-kernels
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        top_k: int,
+        hidden_size: int,
+        intermediate_size: int,
+        layer_id: int,
+        params_dtype: Optional[torch.dtype] = None,
+        renormalize: bool = True,
+        use_grouped_topk: bool = False,
+        num_expert_group: Optional[int] = None,
+        num_fused_shared_experts: int = 0,
+        topk_group: Optional[int] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        tp_size: Optional[int] = None,
+        prefix: str = "",
+        correction_bias: Optional[torch.Tensor] = None,
+        custom_routing_function: Optional[Callable] = None,
+        activation: str = "silu",
+        routed_scaling_factor: Optional[float] = None,
+        dp_size: int = 1,
+        max_tokens: int = 1024,
+    ):
+        if params_dtype == None:
+            params_dtype = torch.bfloat16
+
+        assert (
+            num_fused_shared_experts == 0
+        ), "num_fused_shared_experts is not supported in EP"
+
+        super().__init__(
+            num_experts,
+            top_k,
+            hidden_size,
+            intermediate_size,
+            layer_id,
+            params_dtype,
+            renormalize,
+            use_grouped_topk,
+            num_expert_group,
+            num_fused_shared_experts,
+            topk_group,
+            quant_config,
+            tp_size,
+            prefix,
+            correction_bias,
+            custom_routing_function,
+            activation,
+            routed_scaling_factor,
+        )
+        assert (
+            num_fused_shared_experts == 0
+        ), "num_fused_shared_experts is not supported in EP"
+
+        self._fused_experts = BatchedTritonExperts(
+            max_num_tokens=max_tokens,
+            dp_size=dp_size,
+            use_fp8_w8a8=False,
+            use_int8_w8a8=False,
+            use_int8_w8a16=False,
+            use_int4_w4a16=False,
+            block_shape=None,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        expert_x: torch.Tensor,
+        expert_x_scale: torch.Tensor,
+        expert_num_tokens: torch.Tensor,
+        global_num_experts: int = -1,
+        inplace: bool = False,
+        no_combine: bool = False,
+        routed_scaling_factor: Optional[float] = None,
+    ):
+        # unpack
+        a1 = hidden_states
+        a1q = expert_x
+        a1q_scale = expert_x_scale
+        topk_ids = topk_idx.type(torch.uint32)
+        E, M, N, K, top_k = _moe_problem_size(
+            a1, self.w13_weight, self.w2_weight, topk_ids
+        )
+
+        if global_num_experts == -1:
+            global_num_experts = E
+
+        workspace13_shape, workspace2_shape, workspace_dtype = (
+            self._fused_experts.workspace_shapes(a1, M, N, K, top_k, global_num_experts)
+        )
+
+        # We can reuse the memory between cache1 and cache3 because by the time
+        # we need cache3, we're done with cache1
+        workspace13 = torch.zeros(
+            workspace13_shape, device=a1.device, dtype=workspace_dtype
+        )
+        workspace2 = torch.zeros(
+            workspace2_shape, device=a1.device, dtype=workspace_dtype
+        )
+
+        fused_out = self._fused_experts.apply(
+            a1q,
+            self.w13_weight,
+            self.w2_weight,
+            topk_ids,
+            activation=self.activation,
+            global_num_experts=global_num_experts,
+            expert_map=None,  #
+            w1_scale=self.w13_input_scale,
+            w2_scale=self.w2_input_scale,
+            w1_zp=None,
+            w2_zp=None,
+            a1q_scale=a1q_scale,
+            a2_scale=None,
+            workspace13=workspace13,
+            workspace2=workspace2,
+            expert_num_tokens=expert_num_tokens,
+        )
+
+        return fused_out