add eplb function

docs/CN/source/tutorial/api_server_args.rst

@@ -794,17 +794,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

@@ -825,17 +825,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

@@ -28,7 +28,7 @@
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.dist_utils import get_dp_world_size
-from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory
+from lightllm.utils.profile_max_tokens import profile_model_memory
 from lightllm.utils.envs_utils import (
     get_env_start_args,
     get_llm_data_type,
@@ -109,7 +109,7 @@ def __init__(self, kvargs):
         self._init_quant()
 
         enable_weight_cpu_backup = self.args.enable_weight_cpu_backup
-        with profile_mtp_weight_memory(self):
+        with profile_model_memory(self):
             with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup):
                 self._init_weights()
         with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE):
@@ -148,6 +148,7 @@ def __init__(self, kvargs):
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             self.config = json.load(json_file)
+        self.config["is_mtp_draft_model"] = self.is_mtp_draft_model
         # rename keys
         repair_config(self.config, same_names=["num_attention_heads", "n_head"])
         repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"])
@@ -631,7 +632,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"

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