add deepep

Author: zhaochaoxing

lmdeploy/cli/serve.py

@@ -165,6 +165,7 @@ def add_parser_api_server():
         quant_policy = ArgumentHelper.quant_policy(pt_group)
         ArgumentHelper.dp(pt_group)
         ArgumentHelper.dp_rank(pt_group)
+        ArgumentHelper.ep(pt_group)
 
         # turbomind args
         tb_group = parser.add_argument_group('TurboMind engine arguments')
@@ -296,6 +297,7 @@ def api_server(args):
                                                  tp=args.tp,
                                                  dp=args.dp,
                                                  dp_rank=args.dp_rank,
+                                                 ep=args.ep,
                                                  max_batch_size=max_batch_size,
                                                  cache_max_entry_count=args.cache_max_entry_count,
                                                  block_size=args.cache_block_seq_len,

lmdeploy/cli/utils.py

@@ -157,6 +157,12 @@ def dp(parser):
         """Add argument dp to parser."""
 
         return parser.add_argument('--dp', type=int, default=1, help='data parallelism. dp_rank is required.')
+    
+    @staticmethod
+    def ep(parser):
+        """Add argument ep to parser."""
+
+        return parser.add_argument('--ep', type=int, default=1, help='expert parallelism. Should be 2^n.')
 
     @staticmethod
     def dp_rank(parser):

lmdeploy/messages.py

@@ -294,6 +294,7 @@ class PytorchEngineConfig:
     tp: int = 1
     dp: int = 1
     dp_rank: int = 0
+    ep: int = 1
     session_len: int = None
     max_batch_size: int = None
     cache_max_entry_count: float = 0.8
@@ -318,6 +319,7 @@ def __post_init__(self):
         assert self.dtype in ['auto', 'float16', 'bfloat16']
         assert self.tp >= 1, 'invalid tp'
         assert self.dp >= 1, 'invalid dp'
+        assert self.ep >= 1, 'invalid ep'
         assert 0 < self.cache_max_entry_count < 1, \
             'invalid cache_max_entry_count'
         assert self.num_cpu_blocks >= 0, 'invalid num_cpu_blocks'

lmdeploy/pytorch/backends/cuda/moe.py

@@ -3,10 +3,14 @@
 from typing import List
 
 import torch
-
+import torch.distributed as dist
+from lmdeploy.pytorch.backends.cuda.token_dispatcher import DeepEPDispatcher
+from lmdeploy.pytorch.distributed import get_dist_manager
 from lmdeploy.pytorch.kernels.cuda import fused_moe, fused_moe_w8a8
 from lmdeploy.pytorch.kernels.cuda.blocked_fp8_fused_moe import fused_moe_blocked_fp8
 from lmdeploy.pytorch.kernels.cuda.blocked_gemm_fp8 import quant_fp8
+from lmdeploy.pytorch.kernels.cuda.fused_moe import _renormalize
+from lmdeploy.pytorch.kernels.cuda.ep_moe import grouped_gemm_triton, silu_and_mul_triton_kernel
 from lmdeploy.pytorch.kernels.cuda.w8a8_triton_kernels import per_token_quant_int8
 from lmdeploy.pytorch.models.q_modules import QTensor
 
