update code

Author: Wanglongzhi2001

fastdeploy/config.py

@@ -651,7 +651,8 @@ def __init__(
         self.enable_expert_parallel = False
         self.enable_chunked_moe = False
         self.chunked_moe_size = 256
-
+        self.enable_mega_moe = False
+        
         self.local_data_parallel_id = 0
         # Engine worker queue port
         self.engine_worker_queue_port: Union[int, str, list] = None

fastdeploy/engine/args_utils.py

@@ -369,6 +369,11 @@ class EngineArgs:
     Whether use chunked moe.
     """
 
+    enable_mega_moe: bool = False
+    """
+    Whether use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear.
+    """
+
     chunked_moe_size: int = 256
     """
     Chunk size of moe input.
@@ -1176,6 +1181,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.enable_chunked_moe,
             help="Use chunked moe.",
         )
+        parallel_group.add_argument(
+            "--enable-mega-moe",
+            action="store_true",
+            default=EngineArgs.enable_mega_moe,
+            help="Use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear.",
+        )
         parallel_group.add_argument(
             "--chunked-moe-size",
             type=int,

fastdeploy/engine/engine.py

@@ -685,6 +685,7 @@ def _start_worker_service(self):
         worker_store_true_flag = {
             "enable_expert_parallel": self.cfg.parallel_config.enable_expert_parallel,
             "enable_chunked_moe": self.cfg.parallel_config.enable_chunked_moe,
+            "enable_mega_moe": self.cfg.parallel_config.enable_mega_moe,
             "enable_prefix_caching": self.cfg.cache_config.enable_prefix_caching,
             "enable_chunked_prefill": self.cfg.cache_config.enable_chunked_prefill,
             "do_profile": self.do_profile,

fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py

@@ -39,6 +39,7 @@
     free_tensor,
     set_weight_attrs,
 )
+from fastdeploy.model_executor.layers.moe.ep import EPRunner
 
 from .utils import get_moe_scores
 
@@ -423,7 +424,7 @@ def apply_ep_prefill(
         """
         gate_out = gate(x.cast("float32"))
         # 1. Select topk experts and weights
-        topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out)
+        topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         # 2. Dynamic compute blockwise quantization scales
         if "a_tokenwise_int8" in self.xpu_moe_quant_type:
@@ -518,7 +519,7 @@ def apply_ep_decode(
         gate_out = gate(x.cast("float32"))
 
         # 1. Select topk experts and weights
-        topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out)
+        topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         # 2. EP Dispatch
         if "a_tokenwise_int8" in self.xpu_moe_quant_type:

fastdeploy/model_executor/layers/moe/ep.py

@@ -490,7 +490,8 @@ def __init__(
             top_k=self.top_k,
         )
 
-    def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
+    @staticmethod
+    def moe_select(layer: nn.Layer, gate_out: paddle.Tensor):
         if layer.redundant_table_manger is not None:
             (
                 ep_rank_to_expert_id_list,
@@ -523,7 +524,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
                     expert_in_rank_num_list=expert_in_rank_num_list,
                     tokens_per_expert_stats_list=tokens_per_expert_stats_list,
                     bias=layer.gate_correction_bias,
-                    moe_topk=self.top_k,
+                    moe_topk=layer.top_k,
                     apply_norm_weight=True,
                     enable_softmax_top_k_fused=False,
                     redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1,
@@ -550,7 +551,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
                 topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select(
                     gate_out,
                     layer.gate_correction_bias,
-                    self.top_k,
+                    layer.top_k,
                     True,
                     False,
                 )
@@ -788,3 +789,20 @@ def combine(self, ffn_out, topk_idx, topk_weights, handle, **kwargs):
             combine_hook()
 
         return combined_hidden_states
+
+
+class FakeEPRunner:
+    """ """
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def dispatch(self, *args, **kwargs):
+        """ """
+        pass
+
+    def combine(self, *args, **kwargs):
+        """ """
+        pass
+    
+    def clean_low_latency_buffer(self):
+        pass

fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py

@@ -232,29 +232,24 @@ def apply(
         """
         if layer.ep_size > 1:
             is_moe_start_layer = layer.layer_idx == layer.fd_config.model_config.moe_layer_start_index
-            if envs.FD_ENABLE_MAGE_MOE:
-                return self.apply_mage_moe(
-                    layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj
+            if layer.fd_config.model_config.moe_phase.phase == "prefill":
+                if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
+                    self.ep_prefill_runner.clean_low_latency_buffer()
+                return self.apply_ep_prefill(
+                    layer,
+                    x,
+                    gate,
+                    topk_ids_hookfunc,
+                    shared_experts,
+                    fc1_latent_proj,
+                    fc2_latent_proj,
                 )
             else:
-                if layer.fd_config.model_config.moe_phase.phase == "prefill":
-                    if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
-                        self.ep_prefill_runner.clean_low_latency_buffer()
-                    return self.apply_ep_prefill(
-                        layer,
-                        x,
-                        gate,
-                        topk_ids_hookfunc,
-                        shared_experts,
-                        fc1_latent_proj,
-                        fc2_latent_proj,
-                    )
-                else:
-                    if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
-                        self.ep_decoder_runner.clean_low_latency_buffer()
-                    return self.apply_ep_decode(
-                        layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj
-                    )
+                if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
+                    self.ep_decoder_runner.clean_low_latency_buffer()
+                return self.apply_ep_decode(
+                    layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj
+                )
         else:
             return self.apply_tp(layer, x, gate, topk_ids_hookfunc, fc1_latent_proj, fc2_latent_proj)
 

fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py

@@ -23,7 +23,7 @@
 from paddleformers.utils.log import logger
 
 import fastdeploy
-from fastdeploy.model_executor.layers.moe.ep import deep_ep
+from fastdeploy.model_executor.layers.moe.ep import deep_ep, EPRunner
 from fastdeploy.model_executor.layers.quantization.fp8_utils import (
     deep_gemm,
     paddlefleet_ops,
@@ -642,7 +642,7 @@ def apply_ep_prefill(
                 getattr(layer, "renormalize", True),
             )
         else:
-            topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out)
+            topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         if topk_ids_hookfunc is not None:
             topk_ids_hookfunc(topk_ids=topk_idx)
@@ -964,7 +964,7 @@ def apply_ep_decode(
         gate_out = gate(x)
         gate_out = gate_out.cast("float32")
         # 1. Select topk experts and weights
-        topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out)
+        topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         if topk_ids_hookfunc is not None:
             topk_ids_hookfunc(topk_ids=topk_idx)

fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py

@@ -52,6 +52,7 @@
     set_weight_attrs,
     weight_fully_copied,
 )
+from fastdeploy.model_executor.layers.moe.ep import EPRunner
 
 
 def m_grouped_bf16_gemm_nn_contiguous(x, y, expert_idx_per_token):
@@ -141,7 +142,7 @@ def apply_ep_prefill(
         if fc1_latent_proj is not None:
             x = fc1_latent_proj(x)
         # 1. Select topk experts and weights
-        topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out)
+        topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         if layer.routed_scaling_factor_learnable:
             safe_topk_indices = paddle.clip(topk_idx, min=0)
@@ -304,7 +305,7 @@ def apply_ep_decode(
 
         estimate_total_token_nums = gate_out.shape[0] * layer.top_k
         # 1. Select topk experts and weights
-        topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out)
+        topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out)
 
         if layer.routed_scaling_factor_learnable:
             safe_topk_indices = paddle.clip(topk_idx, min=0)