fix and test EPLB balancedness calculation (#51813)
Signed-off-by: jdebache <jdebache@nvidia.com>
Julien Debache committed
Aug 13, 2026 at 05:27 UTC
6e502b62cb6affb3710b50d6f0f9c74eb6accabd
2 files changed
+31
-3
tests/distributed/test_eplb_algo.py
+20
-1
@@ -5,10 +5,29 @@ import numpy as np
5
import pytest
6
import torch
7
8
-from vllm.distributed.eplb.eplb_state import compute_logical_maps
8
+from vllm.distributed.eplb.eplb_state import (
9
+ _compute_eplb_load_stats, # pyright: ignore[reportPrivateUsage]
10
+ compute_logical_maps,
11
+)
12
from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
13
14
15
+def test_eplb_load_stats_reduce_across_ranks():
16
+ num_tokens_per_rank = torch.tensor(
17
+ [
18
+ [100, 0, 0, 0],
19
+ [100, 0, 0, 0],
20
+ ],
21
+ dtype=torch.float32,
22
+ )
23
+
24
+ avg_tokens, max_tokens = _compute_eplb_load_stats(num_tokens_per_rank)
25
+
26
+ assert avg_tokens.item() == 50
27
+ assert max_tokens.item() == 200
28
+ assert (avg_tokens / max_tokens).item() == 0.25
29
+
30
+
31
def test_basic_rebalance():
32
"""Test basic rebalancing functionality"""
33
# Example from https://github.com/deepseek-ai/eplb
vllm/distributed/eplb/eplb_state.py
+11
-2
@@ -61,6 +61,14 @@ from .rebalance_execute import (
61
logger = init_logger(__name__)
62
63
64
+def _compute_eplb_load_stats(
65
+ num_tokens_per_rank: torch.Tensor,
66
+) -> tuple[torch.Tensor, torch.Tensor]:
67
+ avg_tokens = num_tokens_per_rank.mean(dim=1).sum()
68
+ max_tokens = num_tokens_per_rank.max(dim=1).values.sum()
69
+ return avg_tokens, max_tokens
70
+
71
+
72
@dataclass
73
class EplbStats:
74
"""
@@ -584,8 +592,9 @@ class EplbState:
592
# Compute balancedness ratio:
593
# for each layer:
594
# (mean load across ranks) / (max load across ranks)
587
- avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
588
- max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
595
+ avg_tokens_tensor, max_tokens_tensor = _compute_eplb_load_stats(
596
+ num_tokens_per_rank
597
+ )
598
599
# Just to make type checker happy
600
tokens_tensors: list[float] = torch.stack(