Add explicit EPLB balancedness reporting modes (#34998)

docs/docs/advanced_features/server_arguments.mdx

@@ -1885,10 +1885,10 @@ Please consult the documentation below and [server_args.py](https://github.com/s
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
     </tr>
         <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-expert-distribution-metrics`</td>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--expert-balancedness-report-mode`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable logging metrics for expert balancedness</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Where to report expert balancedness.</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>off</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>off</code>, <code>server_log</code>, <code>prometheus</code>, <code>both</code></td>
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--deepep-config`</td>

docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx

@@ -1709,9 +1709,9 @@ If the value is int8, you must also set the environment variable:DEEP_NORMAL_MOD
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
     </tr>
     <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-expert-distribution-metrics`</td>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--expert-balancedness-report-mode`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>off</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag (set to enable)</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>off</code>, <code>server_log</code>, <code>prometheus</code>, <code>both</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
     </tr>
     <tr>

docs/docs/references/environment_variables.mdx

@@ -1745,7 +1745,7 @@ SGLang supports various environment variables that can be used to configure its
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Emit an EPLB balancedness metric.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Removed. Use <code>--expert-balancedness-report-mode=prometheus</code> to emit the EPLB balancedness metric.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
     <tr>

python/sglang/srt/arg_groups/argparse_actions.py

@@ -30,12 +30,15 @@ def print_deprecated_warning(message: str):
 
 
 class DeprecatedAction(argparse.Action):
-    def __init__(self, option_strings, dest, nargs=0, **kwargs):
+    def __init__(self, option_strings, dest, error_message=None, nargs=0, **kwargs):
+        self.error_message = error_message
         super(DeprecatedAction, self).__init__(
             option_strings, dest, nargs=nargs, **kwargs
         )
 
     def __call__(self, parser, namespace, values, option_string=None):
+        if self.error_message is not None:
+            parser.error(self.error_message)
         print_deprecated_warning(
             f"The command line argument '{option_string}' is deprecated and will be removed in future versions."
         )

python/sglang/srt/environ.py

@@ -939,7 +939,6 @@ class Envs:
     SGLANG_LOG_EXPERT_LOCATION_METADATA = EnvBool(False)
     SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp")
     SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL = EnvInt(0)
-    SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC = EnvBool(False)
     # Chunk size for the rebalance expert-weight P2P exchange; set
     # >= num_physical_experts to submit a single batch_isend_irecv.
     SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE = EnvIntWithAlias(

python/sglang/srt/eplb/expert_distribution.py

@@ -22,7 +22,17 @@ from collections import deque
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+)
 
 import einops
 import torch
@@ -47,16 +57,22 @@ logger = logging.getLogger(__name__)
 # --------------------------------------- Entrypoint -----------------------------------------
 
 _OutputMode = Literal["file", "object"]
+EPLB_BALANCEDNESS_WINDOW_SIZES = (10, 100, 1000)
 
 
 @dataclass
 class ExpertDistributionMetrics:
+    forward_pass_id: int
     eplb_balancedness: torch.Tensor
+    gpu_physical_count_sum: Optional[torch.Tensor]
+    reset_server_log_history: bool
 
     def map_device_tensors(self, fn):
         # Device-tensor fields only; caller injects the copy+safety primitive
         # (see GenerationBatchResult.copy_to_cpu).
         self.eplb_balancedness = fn(self.eplb_balancedness)
+        if self.gpu_physical_count_sum is not None:
+            self.gpu_physical_count_sum = fn(self.gpu_physical_count_sum)
 
 
 class ExpertDistributionRecorder(ABC):
@@ -159,9 +175,10 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder):
             for k in self._accumulator.get_single_pass_gatherer_keys()
         }
 
-        if server_args.enable_expert_distribution_metrics:
+        if server_args.should_report_expert_balancedness():
             logger.info(
-                "ExpertDistributionRecorder auto start record since enable_expert_distribution_metrics"
+                "ExpertDistributionRecorder auto start record since "
+                f"expert_balancedness_report_mode={server_args.expert_balancedness_report_mode}"
             )
             self.start_record()
 
@@ -699,11 +716,12 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self._enable = self._server_args.enable_expert_distribution_metrics
+        self._enable = self._server_args.should_report_expert_balancedness()
 
         if self._enable:
-            self.window_sizes = [10, 100, 1000]
+            self.window_sizes = EPLB_BALANCEDNESS_WINDOW_SIZES
             self._history = _DequeCollection(maxlens=self.window_sizes)
+            self._reset_server_log_history = True
             self._rank = torch.distributed.get_rank()
             expert_dispatch_cls = resolve_collector_class(
                 self._server_args,
@@ -732,6 +750,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
         super().reset()
         if self._enable:
             self._history.clear()
+            self._reset_server_log_history = True
 
     def _append_utilization_rate(
         self,
@@ -757,27 +776,22 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
             should_track_history = not math.isclose(
                 self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0
             )
-            if envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get():
-                outputs["metrics"] = ExpertDistributionMetrics(
-                    eplb_balancedness=utilization_rate_gpu,
-                )
-                if should_track_history:
-                    self._history.append(utilization_rate_gpu.item())
-            else:
-                # TODO maybe refactor this part to also avoid a `.item()` gpu->cpu sync
-                utilization_rate_cpu = utilization_rate_gpu.item()
-                self._history.append(utilization_rate_cpu)
 
-                gpu_physical_count_sum = gpu_physical_count.sum().item()
+            should_log = (
+                self._server_args.should_log_expert_balancedness_to_server_log()
+            )
+            outputs["metrics"] = ExpertDistributionMetrics(
+                forward_pass_id=forward_pass_id,
+                eplb_balancedness=utilization_rate_gpu,
+                gpu_physical_count_sum=(
+                    gpu_physical_count.sum() if should_log else None
+                ),
+                reset_server_log_history=self._reset_server_log_history,
+            )
+            self._reset_server_log_history = False
 
-                logger.info(
+            if should_track_history:
-                    f"[Expert Balancedness] "
+                self._history.append(utilization_rate_gpu.item())
-                    f"forward_pass_id={forward_pass_id} "
-                    f"current_pass_balancedness={utilization_rate_cpu:.03f} "
-                    f"{''.join(f'last_{size}_average_balancedness={value:.03f} ' for size, value in self._history.mean().items())} "
-                    f"gpu_physical_count_sum={gpu_physical_count_sum}"
-                    # f"current_pass_per_layer={[round(x, 2) for x in utilization_rate_tensor.cpu().tolist()]}"
-                )
 
     # TODO refactor
     def _handle_metric_eplb_heatmap(self, gpu_physical_count: torch.Tensor):
@@ -804,7 +818,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
 
 
 class _DequeCollection:
-    def __init__(self, maxlens: List[int]):
+    def __init__(self, maxlens: Sequence[int]):
         self._dequeues = [deque(maxlen=maxlen) for maxlen in maxlens]
 
     def append(self, value):

python/sglang/srt/managers/scheduler_components/metrics_reporter.py

@@ -5,12 +5,13 @@ import logging
 import math
 import tempfile
 import time
-from collections import defaultdict
+from collections import defaultdict, deque
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_distribution import EPLB_BALANCEDNESS_WINDOW_SIZES
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.observability.metrics_collector import (
@@ -112,6 +113,11 @@ class SchedulerMetricsReporter:
         )
         self._init_metrics(self.tp_rank, self.pp_rank, self.dp_rank)
         self._install_device_timer_on_runners()
+        # Keep log history after the existing async result copy so reporting does
+        # not synchronize the model stream once per generated token.
+        self._eplb_balancedness_history = [
+            deque(maxlen=window_size) for window_size in EPLB_BALANCEDNESS_WINDOW_SIZES
+        ]
 
     def _init_metrics(
         self,
@@ -951,16 +957,45 @@ class SchedulerMetricsReporter:
         batch: ScheduleBatch,
         result: Union[GenerationBatchResult, EmbeddingBatchResult],
     ):
-        if not self.enable_metrics:
-            return
         if not isinstance(result, GenerationBatchResult):
             return
 
         if (m := result.expert_distribution_metrics) is not None:
-            self.metrics_collector.increment_eplb_balancedness(
+            balancedness = m.eplb_balancedness.item()
-                forward_mode=batch.forward_mode.name.lower(),
+
-                balancedness=m.eplb_balancedness.item(),
+            if (
-            )
+                self.scheduler.server_args.should_log_expert_balancedness_to_server_log()
+            ):
+                if m.reset_server_log_history:
+                    for history in self._eplb_balancedness_history:
+                        history.clear()
+                for history in self._eplb_balancedness_history:
+                    history.append(balancedness)
+                balancedness_history_means = {
+                    history.maxlen: sum(history) / len(history)
+                    for history in self._eplb_balancedness_history
+                    if len(history) > 0
+                }
+                assert m.gpu_physical_count_sum is not None
+                gpu_physical_count_sum = m.gpu_physical_count_sum.item()
+
+                logger.info(
+                    f"[Expert Balancedness] "
+                    f"forward_pass_id={m.forward_pass_id} "
+                    f"current_pass_balancedness={balancedness:.03f} "
+                    f"{''.join(f'last_{size}_average_balancedness={value:.03f} ' for size, value in balancedness_history_means.items())} "
+                    f"gpu_physical_count_sum={gpu_physical_count_sum}"
+                )
+
+            if (
+                self.enable_metrics
+                and self.scheduler.server_args.should_export_expert_balancedness_to_prometheus()
+            ):
+                assert self.metrics_collector is not None
+                self.metrics_collector.increment_eplb_balancedness(
+                    forward_mode=batch.forward_mode.name.lower(),
+                    balancedness=balancedness,
+                )
 
     def _emit_forward_pass_metrics(
         self,

python/sglang/srt/observability/metrics_collector.py

@@ -24,7 +24,6 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
-from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.observability.utils import exponential_buckets, generate_buckets
 from sglang.srt.server_args import ServerArgs
@@ -240,10 +239,10 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
     def __init__(
         self,
         labels: Dict[str, str],
+        server_args: ServerArgs,
         enable_lora: bool = False,
         enable_hierarchical_cache: bool = False,
         enable_streaming_session: bool = False,
-        server_args: Optional[ServerArgs] = None,
     ) -> None:
         # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
         from prometheus_client import Counter as _PromCounter
@@ -874,7 +873,8 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
         # =================================================================
         if (
             labels["moe_ep_rank"] == 0
-        ) and envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get():
+            and server_args.should_export_expert_balancedness_to_prometheus()
+        ):
             self.eplb_balancedness = Summary(
                 name="sglang:eplb_balancedness",
                 documentation="Balancedness of MoE in expert parallelism.",

python/sglang/srt/server_args.py

@@ -2409,9 +2409,11 @@ class ServerArgs:
         "Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer.",
         NS("exec.moe"),
     ] = None
-    enable_expert_distribution_metrics: A[
+    expert_balancedness_report_mode: A[
-        bool, "Enable logging metrics for expert balancedness", NS("exec.moe")
+        Literal["off", "server_log", "prometheus", "both"],
-    ] = False
+        "Where to report expert balancedness. Options: off, server_log, prometheus, both.",
+        NS("exec.moe"),
+    ] = "off"
     deepep_config: A[
         Optional[str],
         "Tuned DeepEP config suitable for your own cluster. It can be either a string with JSON content or a file path.",
@@ -7198,7 +7200,14 @@ class ServerArgs:
         # ===== END TO BE REFACTORED ====
 
     def _handle_expert_distribution_metrics(self):
-        if self.enable_expert_distribution_metrics and (
+        if "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC" in os.environ:
+            raise ValueError(
+                "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC is no longer supported. Use "
+                "--expert-balancedness-report-mode with one of: off, server_log, "
+                "prometheus, both."
+            )
+
+        if self.should_report_expert_balancedness() and (
             self.expert_distribution_recorder_mode is None
         ):
             self.expert_distribution_recorder_mode = "stat"
@@ -8620,6 +8629,19 @@ class ServerArgs:
         )
 
         # --- Deprecated argument registrations ---
+        parser.add_argument(
+            "--enable-expert-distribution-metrics",
+            action=DeprecatedAction,
+            error_message=(
+                "--enable-expert-distribution-metrics is no longer supported. Use "
+                "--expert-balancedness-report-mode with one of: off, server_log, "
+                "prometheus, both."
+            ),
+            help=(
+                "Removed. Use --expert-balancedness-report-mode with one of: "
+                "off, server_log, prometheus, both."
+            ),
+        )
         parser.add_argument(
             "--stream-output",
             action=DeprecatedStoreTrueAction,
@@ -9593,6 +9615,15 @@ class ServerArgs:
             "dp_size": self.dp_size,
         }
 
+    def should_report_expert_balancedness(self) -> bool:
+        return self.expert_balancedness_report_mode != "off"
+
+    def should_log_expert_balancedness_to_server_log(self) -> bool:
+        return self.expert_balancedness_report_mode in ("server_log", "both")
+
+    def should_export_expert_balancedness_to_prometheus(self) -> bool:
+        return self.expert_balancedness_report_mode in ("prometheus", "both")
+
 
 def m3_fp8_attn_gemm_enabled(args) -> bool:
     """Whether MiniMax-M3 attention GEMMs run in fp8 (no opt-in flag; active

test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py

@@ -51,7 +51,8 @@ class TestEplbMinRebalancingUtilizationThresholdBase(CustomTestCase):
         50,
         "--expert-distribution-recorder-buffer-size",
         50,
-        "--enable-expert-distribution-metrics",
+        "--expert-balancedness-report-mode",
+        "server_log",
         "--eplb-rebalance-layers-per-chunk",
         "1",
     ]