add eplb function

docs/CN/source/tutorial/api_server_args.rst

@@ -641,17 +641,6 @@ MTP 多预测参数
     增加此值允许更多预测，但确保模型与指定的步数兼容。
     目前 deepseekv3/r1 模型仅支持 1 步
 
-DeepSeek 冗余专家参数
----------------------
-
-.. option:: --ep_redundancy_expert_config_path
-
-    冗余专家配置的路径。可用于 deepseekv3 模型。
-
-.. option:: --auto_update_redundancy_expert
-
-    是否通过在线专家使用计数器为 deepseekv3 模型更新冗余专家。
-
 监控和日志参数
 --------------
 

docs/EN/source/tutorial/api_server_args.rst

@@ -644,17 +644,6 @@ MTP Multi-Prediction Parameters
     Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps.
     Currently deepseekv3/r1 models only support 1 step
 
-DeepSeek Redundant Expert Parameters
-------------------------------------
-
-.. option:: --ep_redundancy_expert_config_path
-
-    Path to redundant expert configuration. Can be used for deepseekv3 models.
-
-.. option:: --auto_update_redundancy_expert
-
-    Whether to update redundant experts for deepseekv3 models through online expert usage counters.
-
 Monitoring and Logging Parameters
 ---------------------------------
 

lightllm/common/basemodel/basemodel.py

@@ -73,6 +73,8 @@ class TpPartBaseModel:
 
     def __init__(self, kvargs):
         self.args = get_env_start_args()
+        self.eplb_manager = None
+        self.ep_balance_monitor = None
         self.run_mode = kvargs["run_mode"]
         self.weight_dir_ = kvargs["weight_dir"]
         self.max_total_token_num = kvargs["max_total_token_num"]
@@ -374,9 +376,16 @@ def forward(self, model_input: ModelInput):
         assert model_input.mem_indexes.is_cuda
 
         if model_input.is_prefill:
-            return self._prefill(model_input=model_input)
-        else:
-            return self._decode(model_input)
+            model_output = self._prefill(model_input=model_input)
+            self._after_prefill()
+            return model_output
+        return self._decode(model_input)
+
+    def _after_prefill(self):
+        if self.ep_balance_monitor is not None:
+            self.ep_balance_monitor.record_prefill_round()
+        if self.eplb_manager is not None:
+            self.eplb_manager.step()
 
     def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
         infer_state = self.infer_state_class()
@@ -863,6 +872,7 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input
         dist_group_manager.clear_deepep_buffer()
         model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event
         model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event
+        self._after_prefill()
         return model_output0, model_output1
 
     @torch.no_grad()

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_balance.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class PrefillEPBalanceCounters:
+    """Cumulative CPU loads for one EP MoE layer's completed prefill dispatches."""
+
+    route_load: int = 0
+    compute_load: int = 0
+
+    def accumulate(self, route_load: int, compute_load: int):
+        """Accumulate exact route and alignment-expanded compute loads for one prefill dispatch."""
+        self.route_load += route_load
+        self.compute_load += compute_load