执行摘要

• 一句话：修复 EPLB 平衡度统计轴错误并补回归测试
• 推荐动作：值得快速精读。该 PR 展示了“指标统计轴选错导致误报”的典型修复路径：抽取纯函数、修正 dtype/维度归约、并用一个构造极端的回归测试锁定行为。对从事分布式训练推理或负载均衡相关工作的工程师，这个测试用例是很好的维度错误捕捉范例。

功能与动机

EPLB 的 step() 会在固定间隔输出每个模型的 avg_tokens、max_tokens 与 balancedness，用于观测专家并行负载。PR body 明确指出当前归约使用 layer 轴，因此当每个 EP rank 在所有层收到全部 token 时，会报告完美平衡，掩盖真实负载倾斜。修复目标是让指标真实反映“每个 rank 的负载差异”，而不是“每层平均负载是否一致”。

实现拆解

1. 提取统计辅助函数：在 vllm/distributed/eplb/eplb_state.py 中新增模块级函数 _compute_eplb_load_stats(num_tokens_per_rank)，接收形状为 (num_moe_layers, num_ranks) 的张量，统一返回 (avg_tokens, max_tokens) 两个标量张量。
2. 修正聚合轴：旧代码 mean(dim=0).sum(dim=0) 与 max(dim=0).values.sum(dim=0) 沿层轴（dim=0）归约；新代码改为 mean(dim=1).sum() 与 max(dim=1).values.sum()，先沿 rank 轴计算每层的均值/最大值，再对层求和。这样 avg 等价于“总 token 数 / rank 数”，max 等价于“各层最大 rank 负载之和”，能真实反映各 rank 的负载差异。
3. 更新调用点：EplbState.step() 中原来的四行内联统计替换为对 _compute_eplb_load_stats 的一次调用，其余日志输出与 balancedness = avg/max 逻辑保持不变。
4. 补回归测试：在 tests/distributed/test_eplb_algo.py 中新增 test_eplb_load_stats_reduce_across_ranks，构造两层各 100 个 token 全部路由到 rank 0（共 4 个 rank）的场景，断言修复后 avg_tokens == 50、max_tokens == 200、balancedness == 0.25；同时调整 import 以引用辅助函数。
5. 测试与 CI 配套：本地执行 pytest -q tests/distributed/test_eplb_algo.py 得到 15 passed, 1 skipped。PR 触发 Buildkite CI 后有一次失败，作者声明为基础设施间歇性问题；三名维护者批准后合并。

关键文件：

• vllm/distributed/eplb/eplb_state.py（模块 负载均衡；类别 source；类型 core-logic；符号 _compute_eplb_load_stats）: 核心修复文件：新增 _compute_eplb_load_stats 辅助函数并将 step() 中的统计逻辑归约轴从层维度（dim=0）改为 rank 维度（dim=1），是针对 EPLB 指标误报的关键逻辑修正。
• tests/distributed/test_eplb_algo.py（模块 负载均衡；类别 test；类型 test-coverage；符号 test_eplb_load_stats_reduce_across_ranks）: 新增 test_eplb_load_stats_reduce_across_ranks 回归用例，构造每层 token 全部命中同一 rank 的极端场景，锁定修复后的正确统计值，防止维度归约问题再次回归。

关键符号：_compute_eplb_load_stats, test_eplb_load_stats_reduce_across_ranks

关键源码片段

vllm/distributed/eplb/eplb_state.py

核心修复文件：新增 _compute_eplb_load_stats 辅助函数并将 step() 中的统计逻辑归约轴从层维度（dim=0）改为 rank 维度（dim=1），是针对 EPLB 指标误报的关键逻辑修正。

