[EP] Refactor DeepEP Engine Organization for Mixed Mode & Buffer Management Optimization (#3182)

* Add support for mixed-ep across multi nodes

* code refine

---------

Co-authored-by: yuanxiaolan <[email protected]>

Author: RichardWooSJTU

fastdeploy/model_executor/layers/moe/ep.py

@@ -68,25 +68,20 @@ def __init__(
         self.num_local_experts = num_experts // ep_size
         self.async_finish = async_finish
 
-        self.prefill_deepep_engine = None
-        self.decode_deepep_engine = None
+        self.deepep_engine = None
 
         self.ep_config = Config(24, 6, 256)
         self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
 
         # In mixed EP mode on a single node, we dynamically switch between
         # high throughput and low latency modes.
         if splitwise_role == "mixed":
-            # decode engine
-            logger.info("Initializing Low Latency Buffer")
-            self.get_low_latency_buffer()
-            # prefill engine
-            self.prefill_deepep_engine = deep_ep.Buffer(
+            self.deepep_engine = deep_ep.Buffer(
                 self.group,
-                int(5e8),
-                0,
-                low_latency_mode=False,
-                num_qps_per_rank=1,
+                int(2e9),
+                int(5e9),
+                low_latency_mode=True,
+                num_qps_per_rank=24,
             )
         # In disaggregated mode on mutiple nodes, we either use
         # high throughput mode or low latency mode.
@@ -95,7 +90,7 @@ def __init__(
                 logger.info("Initializing Low Latency Buffer")
                 self.get_low_latency_buffer()
             elif moe_phase.phase == "prefill":
-                self.prefill_deepep_engine = deep_ep.Buffer(
+                self.deepep_engine = deep_ep.Buffer(
                     self.group,
                     int(5e8),
                     0,
@@ -124,14 +119,14 @@ def get_low_latency_buffer(self):
         )
         # Allocate a buffer if not existed or not enough buffer size
         if (
-            self.decode_deepep_engine is None
-            or self.decode_deepep_engine.group != self.group
-            or not self.decode_deepep_engine.low_latency_mode
-            or self.decode_deepep_engine.num_rdma_bytes < num_rdma_bytes
+            self.deepep_engine is None
+            or self.deepep_engine.group != self.group
+            or not self.deepep_engine.low_latency_mode
+            or self.deepep_engine.num_rdma_bytes < num_rdma_bytes
         ):
             # NOTES: for best performance, the QP number **must** be equal to the number of the local experts
             assert self.num_experts % self.ep_size == 0
-            self.decode_deepep_engine = deep_ep.Buffer(
+            self.deepep_engine = deep_ep.Buffer(
                 self.group,
                 0,
                 num_rdma_bytes,
@@ -168,7 +163,7 @@ def low_latency_dispatch(
             handle,
             _,
             dispatch_hook,
-        ) = self.decode_deepep_engine.low_latency_dispatch(
+        ) = self.deepep_engine.low_latency_dispatch(
             hidden_states,
             topk_idx,
             expertwise_scale,
@@ -210,7 +205,7 @@ def low_latency_combine(
                 num_experts,
             )
 
-        combined_hidden_states, _, combine_hook = self.decode_deepep_engine.low_latency_combine(
+        combined_hidden_states, _, combine_hook = self.deepep_engine.low_latency_combine(
             hidden_states,
             topk_idx,
             topk_weights,
@@ -224,19 +219,15 @@ def clean_low_latency_buffer(self):
         """
         clean_low_latency_buffer
         """
-        self.decode_deepep_engine.clean_low_latency_buffer(
+        self.deepep_engine.clean_low_latency_buffer(
             self.num_max_dispatch_tokens_per_rank, self.hidden, self.num_experts
         )
 
     def barrier_all(self):
         """
         barrier_all
         """
-        if self.prefill_deepep_engine is not None:
-            self.prefill_deepep_engine.barrier_all()
-
-        if self.decode_deepep_engine is not None:
-            self.decode_deepep_engine.barrier_all()
+        self.deepep_engine.barrier_all()
 
 
 class EPRunner:
@@ -316,6 +307,9 @@ def combine(self, *args, **kwargs):
         """
         raise NotImplementedError
 
+    def clean_low_latency_buffer(self):
+        self.ep_engine.clean_low_latency_buffer()
+
 
 class EPPrefillRunner(EPRunner):
     """
@@ -328,6 +322,7 @@ def __init__(
         hidden: int,
         num_experts: int,
         splitwise_role: str,
+        num_max_dispatch_tokens_per_rank: int,
         ep_size: int = 1,
         ep_rank: int = 0,
         redundant_experts_num: int = 0,
@@ -339,7 +334,7 @@ def __init__(
             num_experts,
             splitwise_role,
             moe_phase,
-            num_max_dispatch_tokens_per_rank=256,
+            num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank,
             ep_size=ep_size,
             ep_rank=ep_rank,
             redundant_experts_num=redundant_experts_num,
@@ -359,7 +354,7 @@ def dispatch(
             num_tokens_per_expert,
             is_token_in_rank,
             _,
-        ) = self.ep_engine.prefill_deepep_engine.get_dispatch_layout(topk_idx, self.num_experts)
+        ) = self.ep_engine.deepep_engine.get_dispatch_layout(topk_idx, self.num_experts)
 
         x_scale_tensor = kwargs.get("x_scale_tensor", None)
         dispatch_args = {
@@ -372,7 +367,7 @@ def dispatch(
             "topk_idx": topk_idx,
             "topk_weights": topk_weights,
         }
-        return self.ep_engine.prefill_deepep_engine.dispatch(**dispatch_args)
+        return self.ep_engine.deepep_engine.dispatch(**dispatch_args)
 
     def combine(
         self,
@@ -387,14 +382,14 @@ def combine(
             "async_finish": self.ep_engine.async_finish,
             "topk_weights": recv_topk_weights,
         }
-        fused_moe_out, _, _ = self.ep_engine.prefill_deepep_engine.combine(**combine_args)
+        fused_moe_out, _, _ = self.ep_engine.deepep_engine.combine(**combine_args)
 
         return fused_moe_out
 
 
 class EPDecoderRunner(EPRunner):
     """
-    EPPrefillRunner
+    EPDecoderRunner
     """
 
     def __init__(

fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py

@@ -51,6 +51,7 @@ def init_ep(self, layer: nn.Layer) -> None:
                     layer.hidden_size,
                     layer.num_experts,
                     layer.fd_config.parallel_config.splitwise_role,
+                    layer.fd_config.model_config.num_max_dispatch_tokens_per_rank,
                     layer.ep_size,
                     layer.ep_rank,
                     layer.fd_config.model_config.redundant_experts_num,
@@ -74,6 +75,7 @@ def init_ep(self, layer: nn.Layer) -> None:
                         layer.hidden_size,
                         layer.num_experts,
                         layer.fd_config.parallel_config.splitwise_role,
+                        layer.fd_config.model_config.num_max_dispatch_tokens_per_rank,
                         layer.ep_size,
                         layer.ep_rank,
                         layer.fd_config.model_config.redundant_experts_num,
@@ -165,8 +167,10 @@ def apply(
         """
         if layer.ep_size > 1:
             if layer.fd_config.parallel_config.moe_phase.phase == "prefill":
+                self.ep_prefill_runner.clean_low_latency_buffer()
                 return self.apply_ep_prefill(layer, x, gate_out)
             else:
+                self.ep_decoder_runner.clean_low_latency_buffer()
                 return self.apply_ep_decode(layer, x, gate_out)
         else:
             return self.apply_tp(layer, x, gate_out)