@@ -227,18 +231,188 @@ def forward(self,
         return output
 
 
+
+class DeepEPMoE:
+    """
+    MoE Expert Parallel Impl based on DeepEP (https://github.com/deepseek-ai/DeepEP/tree/main)
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        ep_size: int,
+        block_shape: list[int],
+    ):
+        self.num_experts = num_experts
+        self.ep_size = ep_size
+        assert self.num_experts % self.ep_size == 0
+        self.num_experts_per_partition = self.num_experts // self.ep_size
+        self.block_shape = block_shape
+        self.use_fp8_w8a8 = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        tokens_per_expert: torch.Tensor,
+        gate_up_weight:torch.Tensor,
+        gate_up_scale:torch.Tensor,
+        gate_down_weight:torch.Tensor,
+        gate_down_scale:torch.Tensor
+    ):
+        seg_indptr_cur_rank = torch.cat(
+            [
+                torch.zeros(
+                    1, device=tokens_per_expert.device, dtype=tokens_per_expert.dtype
+                ),
+                torch.cumsum(tokens_per_expert, dim=0),
+            ]
+        )
+        reorder_topk_ids = torch.repeat_interleave(tokens_per_expert)
+        weight_indices_cur_rank = torch.arange(
+            0,
+            self.num_experts_per_partition,
+            device=hidden_states.device,
+            dtype=torch.int64,
+        )
+
+        # GroupGemm-0
+        gateup_output = torch.empty(
+            hidden_states.shape[0],
+            gate_up_weight.shape[1],
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        if hidden_states.shape[0] > 0:
+            input, input_scale = quant_fp8(hidden_states, 128, dtype=gate_up_weight.dtype)
+            gateup_output = grouped_gemm_triton(
+                a=input,
+                b=gate_up_weight,
+                c=gateup_output,
+                batch_size=self.num_experts_per_partition,
+                weight_column_major=True,
+                seg_indptr=seg_indptr_cur_rank,
+                weight_indices=weight_indices_cur_rank,
+                use_fp8_w8a8=self.use_fp8_w8a8,
+                scale_a=input_scale,
+                scale_b=gate_up_scale,
+                block_shape=self.block_shape,
+            )
+
+        # Act
+        down_input = torch.empty(
+            gateup_output.shape[0],
+            gateup_output.shape[1] // 2,
+            device=gateup_output.device,
+            dtype=hidden_states.dtype,
+        )
+        silu_and_mul_triton_kernel[(gateup_output.shape[0],)](
+            gateup_output,
+            down_input,
+            gateup_output.shape[1],
+            reorder_topk_ids,
+            None,
+            0,
+            self.num_experts_per_partition - 1,
+            BLOCK_SIZE=512,
+        )
+
+        # GroupGemm-1
+        down_output = torch.empty(
+            down_input.shape[0],
+            gate_down_weight.shape[1],
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        if down_input.shape[0] > 0:
+            down_input, down_input_scale = quant_fp8(down_input, 128, dtype=gate_down_weight.dtype)
+            down_output = grouped_gemm_triton(
+                a=down_input,
+                b=gate_down_weight,
+                c=down_output,
+                batch_size=self.num_experts_per_partition,
+                weight_column_major=True,
+                seg_indptr=seg_indptr_cur_rank,
+                weight_indices=weight_indices_cur_rank,
+                use_fp8_w8a8=self.use_fp8_w8a8,
+                scale_a=down_input_scale,
+                scale_b=gate_down_scale, 
+                block_shape=self.block_shape,
+            )
+        return down_output
+
+
+class FusedDeepEpMoEBlockedF8Impl(TritonFusedMoEBlockedF8Impl):
+    def __init__(self,
+                 ep_size: int,
+                 ep_group:dist.ProcessGroup,
+                 top_k: int,
+                 num_experts: int,
+                 hidden_dim: int,
+                 renormalize: bool = False,
+                 block_size: int = 128,
+                 out_dtype: torch.dtype = torch.bfloat16):
+        super().__init__(top_k, num_experts, renormalize, block_size, out_dtype)
+        self.token_dispatcher = DeepEPDispatcher(
+                group=ep_group,
+                router_topk=self.top_k,
+                permute_fusion=True,
+                num_experts=self.num_experts,
+                num_local_experts=self.num_experts // ep_size,
+                hidden_size=hidden_dim,
+                params_dtype=out_dtype,
+            )
+        self.experts = DeepEPMoE(num_experts, ep_size, [block_size,block_size])
+    
+    def forward(self,
+                hidden_states: torch.Tensor,
+                topk_weights: torch.Tensor,
+                topk_ids: torch.LongTensor,
+                gate_up_weights: torch.Tensor,
+                gate_up_scale: torch.Tensor,
+                down_weights: torch.Tensor,
+                down_scale: torch.Tensor,
+                expert_list: List[int] = None):
+        """forward."""
+        topk_weights = _renormalize(topk_weights, self.renormalize)
+        recv_hidden_states, recv_topk_ids, recv_topk_weights, tokens_per_expert = (
+            self.token_dispatcher.dispatch(
+                hidden_states,
+                topk_ids.to(torch.int32),
+                topk_weights.to(torch.float32),
+                self.num_experts,
+            )
+        )
+        out_states = self.experts.forward(recv_hidden_states, tokens_per_expert, gate_up_weights, gate_up_scale,
+                                 down_weights, down_scale)
+        out_states = self.token_dispatcher.combine(out_states)
+        return out_states
+
 class TritonFusedMoEBlockedF8Builder(FusedMoEBlockedF8Builder):
     """triton fused moe blocked f8 builder."""
 
     @staticmethod
     def build(top_k: int,
               num_experts: int,
+              hidden_dim: int,
               renormalize: bool = False,
               block_size: int = 128,
+              ep_size: int = 1,
+              ep_group: dist.ProcessGroup = None,
               out_dtype: torch.dtype = torch.float16):
         """build from mlp."""
-        return TritonFusedMoEBlockedF8Impl(top_k=top_k,
-                                           num_experts=num_experts,
-                                           renormalize=renormalize,
-                                           block_size=block_size,
-                                           out_dtype=out_dtype)
+        if ep_size > 1:
+            return FusedDeepEpMoEBlockedF8Impl(ep_size=ep_size,
+                                            ep_group=ep_group,
+                                            top_k=top_k,
+                                            num_experts=num_experts,
+                                            hidden_dim=hidden_dim,
+                                            renormalize=renormalize,
+                                            block_size=block_size,
+                                            out_dtype=out_dtype)
+        else:
+            return TritonFusedMoEBlockedF8Impl(top_k=top_k,
+                                            num_experts=num_experts,
+                                            renormalize=renormalize,
+                                            block_size=block_size,
+                                            out_dtype=out_dtype)
+