def _compute_eplb_load_stats(
    num_tokens_per_rank: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    # num_tokens_per_rank: (num_moe_layers, num_ranks)
    # 正确聚合方式：先沿 rank 维度（dim=1）计算均值 / 最大值，再对层求和。
    # 旧实现沿 layer 轴（dim=0）归约，当每层 token 全路由到同一 rank 时
    # 会误报 avg=100、max=100、balancedness=1.0 的完美平衡假象。
    avg_tokens = num_tokens_per_rank.mean(dim=1).sum()
    max_tokens = num_tokens_per_rank.max(dim=1).values.sum()
    return avg_tokens, max_tokens

# EplbState.step() 中的调用点，替换原来四行内联统计：
avg_tokens_tensor, max_tokens_tensor = _compute_eplb_load_stats(
    num_tokens_per_rank
)
​
# 后续日志逻辑保持不变：
# balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0

tests/distributed/test_eplb_algo.py

新增 test_eplb_load_stats_reduce_across_ranks 回归用例，构造每层 token 全部命中同一 rank 的极端场景，锁定修复后的正确统计值，防止维度归约问题再次回归。

def test_eplb_load_stats_reduce_across_ranks():
    # 两层各 100 个 token 全部路由到 rank 0（共 4 个 EP rank）
    num_tokens_per_rank = torch.tensor(
        [
            [100, 0, 0, 0],
            [100, 0, 0, 0],
        ],
        dtype=torch.float32,
    )
​
    avg_tokens, max_tokens = _compute_eplb_load_stats(num_tokens_per_rank)
​
    # 修复前误报 avg=100 / max=100（balancedness=1.0），
    # 修复后应正确反映 rank 0 负载集中：avg=50、max=200、balancedness=0.25
    assert avg_tokens.item() == 50
    assert max_tokens.item() == 200
    assert (avg_tokens / max_tokens).item() == 0.25

评论区精华

审核人 ilmarkov 在批准意见中标注“Basically revival of #39178 without additional changes”，确认这是对既有修复方案的复用，没有新增设计分歧。作者 jdebache 在 CI 失败后评论指出失败与本改动无关，疑似基础设施间歇性问题，并请求维护者同意合并；随后 tlrmchlsmth、SageMoore 等批准，PR 合入 main。claude[bot] 的自动审核因 fork 来源被跳过，未产生实质技术讨论。

• CI 失败是否为基础设施问题 (question): 维护者 tlrmchlsmth 触发 CI 并批准合并，PR 最终合入 main，说明社区认可该失败与改动无关。
• 与 #39178 的关系 (other): 确认本 PR 是对既有修复方案 #39178 的复兴，无新增设计争议。

风险与影响

• 风险：影响范围仅限 EPLB step() 中 log_stats 分支的指标计算与日志输出，不改变重排算法、通信协议或 KV 缓存行为。num_tokens_per_rank 的形状 (num_moe_layers, num_ranks) 未变化，dim=1 始终是 EP rank 维度，聚合逻辑稳定；新增回归测试覆盖了最关键的维度陷阱场景。主要风险是“指标口径变化”：依赖 EPLB 日志数值的监控面板或自动化脚本需要按新统计口径调整预期值，短时间内可能对长期观察 balancedness=1.0 的团队造成困惑。
• 影响：直接影响面为使用 EPLB 的多 GPU 分布式 MoE 推理场景，用户会在日志中看到更真实的 avg_tokens、max_tokens 与 balancedness 数值，不再出现单 rank 满载却仍显示完美平衡的假象。对系统运行本身无行为影响，对依赖日志做可视化或告警的运维侧有指标口径迁移成本。团队内部收获一个清晰的维度归约回归测试用例，可作为未来张量统计逻辑的参考模板。
• 风险标记：统计路径变更, 有回归测试覆盖, 不影响重排决策, 日志指标口径变化

关联脉络

• PR #39178 （讨论中提及的 EPLB 平衡度统计修复）: 审核人 ilmarkov 明确表示本 PR 是 #39178 的复兴版本，本次补充了回归测试，两者同属 EPLB 指标统计修复系列。