feat: imbalance statistics

lightllm/common/basemodel/basemodel.py

@@ -63,6 +63,7 @@ class TpPartBaseModel:
 
     def __init__(self, kvargs):
         self.args = get_env_start_args()
+        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"]
@@ -354,9 +355,14 @@ def forward(self, model_input: ModelInput):
         assert model_input.mem_indexes.is_cuda
 
         if model_input.is_prefill:
-            return self._prefill(model_input)
-        else:
-            return self._decode(model_input)
+            model_output = self._prefill(model_input)
+            self._record_prefill_ep_balance()
+            return model_output
+        return self._decode(model_input)
+
+    def _record_prefill_ep_balance(self):
+        if self.ep_balance_monitor is not None:
+            self.ep_balance_monitor.record_prefill_round()
 
     def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
         infer_state = self.infer_state_class()
@@ -842,6 +848,7 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
         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._record_prefill_ep_balance()
         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

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py

@@ -20,6 +20,10 @@
 
 
 class FuseMoeDeepGEMM(FuseMoeTriton):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ep_balance_counters = None
+
     def _select_experts(
         self,
         input_tensor: torch.Tensor,
@@ -87,6 +91,7 @@ def _fused_experts(
             quant_method=self.quant_method,
             is_prefill=is_prefill,
             previous_event=None,  # for overlap
+            ep_balance_counters=self.ep_balance_counters,
         )
         return output
 
@@ -181,8 +186,23 @@ def dispatch(
             use_tma_aligned_col_major_sf=True,
         )
 
-        def hook():
-            event.current_stream_wait()
+        counters = self.ep_balance_counters
+        if counters is None:
+
+            def hook():
+                event.current_stream_wait()
+
+        else:
+            # Sent routes are globally conserved by all-to-all; recv_x[0] is the 128-aligned expanded compute load.
+            route_load = topk_idx.numel()
+            compute_load = recv_x[0].shape[0]
+
+            def hook():
+                event.current_stream_wait()
+                counters.accumulate(
+                    route_load=route_load,
+                    compute_load=compute_load,
+                )
 
         return recv_x, recv_topk_idx, recv_topk_weights, handle.num_recv_tokens_per_expert_list, handle, hook
 

lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py

@@ -19,6 +19,7 @@
     ep_gather_chunk,
     ep_zero_padding,
 )
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters
 from lightllm.utils.envs_utils import (
     get_deepep_num_max_dispatch_tokens_per_rank_prefill,
     get_deepep_num_max_dispatch_tokens_per_rank_decode,
@@ -201,6 +202,7 @@ def fused_experts(
     quant_method: Any,
     is_prefill: Optional[bool],
     previous_event: Optional[Any] = None,
+    ep_balance_counters: Optional[PrefillEPBalanceCounters] = None,
 ):
     check_ep_expert_dtype(quant_method)
     if use_sm100_mega_moe(quant_method):
@@ -222,6 +224,7 @@ def fused_experts(
         w1_scale=w13.weight_scale,
         w2_scale=w2.weight_scale,
         previous_event=previous_event,
+        ep_balance_counters=ep_balance_counters,
     )
 
 
@@ -240,6 +243,7 @@ def fused_experts_impl(
     w1_scale: Optional[torch.Tensor] = None,
     w2_scale: Optional[torch.Tensor] = None,
     previous_event: Optional[Any] = None,
+    ep_balance_counters: Optional[PrefillEPBalanceCounters] = None,
 ):
     # Check constraints.
     assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
@@ -290,6 +294,12 @@ def fused_experts_impl(
             do_expand=True,
             use_tma_aligned_col_major_sf=True,
         )
+        if ep_balance_counters is not None:
+            # Sent routes are globally conserved by all-to-all; recv_x[0] is the 128-aligned expanded compute load.
+            ep_balance_counters.accumulate(
+                route_load=topk_idx.numel(),
+                compute_load=recv_x[0].shape[0],
+            )
         # Dispatch is synchronous in this path.  Its FP8 source is no longer
         # needed once the received tensors have been produced.
         del qinput_tensor, input_scale

lightllm/distributed/communication_op.py

@@ -107,6 +107,7 @@ def all_gather_into_tensor(self, output_: torch.Tensor, input_: torch.Tensor, as
 class DistributeGroupManager:
     def __init__(self):
         self.groups = []
+        self.ep_balance_monitor_group = None
         self.ep_buffer = None
         self.ep_low_latency_buffer = None
         self.ep_mega_moe_buffer = None
@@ -124,6 +125,14 @@ def create_groups(self, group_size: int):
             if not args.disable_flashinfer_allreduce:
                 group.init_flashinfer_reduce()
             self.groups.append(group)
+        if (
+            getattr(args, "enable_ep_moe", False)
+            and not getattr(args, "disable_ep_balance_monitor", False)
+            and getattr(args, "run_mode", "normal") != "decode"
+            and not getattr(args, "enable_prefill_cudagraph", False)
+            and not is_sm100_gpu()
+        ):
+            self.ep_balance_monitor_group = dist.new_group(ranks=list(range(get_global_world_size())), backend="gloo")
         return
 
     def get_default_group(self) -> CustomProcessGroup:

lightllm/server/api_cli.py

@@ -710,6 +710,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
         action="store_true",
         help="""Whether to enable ep moe for deepseekv3 model.""",
     )
+    parser.add_argument(
+        "--disable_ep_balance_monitor",
+        action="store_true",
+        help="""Disable the prefill expert balance monitor enabled by default for EP-MoE.""",
+    )
     parser.add_argument(
         "--ep_redundancy_expert_config_path",
         type=str,

lightllm/server/core/objs/start_args_type.py

@@ -176,6 +176,7 @@ class StartArgs:
         default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]}
     )
     enable_ep_moe: bool = field(default=False)
+    disable_ep_balance_monitor: bool = field(default=False)
     ep_redundancy_expert_config_path: Optional[str] = field(default=None)
     auto_update_redundancy_expert: bool = field(default=False)
     enable_fused_shared_experts: bool = field(default=False)

lightllm/server/metrics/metrics.py

@@ -32,6 +32,15 @@
     "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request",
     "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)",
     "lightllm_num_running_reqs": "Number of running requests",
+    "lightllm_prefill_ep_critical_overhead_gflops_per_routed_token": (
+        "Estimated critical-path excess compute per logical routed source token, GFLOPs/token"
+    ),
+    "lightllm_prefill_ep_compute_critical_overhead_ratio": (
+        "Estimated excess critical compute divided by balanced compute; 0.3 means +30%"
+    ),
+    "lightllm_prefill_ep_placement_pressure_drift": (
+        "Normalized temporal drift of overloaded-rank pressure from the latest complete prefill report"
+    ),
 }
 
 
@@ -111,6 +120,9 @@ def init_metrics(self, args):
         self.create_gauge("lightllm_cache_hit_rate")
         self.create_gauge("lightllm_gen_throughput")
         self.create_gauge("lightllm_num_running_reqs")
+        self.create_gauge("lightllm_prefill_ep_critical_overhead_gflops_per_routed_token")
+        self.create_gauge("lightllm_prefill_ep_compute_critical_overhead_ratio")
+        self.create_gauge("lightllm_prefill_ep_placement_pressure_drift")
 
     def create_histogram(self, name, buckets, labelnames=None):
         all_labels = ["model_name"] + (labelnames or [])