diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 28cdf46dd3f..aed8f45644d 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -44,6 +44,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.v1.metrics.stats import EplbMetricsStats from .async_worker import start_async_worker from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy @@ -311,6 +312,20 @@ class EplbState: newly started EP ranks may not have physical experts mapped yet. """ + self.last_eplb_stats: EplbMetricsStats | None = None + """ + Most recent EPLB balancedness stats for Prometheus export. + Computed locally (no inter-rank sync) each step. + """ + self.rearrangements_since_last_report: int = 0 + """ + Number of rearrangements since the last stats were consumed. + """ + self.last_rearrangement_seconds: float = 0.0 + """ + Duration of the most recent rearrangement in seconds. + """ + if self.device.type == "cuda": self.cuda_device_index = self.device.index if self.cuda_device_index is None and torch.cuda.is_available(): @@ -606,20 +621,12 @@ class EplbState: # per-rank token counts for the worst layer worst_layer = int(per_layer_balance.argmin().item()) best_layer = int(per_layer_balance.argmax().item()) - worst_balance = float( - per_layer_balance[worst_layer].item() - ) - best_balance = float( - per_layer_balance[best_layer].item() - ) + worst_balance = float(per_layer_balance[worst_layer].item()) + best_balance = float(per_layer_balance[best_layer].item()) worst_layer_ranks = num_tokens_per_rank[worst_layer] - worst_min_rank = int( - worst_layer_ranks.argmin().item() - ) - worst_max_rank = int( - worst_layer_ranks.argmax().item() - ) + worst_min_rank = int(worst_layer_ranks.argmin().item()) + worst_max_rank = int(worst_layer_ranks.argmax().item()) logger.info( "EPLB balance breakdown: " @@ -630,26 +637,17 @@ class EplbState: worst_layer, worst_balance, worst_min_rank, - float( - worst_layer_ranks[worst_min_rank].item() - ), + float(worst_layer_ranks[worst_min_rank].item()), worst_max_rank, - float( - worst_layer_ranks[worst_max_rank].item() - ), + float(worst_layer_ranks[worst_max_rank].item()), best_layer, best_balance, num_tokens_per_rank.shape[0], ) # Log replica distribution for debug - replica_count = ( - eplb_model_state.logical_replica_count - ) - if ( - replica_count is not None - and replica_count.numel() > 0 - ): + replica_count = eplb_model_state.logical_replica_count + if replica_count is not None and replica_count.numel() > 0: rc_float = replica_count.float() logger.debug( "EPLB replica stats (layer avg): " @@ -658,15 +656,16 @@ class EplbState: float(rc_float.min().item()), float(rc_float.max().item()), float(rc_float.mean().item()), - int( - (rc_float > 1) - .any(dim=0) - .sum() - .item() - ), + int((rc_float > 1).any(dim=0).sum().item()), replica_count.shape[-1], ) + # Compute local balancedness stats for Prometheus (no inter-rank sync). + # Uses only the driver rank's expert_load_pass which records routing + # decisions for all physical experts across all EP ranks. + if not is_dummy: + self._compute_local_balancedness_stats() + # Update the expert load sliding window if not is_dummy: for eplb_model_state in self.model_states.values(): @@ -708,6 +707,7 @@ class EplbState: return self.expert_rearrangement_step = 0 self.rearrange() + self.rearrangements_since_last_report += 1 def rearrange( self, @@ -746,11 +746,7 @@ class EplbState: if is_main_rank: # Log window utilization diagnostics nonzero_slots = sum( - int( - (ms.expert_load_window.sum(dim=(1, 2)) > 0) - .sum() - .item() - ) + int((ms.expert_load_window.sum(dim=(1, 2)) > 0).sum().item()) for ms in self.model_states.values() ) logger.info( @@ -764,10 +760,7 @@ class EplbState: nonzero_slots, self.expert_load_window_size, ) - if ( - self.expert_load_window_size - > self.expert_rearrangement_step_interval - ): + if self.expert_load_window_size > self.expert_rearrangement_step_interval: logger.warning( "EPLB: window_size (%d) > step_interval (%d). " "Stale window entries from before the last " @@ -859,17 +852,12 @@ class EplbState: ) # Top-5 hottest experts (averaged across layers) avg_load = load.mean(dim=0) - top5_vals, top5_ids = avg_load.topk( - min(5, avg_load.shape[0]) - ) + top5_vals, top5_ids = avg_load.topk(min(5, avg_load.shape[0])) logger.info( - "EPLB top-5 hottest logical experts " - "(avg across layers): %s", + "EPLB top-5 hottest logical experts (avg across layers): %s", ", ".join( f"e{int(eid)}={float(val):.0f}" - for eid, val in zip( - top5_ids.tolist(), top5_vals.tolist() - ) + for eid, val in zip(top5_ids.tolist(), top5_vals.tolist()) ), ) @@ -891,12 +879,8 @@ class EplbState: if is_main_rank and not is_profile: # Log what the algorithm decided old_p2l = eplb_model_state.physical_to_logical_map - new_p2l = new_physical_to_logical_map.to( - old_p2l.device - ) - changed_slots = int( - (old_p2l != new_p2l).sum().item() - ) + new_p2l = new_physical_to_logical_map.to(old_p2l.device) + changed_slots = int((old_p2l != new_p2l).sum().item()) total_slots = old_p2l.numel() rc = new_logical_replica_count.float() logger.info( @@ -914,9 +898,7 @@ class EplbState: # Simulate new per-rank load to preview # balancedness - new_rc = new_logical_replica_count.to( - load.device - ).float() + new_rc = new_logical_replica_count.to(load.device).float() per_expert_load = load / new_rc.clamp(min=1) phys_load = per_expert_load.gather( dim=-1, @@ -936,16 +918,10 @@ class EplbState: "EPLB predicted post-rearrange " "balancedness: mean=%.4f, " "min=%.4f (layer %d), max=%.4f", - float( - predicted_balance.mean().item() - ), + float(predicted_balance.mean().item()), float(predicted_balance.min().item()), - int( - predicted_balance.argmin().item() - ), - float( - predicted_balance.max().item() - ), + int(predicted_balance.argmin().item()), + float(predicted_balance.max().item()), ) # Update expert weights @@ -998,6 +974,8 @@ class EplbState: end_event.record() end_event.synchronize() gpu_elapsed = start_event.elapsed_time(end_event) / 1000.0 + if not is_profile: + self.last_rearrangement_seconds = gpu_elapsed logger.info( "Rearranged experts %s in %.2f s.", " (profile) " if is_profile else " ", @@ -1220,6 +1198,47 @@ class EplbState: offset += shape[0] return all_reduce_list + def _compute_local_balancedness_stats(self) -> None: + """Compute per-layer balancedness from this rank's expert_load_pass. + + No inter-rank communication. The expert_load_pass tensor records + routing decisions for all physical experts (partitioned across EP + ranks), so reshaping by rank gives per-rank token loads. + """ + ep_group = get_ep_group().device_group + num_ranks = ep_group.size() + + # Use the first model's expert_load_pass (main model, not drafter) + eplb_model_state = next(iter(self.model_states.values())) + expert_load = eplb_model_state.expert_load_pass + # expert_load: (num_moe_layers, num_physical_experts) + num_layers = expert_load.shape[0] + + # Reshape to (num_moe_layers, num_ranks, experts_per_rank) + # and sum per-rank token loads + per_rank = expert_load.reshape(num_layers, num_ranks, -1).sum(dim=2).float() + avg_per_layer = per_rank.mean(dim=1) # (num_layers,) + max_per_layer = per_rank.max(dim=1).values # (num_layers,) + layer_balancedness = torch.where( + max_per_layer > 0, + avg_per_layer / max_per_layer, + torch.ones_like(max_per_layer), + ) + + # Compute quantile stats from the per-layer balancedness values + bal = layer_balancedness.cpu().numpy() + rearrangements = self.rearrangements_since_last_report + self.rearrangements_since_last_report = 0 + + self.last_eplb_stats = EplbMetricsStats( + min_balancedness=float(bal.min()), + p50_balancedness=float(np.median(bal)), + p90_balancedness=float(np.percentile(bal, 10)), + avg_balancedness=float(bal.mean()), + rearrangements=rearrangements, + last_rearrangement_seconds=self.last_rearrangement_seconds, + ) + def _sync_load_pass(self) -> list[torch.Tensor]: """ Sync the expert load pass across all ranks for log stats. diff --git a/vllm/distributed/eplb/metrics.py b/vllm/distributed/eplb/metrics.py new file mode 100644 index 00000000000..59ed6f790bb --- /dev/null +++ b/vllm/distributed/eplb/metrics.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import prometheus_client + +from vllm.config import ParallelConfig +from vllm.v1.metrics.stats import EplbMetricsStats + + +def _make_per_engine( + metric: prometheus_client.Gauge | prometheus_client.Counter, + per_engine_labelvalues: dict[int, list[object]], +) -> dict[int, prometheus_client.Gauge | prometheus_client.Counter]: + return { + idx: metric.labels(*labelvalues) + for idx, labelvalues in per_engine_labelvalues.items() + } + + +class EplbProm: + """Record EPLB balancedness metrics in Prometheus. + + Balancedness is defined per MoE layer as: + avg(tokens_per_rank) / max(tokens_per_rank) + + A value of 1.0 means perfectly balanced load across EP ranks. + + Example PromQL queries: + + # Alert on worst-layer imbalance + vllm:eplb_balancedness_min < 0.7 + + # Median vs worst-layer gap + vllm:eplb_balancedness_p50 - vllm:eplb_balancedness_min + """ + + _gauge_cls = prometheus_client.Gauge + _counter_cls = prometheus_client.Counter + + def __init__( + self, + parallel_config: ParallelConfig, + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ): + self.enabled = parallel_config.enable_eplb + if not self.enabled: + return + + gauge_min = self._gauge_cls( + name="vllm:eplb_balancedness_min", + documentation=( + "Minimum per-layer EPLB balancedness " + "(worst MoE layer, avg/max token load across EP ranks)." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.gauge_balancedness_min = _make_per_engine( + gauge_min, per_engine_labelvalues + ) + + gauge_p50 = self._gauge_cls( + name="vllm:eplb_balancedness_p50", + documentation=("Median per-layer EPLB balancedness across MoE layers."), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.gauge_balancedness_p50 = _make_per_engine( + gauge_p50, per_engine_labelvalues + ) + + gauge_p90 = self._gauge_cls( + name="vllm:eplb_balancedness_p90", + documentation=( + "10th-percentile EPLB balancedness " + "(90%% of MoE layers are at least this balanced)." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.gauge_balancedness_p90 = _make_per_engine( + gauge_p90, per_engine_labelvalues + ) + + gauge_avg = self._gauge_cls( + name="vllm:eplb_balancedness_avg", + documentation=("Mean EPLB balancedness across MoE layers."), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.gauge_balancedness_avg = _make_per_engine( + gauge_avg, per_engine_labelvalues + ) + + counter_rearrangements = self._counter_cls( + name="vllm:eplb_rearrangements_total", + documentation="Total number of EPLB expert rearrangements.", + labelnames=labelnames, + ) + self.counter_rearrangements = _make_per_engine( + counter_rearrangements, per_engine_labelvalues + ) + + gauge_rearrangement_seconds = self._gauge_cls( + name="vllm:eplb_rearrangement_seconds", + documentation=( + "Duration of the most recent EPLB expert rearrangement in seconds." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.gauge_rearrangement_seconds = _make_per_engine( + gauge_rearrangement_seconds, per_engine_labelvalues + ) + + def observe(self, eplb_stats: EplbMetricsStats, engine_idx: int = 0): + if not self.enabled: + return + self.gauge_balancedness_min[engine_idx].set(eplb_stats.min_balancedness) + self.gauge_balancedness_p50[engine_idx].set(eplb_stats.p50_balancedness) + self.gauge_balancedness_p90[engine_idx].set(eplb_stats.p90_balancedness) + self.gauge_balancedness_avg[engine_idx].set(eplb_stats.avg_balancedness) + if eplb_stats.rearrangements > 0: + self.counter_rearrangements[engine_idx].inc(eplb_stats.rearrangements) + self.gauge_rearrangement_seconds[engine_idx].set( + eplb_stats.last_rearrangement_seconds + ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index cb99de93b6f..2774559540d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -50,7 +50,7 @@ from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.metrics.perf import ModelMetrics, PerfStats -from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats +from vllm.v1.metrics.stats import EplbMetricsStats, PrefixCacheStats, SchedulerStats from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats @@ -1287,6 +1287,7 @@ class Scheduler(SchedulerInterface): num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output cudagraph_stats = model_runner_output.cudagraph_stats + eplb_stats = model_runner_output.eplb_stats perf_stats: PerfStats | None = None if self.perf_metrics and self.perf_metrics.is_enabled(): @@ -1517,7 +1518,11 @@ class Scheduler(SchedulerInterface): if ( stats := self.make_stats( - spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats + spec_decoding_stats, + kv_connector_stats, + cudagraph_stats, + perf_stats, + eplb_stats, ) ) is not None: # Return stats to only one of the front-ends. @@ -1876,6 +1881,7 @@ class Scheduler(SchedulerInterface): kv_connector_stats: KVConnectorStats | None = None, cudagraph_stats: CUDAGraphStat | None = None, perf_stats: PerfStats | None = None, + eplb_stats: EplbMetricsStats | None = None, ) -> SchedulerStats | None: if not self.log_stats: return None @@ -1906,6 +1912,7 @@ class Scheduler(SchedulerInterface): kv_connector_stats=connector_stats_payload, cudagraph_stats=cudagraph_stats, perf_stats=perf_stats, + eplb_stats=eplb_stats, ) def _get_encoder_cache_usage(self) -> float: diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index f20d7854224..bdbcc63459f 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -12,6 +12,7 @@ from prometheus_client import Counter, Gauge, Histogram import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphLogging from vllm.config import SupportsMetricsInfo, VllmConfig +from vllm.distributed.eplb.metrics import EplbProm from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( KVConnectorLogging, KVConnectorPrometheus, @@ -393,6 +394,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase): _spec_decoding_cls = SpecDecodingProm _kv_connector_cls = KVConnectorPrometheus _perf_metrics_cls = PerfMetricsProm + _eplb_cls = EplbProm def __init__( self, vllm_config: VllmConfig, engine_indexes: list[int] | None = None @@ -428,6 +430,9 @@ class PrometheusStatLogger(AggregateStatLoggerBase): self.perf_metrics_prom = self._perf_metrics_cls( vllm_config, labelnames, per_engine_labelvalues ) + self.eplb_prom = self._eplb_cls( + vllm_config.parallel_config, labelnames, per_engine_labelvalues + ) # # Scheduler state @@ -1072,6 +1077,9 @@ class PrometheusStatLogger(AggregateStatLoggerBase): if scheduler_stats.perf_stats is not None: self.perf_metrics_prom.observe(scheduler_stats.perf_stats, engine_idx) + if scheduler_stats.eplb_stats is not None: + self.eplb_prom.observe(scheduler_stats.eplb_stats, engine_idx) + if ( self.kv_cache_metrics_enabled and scheduler_stats.kv_cache_eviction_events diff --git a/vllm/v1/metrics/ray_wrappers.py b/vllm/v1/metrics/ray_wrappers.py index abc53f3802e..55bf829d839 100644 --- a/vllm/v1/metrics/ray_wrappers.py +++ b/vllm/v1/metrics/ray_wrappers.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +from vllm.distributed.eplb.metrics import EplbProm from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorPrometheus from vllm.v1.metrics.loggers import PrometheusStatLogger from vllm.v1.metrics.perf import PerfMetricsProm @@ -190,6 +191,17 @@ class RayPerfMetricsProm(PerfMetricsProm): _counter_cls = RayCounterWrapper +class RayEplbProm(EplbProm): + """ + RayEplbProm is used by RayMetrics to log Ray metrics. + Provides the same EPLB balancedness metrics as EplbProm + but uses Ray's util.metrics library. + """ + + _gauge_cls = RayGaugeWrapper + _counter_cls = RayCounterWrapper + + class RayPrometheusStatLogger(PrometheusStatLogger): """RayPrometheusStatLogger uses Ray metrics instead.""" @@ -199,6 +211,7 @@ class RayPrometheusStatLogger(PrometheusStatLogger): _spec_decoding_cls = RaySpecDecodingProm _kv_connector_cls = RayKVConnectorPrometheus _perf_metrics_cls = RayPerfMetricsProm + _eplb_cls = RayEplbProm @staticmethod def _unregister_vllm_metrics(): diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 4a1e8b6f35c..d1944f7f6a2 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -167,6 +167,18 @@ class KVCacheEvictionEvent: reuse_gaps_seconds: tuple[float, ...] +@dataclass +class EplbMetricsStats: + """EPLB balancedness stats computed per step for Prometheus export.""" + + min_balancedness: float + p50_balancedness: float + p90_balancedness: float + avg_balancedness: float + rearrangements: int + last_rearrangement_seconds: float + + @dataclass class SchedulerStats: """Stats associated with the scheduler.""" @@ -196,6 +208,8 @@ class SchedulerStats: perf_stats: PerfStats | None = None + eplb_stats: EplbMetricsStats | None = None + @dataclass class RequestStateStats: diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 22b06f0e2d9..509c0c49b18 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -15,9 +15,11 @@ from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: from vllm.distributed.kv_events import KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats + from vllm.v1.metrics.stats import EplbMetricsStats else: KVConnectorStats = object KVConnectorKVEvents = object + EplbMetricsStats = object class LogprobsLists(NamedTuple): @@ -247,6 +249,9 @@ class ModelRunnerOutput: # information related to cudagraph execution cudagraph_stats: CUDAGraphStat | None = None + # EPLB balancedness stats + eplb_stats: "EplbMetricsStats | None" = None + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index abeb1073512..4b968525b05 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3865,6 +3865,10 @@ class GPUModelRunner( else: logger.error("RoutedExpertsCapturer not initialized.") + eplb_stats = ( + self.eplb_state.last_eplb_stats if self.eplb_state is not None else None + ) + output = ModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -3877,6 +3881,7 @@ class GPUModelRunner( else None, num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, + eplb_stats=eplb_stats, ) if not self.use_async_scheduling: