karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Misc] Add EPLB balancedness to Prometheus metrics

Browse Source 
Expose EPLB load balance metrics via Prometheus with zero inter-rank
synchronization. The driver rank's local expert_load_pass tensor is
reshaped to compute per-layer balancedness (avg/max token load across
EP ranks) using exact quantiles.

New metrics:
- vllm:eplb_balancedness_min  (worst MoE layer)
- vllm:eplb_balancedness_p50  (median across layers)
- vllm:eplb_balancedness_p90  (90th percentile)
- vllm:eplb_balancedness_avg  (mean across layers)
- vllm:eplb_rearrangements_total (rearrangement counter)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
Tyler Michael Smith
2026-03-07 00:39:50 -05:00
co-authored by Claude Opus 4.6
parent 0ccb2ef093
commit 53c8d72f71
8 changed files with 269 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/distributed/eplb/eplb_state.py
+85 -66
View File
@@ -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.
vllm/distributed/eplb/metrics.py
+130
View File
@@ -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
)
vllm/v1/core/sched/scheduler.py
+9 -2
View File
@@ -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:
vllm/v1/metrics/loggers.py
+8
View File
@@ -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
vllm/v1/metrics/ray_wrappers.py
+13
View File
@@ -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():
vllm/v1/metrics/stats.py
+14
View File
@@ -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:
vllm/v1/outputs.py
+5
View File
@@ -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):
vllm/v1/worker/gpu_model_runner.py
+5
View File
@@ -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: