DeepInfer 源码完全拆解（四）：MoE 专家并行与 EPLB——冗余专家的复制、放置与在线重排

为什么 MoE 推理必须并行

先交代 MoE 层的形状。DeepSeek 系的 MoE 层里，一个 token 经过 gate 打分后只被路由到 top-k 个专家（V3 从 256 个路由专家中选 top-8，V4.1-Flash 的具体配置以模型卡为准），每个专家是一个独立的 SwiGLU MLP。整个模型的参数量因此可以远大于「每 token 实际激活的参数」。DeepSeek-V3 的技术报告给出了这套结构的全景：注意力侧用 MLA，前馈侧用 DeepSeekMoE（少量共享专家 + 大量细粒度路由专家）。

（图源：DeepSeek-V3 Technical Report 图 3 basic_arch）

这个设计在显存上是省钱的，在并行上却是麻烦的：既然每次前向只用到少数专家，那就不必让每张卡都持有全部专家权重。专家并行（Expert Parallel, EP）的做法是把专家集合切成 ep_size 份，每张卡只持有 $N_{logical} / ep\_size$ 个专家；一次前向里，本卡先算出每个 token 该去哪些专家，再把 token 发给持有那些专家的卡（dispatch），算完再收回来（combine）。

这条链路的代价模型很直白：一次 MoE 层的时延约等于

其中 $T_{expert}$ 是本卡要算的专家前向时间，取决于发到本卡的 token 数；$T_{dispatch}$/$T_{combine}$ 是通信时间，取决于要跨卡搬运的 token 数。注意这两项都是按整批最慢的那张卡结算的：MoE 前向是一次集合通信，任何一张卡慢，其他卡都得等着。

于是问题的形状就出来了：

均衡的目标不是让平均负载低，而是让最大负载低。 这是 EPLB 全部算法设计的出发点——后面 policy.py 里的目标函数、_batch_greedy_indices 里的 scores、日志里的 rank_balance 指标，全都是围绕「峰值的均值」而不是「均值」定义的。

不均衡有多严重

如果路由是完全均匀的，每个专家拿到的 token 数都一样，那按 expert 均匀切分就完事了。真实情况不是这样：

• 专家热度天然不均匀。不同专家学到的功能不同，在具体流量分布下被选中的频率可能差好几倍。DeepSeek-V3 的技术报告里给出的路由负载统计显示，少数专家的 token 数是平均值的数倍。
• 流量随场景漂移。同一个部署上午在服务代码补全、下午在服务长文问答，热点专家集合会变。
• 长度分布放大不均衡。prefill 阶段一个请求就是几千 token，它的路由分布会整块地压到某些专家上；decode 阶段的步长小，抖动被平均掉。同一份权重放置对两个阶段的优劣可能完全相反。

如果听任不均衡，后果是显存和算力两头的浪费：热点专家所在的卡成为瓶颈，其余卡的算力闲置；同时为了兜住最坏情况，调度器必须按峰值预留 token 缓冲，进一步吃掉本可以给 KV 缓存的显存。DeepSeek 在线上系统里用三个负载均衡器（专家级、节点级、序列级）解决这件事的思路，站内已有专文拆过，可参考 DeepSeek-V3/R1 推理系统完全拆解：跨节点专家并行、双批重叠与负载均衡。

下面这张图是 DeepSeek-V3 技术报告给出的实测专家负载分布：横轴是专家编号，纵轴是相对负载（以平均值为 1），不同曲线对应不同层的不同专家组。可以看到有明显超过平均值的尖峰。

（图源：DeepSeek-V3 Technical Report 图 5 relative_expert_load_multi）

EPLB 的两段式答案

EPLB（Expert Parallel Load Balancer）是这套问题的标准答案，分两段：

1. 冗余副本（redundancy）：给热点专家多放几份副本，分散到不同的卡上，让热点专家的计算量被多张卡分摊。
2. 重排（placement / rearrangement）：把「哪个专家副本放在哪张卡」的映射定期重算，跟上流量漂移。

第一段改变的是副本数，第二段改变的是副本的位置。DeepInfer 的实现把这两段都放进了一个函数：rebalance_experts_global。而在这之前，先有一个取整问题必须解决——副本数不能随便选。

DeepSeek 官方开源的 EPLB 仓库里有一张说明图，把「复制 + 放置」这件事的画面表达得很清楚：同一层的专家里，热门的被复制成多份，再连同其他专家一起被摊到各张 GPU 上。

（图源：DeepSeek 官方 EPLB 仓库 example.png。该仓库的 rebalance_experts 提供 hierarchical 与 global 两种策略：前者先按专家组对齐节点、再在节点内复制，适配 prefill 阶段较小的 EP 规模；后者跨组全局复制，适配 decode 阶段较大的 EP 规模。DeepInfer 的 rebalance_experts_global 与后者同源。）

三个数字的定义

DeepInfer 在 eplb/policy.py 的模块 docstring 里把术语钉死：

1
术语：
2
- logical expert：模型定义的 expert，简写 log
3
- physical expert：GPU 上实际放置的 expert 副本，简写 phy，数量 = num_logical + num_redundant

再往下拆，一共有三个数字，moe.py 的模块 docstring 把第四个也补上了：

• logical = cfg.moe.n_routed_experts，即 gate 选择的范围（具体数值由模型的 config.json 给出，DeepInfer 从 HF 配置读取）。
• extra = cfg.moe.n_redundant_experts，由 model_loader._configure_moe_redundancy 写入。
• physical = logical + extra，实际分配的参数行数，必须能被 ep_size 整除。
• local = physical // ep_size，本卡持有的专家槽位数。

这四个数不是四种东西，而是同一件事在三个不同层面上的名字：模型定义层面（logical）、存储层面（physical）、设备层面（local）。搞混它们就会看不懂代码——因为 DeepInfer 的变量名里同时出现 physical_to_logical_map、logical_to_physical_map、num_local_physical_experts 三种命名，它们分别描述「物理槽位指向谁」「某个 logical expert 的副本都在哪些槽位」「本卡有几个槽位」。

用户给的是下限，引擎算的是取整结果

冗余专家数不是直接配置的。环境变量 DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS 默认 0，它的语义是下限，而不是最终值。真正落地的值由一个取整函数算出：

1
def resolve_num_extra_experts(
2
    num_routed_experts: int,
3
    min_num_extra_experts: int,
4
    ep_size: int,
5
) -> int:
6
    """取不小于下限、且使 num_physical_experts 能被 ep_size 整除的冗余 expert 数。"""
7
    padding = -(num_routed_experts + min_num_extra_experts) % ep_size
8
    return min_num_extra_experts + padding

这几行值得逐项拆开。Python 的取模对正数返回非负余数，所以 -(a) % m 等于 (m - a % m) % m，即「把 a 补到下一个 m 的倍数还差多少」。设 $N$ 为 routed expert 数、$E$ 为下限、$P$ 为 ep_size，公式是

用一组量级接近 V3 配置的参数算两遍就清楚了。先看不需补齐的情形：$N = 256$，下限 $E = 32$，$P = 32$，则 $N + E = 288$，$288 \bmod 32 = 0$，所以 padding = 0、extra = 32、physical = 288，每卡 9 个槽位。

再看一个不对齐的例子：$N = 256$、$E = 32$、$P = 64$。$N + E = 288$，$288 \bmod 64 = 32$，padding = 32，extra = 64，physical = 320，$320 / 64 = 5$ 个 local 槽位。用户要 32 个冗余，引擎给了他 64 个。 这不是 bug，是被迫的：如果只有 288 个 physical 槽位，288 不能被 64 整除，各卡拿到的槽位数就不一样，num_local_physical_experts 这个「每卡槽位数」的标量就不存在了——而整个代码库到处都在用它做整数除法索引（例如 physical_to_logical_map[:, rank * num_local : (rank+1) * num_local]）。整除约束是这里的硬约束，不是调优参数。

这个行为在装载侧有显式告警：

1
def resolve(moe_cfg: AttrDict, layout: str) -> int:
2
    num_extra_experts = resolve_num_extra_experts(moe_cfg.n_routed_experts, min_num_extra_experts, ep_size)
3
    if num_extra_experts != min_num_extra_experts:
4
        logger.warning(
5
            f"{layout} DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS={min_num_extra_experts} 已向上取整为 {num_extra_experts}，"
6
            f"以使 num_routed_experts={moe_cfg.n_routed_experts} + num_extra_experts 能被 ep_size={ep_size} 整除"
7
        )
8
    return num_extra_experts

两个 MoE layout 各自取整

_configure_moe_redundancy 里有一处容易漏掉的细节：

1
main_moe_cfg.n_redundant_experts = resolve(main_moe_cfg, "backbone")
2
dflash_cfg = gpt_cfg.get("dflash_cfg")
3
dflash_moe_override = dflash_cfg.get("moe") if dflash_cfg is not None else None
4
if dflash_moe_override is not None:
5
    dflash_moe_override.n_redundant_experts = resolve(main_moe_cfg + dflash_moe_override, "DFlash")

主干模型（backbone）和 DFlash（投机解码的草稿头，见第一篇讲的 MTP/DSpark 机制）各自做一次取整。DFlash 的配置是叠加在主干配置之上的：

1
def resolve(moe_cfg: AttrDict, layout: str) -> int:

这里的 moe_cfg 参数就是 main_moe_cfg + dflash_moe_override —— 这是一个配置字典相加的结果，也就是说 DFlash 的 MoE 布局可能与主干不同（不同的 routed expert 数），必须用它自己的数字去算整除。两个 layout 的 physical 数是独立的，后面的 EPLB 状态也因此按层分组管理（ExpertLoadBalancerManager 接收 groups: list[tuple[ExpertLoadBalancer, list[int]]]，把不同 balancer 映射到全局层号）。

另外注意 finalize_model_config 里这两行：

1
model_cfg.gpt.max_num_decoding_tokens = model_cfg.gpt.max_num_seqs * (1 + envs.DEEPINFER_USE_NEXTN_DECODE)
2
model_cfg.gpt.max_num_prefilling_tokens = model_cfg.gpt.max_num_batched_tokens // PG.context_parallel_group().size()

专家布局的取整发生在 model_loader.finalize_model_config 阶段，也就是在 meta 设备上构建模型容器之前。为什么必须在构建模型之前？因为 physical expert 数决定了 FusedExperts 里 ParallelRoutedLinear 的参数形状：

1
self.n_routed_experts = self.n_logical_routed_experts + cfg.moe.get("n_redundant_experts", 0)
2
...
3
self.n_local_routed_experts = self.n_routed_experts // self.ep_size
4
assert self.n_routed_experts % self.ep_size == 0
5

6
self.linear1 = ParallelRoutedLinear(
7
    cfg.hidden_size,
8
    self.hidden_dim * 2,
9
    self.n_local_routed_experts,
10
    use_fp8_gemm=cfg.fp8_mlp,
11
)

ParallelRoutedLinear 的第三个参数就是本卡要分配多少组专家权重。这个数字一旦定死，模型容器就按它分配显存。EPLB 的重排只会在这些槽位之间搬内容，不会改变槽位数——这是理解整个系统的一个关键前提：EPLB 是「在固定的槽位集合上重排内容」，不是「动态增减显存」。这也是为什么运行时重排可以做得很轻：没有分配器的介入，只有拷贝。

初始布局与第一张映射表

物理槽位在装载时的初始内容是确定性的：

1
def get_initial_global_physical_to_logical_map(
2
    num_routed_experts: int,
3
    num_extra_experts: int,
4
) -> list[int]:
5
    """初始布局：[原始 routed experts, 复制的 routed experts]。"""
6
    return [
7
        *range(num_routed_experts),
8
        *(expert_index % num_routed_experts for expert_index in range(num_extra_experts)),
9
    ]

即：前 num_routed_experts 个槽位放全部 logical expert 本体，剩下的 num_extra_experts 个槽位放循环复制的 expert（用 % 取模实现）。以 $N = 256$、extra = 32 为例，就是前 256 个槽放本体、后 32 个槽放 expert 0…31。这个初始分配显然不是最优的——它只是复用了编号最小的那批 expert，与真实热度毫无关系。ExpertLoadBalancer.__init__ 的注释也承认这一点，所以第一次重排被刻意提前：

1
# 首次 rearrange 提前到 1/4 周期后，尽早纠正冷启动的不均衡。
2
initial_rearrangement_delay = max(1, self._rearrangement_interval // 4)
3
self._rearrangement_step = max(0, self._rearrangement_interval - initial_rearrangement_delay)

如果重排周期是 10000 步（DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL 未设时默认为 0，即关闭；开启后由运维配置一个周期），第一次重排会在 2500 步时发生，而不是等到 10000 步。冷启动窗口缩短到四分之一，这个系数是拍出来的，但方向是明确的：权重装载时的均匀假定一定不成立，越早纠正越好。

由 physical_to_logical_map 再派生两张表：

1
logical_to_physical_map = torch.full((num_logical_experts, num_redundant_experts + 1), -1, dtype=torch.int32)
2
logical_count = torch.zeros((num_logical_experts,), dtype=torch.int32)
3
for physical_index in range(num_physical_experts):
4
    logical_index = physical_to_logical_map[physical_index]
5
    logical_to_physical_map[logical_index, logical_count[logical_index]] = physical_index
6
    logical_count[logical_index] += 1

• logical_to_physical_map[L]：一个长度 num_redundant_experts + 1 的向量，列出 logical expert $L$ 的所有副本槽位；无效位是 -1。
• logical_count[L]：$L$ 的副本数。

为什么两张表都要？因为它们服务的读者不同。

• dispatch 侧需要 logical_to_physical_map：路由算出「这个 token 该去 expert 37」，然后必须回答「37 的副本里，哪一个离我最近/负载最低」——这是按 logical 查 physical。
• 权重装载与 shuffle 侧需要 physical_to_logical_map：本卡的槽位 3 里装的到底是哪个 expert，才能把对应的权重行搬过来——这是按 physical 查 logical。

两张表互为逆映射（在 logical_count 限定的有效位内）。单元测试 test_eplb_policy.py 里的 _check_structure 正是把这个逆关系当作不变量来断言：

1
# log_count == the number of valid slots in log2phy
2
assert log_count.equal((log2phy >= 0).sum(-1))
3
# every physical slot can be found in its logical expert's log2phy
4
for layer in range(num_layers):
5
    for physical_index, logical_index in enumerate(phy2log[layer].tolist()):
6
        assert physical_index in log2phy[layer][logical_index].tolist()
7
# total replica count == physical expert count
8
assert log_count.sum(-1).equal(torch.full((num_layers,), num_physical_experts, dtype=torch.int))

三条断言分别锁死「count 与有效位一致」「反查必存在」「副本总数守恒」。这类结构不变量在数值型代码里非常重要：算法只保证近似最优，但结构必须严格自洽，否则一次映射错位就会让某个 expert 的权重被搬进错误的槽位，表现为输出乱码而不是报错。测试里还有 test_rebalance_experts_global_deterministic（同输入必须同输出）和 test_rebalance_does_not_modify_history（算法不能就地改写历史统计张量）两条，也是一样的防守思路。

目标：让每副本负载的最大值最小

先看输入的形状：

1
def rebalance_experts_global(
2
    tokens_per_expert: torch.Tensor,   # [num_steps, num_moe_layers, num_logical_experts]
3
    num_physical_experts: int,
4
    num_local_physical_experts: int,
5
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

tokens_per_expert 是一个三维历史张量：步数 × 层数 × logical expert 数。也就是说，算法看的不只是「最近一步每个专家多少 token」，而是最近 $S$ 步的完整轨迹，逐层独立求解。为什么是「轨迹」而不是「平均」？因为目标函数里出现了 max：

1
# score[..., j] 是 expert j 当前的每副本负载，score_added 是给 j 加一个副本后的每副本负载；
2
# 把 argmax 位置替换成其 score_added 后，values[..., j] 即“副本给 j”时的最大每副本负载
3
# （j 非 argmax 时最大值不变，j 为 argmax 时是替换后重算的 max），按 step 求和取 argmin 得最优复制对象。

翻译成公式：对每个候选的复制对象 $j$，定义

其中 $\tilde{L}$ 是「假如把副本给 $j$ 之后」各 expert 的每副本负载：

$t(s, e)$ 是历史第 $s$ 步 expert $e$ 的 token 数，$c(e)$ 是它当前的副本数。然后选 $\arg\min_j \text{cost}(j)$。

这个目标函数的含义是：副本应该给那个「加副本之后，能把整层最热的那个每副本负载压得最低」的专家。逐 step 取 max 再求和，意味着算法关心的是历史每一小步（尤其是 prefill 那种大块投入的步）里的峰值，而不是把历史平均掉。

一次「替换式最大值」的向量化技巧

直接实现这个目标需要 $O(N)$ 次「假设加副本」的重算，每次 $O(S \cdot N)$，总体 $O(S \cdot N^2)$。以 $N = 256$、窗口 100 步计，这是千万量级的元素操作——勉强能跑，但代码里用了一个更省的技巧：

1
tokens_per_expert_all_diff = tokens_per_expert + arange_logical * 1e-4
2
for i in range(num_redundant_experts):
3
    score = tokens_per_expert_all_diff / logical_count  # NOTE: Values in score must be different from each other
4
    score_added = tokens_per_expert / (logical_count + 1)
5
    values, indices = score.max(-1, keepdim=True)
6
    values = values.expand_as(score).contiguous()
7
    score.scatter_(-1, indices, score_added.gather(-1, indices))
8
    values.scatter_(-1, indices, score.max(-1, keepdim=True).values)
9
    redundancy_indices = values.sum(0).argmin(-1)  # [num_moe_layers]，本轮被复制的 logical 编号

逐步读：

1. score 是当前每副本负载，形状 [steps, layers, experts]。
2. values, indices = score.max(-1) 找到每一 (step, layer) 上的最热专家及其负载。
3. values.expand_as(score) 把最大值广播开，于是 values[..., j] 目前都是「不加副本时的最大值」。
4. score.scatter_(-1, indices, score_added.gather(-1, indices))：只把最热那一个位置替换成「加了副本之后的每副本负载」。其他位置的每副本负载加副本不变（因为加副本只影响被复制那一个专家）。
5. values.scatter_(-1, indices, score.max(-1, keepdim=True).values)：重新求一次 max，写回最热位置。因为只有那个位置变了，新 max 要么是新的替换值，要么是原来的次大值——一次 max 就够。
6. values.sum(0) 对 step 求和，得到 [layers, experts] 的 cost 矩阵，argmin(-1) 选出每层该复制的 expert。

这个技巧成立的前提是：每轮只加一个副本，且每轮的最大值只有一个来源。 所以每轮循环只做一次 max/scatter，复杂度降到 $O(S \cdot N)$ 每轮，总共 $O(S \cdot N \cdot R)$（$R$ 为冗余副本数，是 32 / 64 这一档）。拿 $N = 256$、$S = 100$、$R = 32$ 估算，约 80 万个元素的操作量，对 GPU 张量来说只是毫秒级。

那行 tokens_per_expert + arange_logical * 1e-4 是干嘛的？注释写了 NOTE: Values in score must be different from each other。当多个专家的负载完全相等时，argmax 会固定返回最小索引，导致「只有一个专家被反复选中」。加一个与 expert 索引成正比的极小扰动，就把平局打破了，而且破法是确定性的（总是偏向索引大的），不会引入随机性——这是确定性要求下处理平局的常见手法。1e-4 这个量级要远小于 token 数的量化单位（整数 token 计数，差值至少 1），所以它只影响平局，不影响任何真实比较。

redundant_indices 拿到之后，写入两张表并更新副本数：

1
physical_index = num_logical_experts + i  # 本轮冗余副本占用的 physical 槽位
2
physical_to_logical_map[:, physical_index] = redundancy_indices
3
redundancy_count = logical_count.gather(-1, redundancy_indices.unsqueeze(-1)).squeeze(-1)
4
logical_to_physical_map[arange_layers, redundancy_indices, redundancy_count] = physical_index
5
logical_count[arange_layers, redundancy_indices] += 1

注意这里没有把副本放在「最优的卡」上——它只是把冗余槽位从 num_logical 开始按复制顺序填进 physical 编号里。放置是下一段 _batch_greedy_indices 的事。这种拆法是刻意的：先决定「复制谁」，再决定「放哪」，两个子问题各自独立、各自贪心，整体是可解的——如果合并成一个「给哪些专家加副本 + 放在哪张卡」的联合优化，就是一个二次分配问题，代价会高得多。

一个边界情况：num_local_physical_experts == 1 时跳过放置：

1
if num_local_physical_experts > 1:
2
    assert num_physical_experts % num_local_physical_experts == 0
3
    indices = _batch_greedy_indices(...)

如果每卡只有一个槽位，那么「哪个槽位在哪张卡」就是固定的（第 $r$ 个槽位在第 $r$ 张卡），放置问题退化为「哪个专家该拥有那个槽位」，没有搬运空间。测试用例里专门有 (1, 1, 8, 8, 1) 与 (2, 1, 4, 16, 1) 两例覆盖这条分支。

放置：逐层独立的贪婪

接下来是 _batch_greedy_indices，它是整个 EPLB 里最需要读透的函数。

1
weights = tokens_per_expert.sum(dim=0, dtype=torch.float32) / logical_count
2
# 按每副本历史总负载降序排列，同负载时保留原物理槽位顺序
3
physical_experts = physical_to_logical_map.to(torch.int64)
4
physical_order = weights.gather(-1, physical_experts).argsort(dim=-1, descending=True, stable=True)
5
expert_order = physical_experts.gather(-1, physical_order)
6
vectors = tokens_per_expert.permute(1, 2, 0).to(dtype=torch.float32, memory_format=torch.contiguous_format, copy=True)
7
vectors.div_(logical_count.unsqueeze(-1))

四步：

• weights 把历史 step 求和，得到每个 logical expert 的每副本总负载（/ logical_count 是「分摊到副本上」）。
• physical_order 是把物理槽位按它对应的每副本负载降序排列后的序号。stable=True 保证相同负载时保持原槽位顺序——这是确定性的一部分。
• expert_order 是同一排列作用在 expert 编号上的结果。
• vectors 是每副本、逐步的负载向量，形状 [layers, experts, steps]。注意 .to(..., memory_format=torch.contiguous_format, copy=True)——强制拷贝，因为下一行要 div_（就地除），不能污染输入。测试 test_rebalance_does_not_modify_history 专门验证了这一点。

然后是目标函数与贪心循环：

1
# 优化目标是逐 layer 独立的，对固定的 layer 而言，目标是将这层的 physical expert 分到各个 rank 上
2
# 并且使得 sum_{step i} max_{rank j} (sum_{physical expert e} [e 被分到 rank j] e 在 step i 的负载) 尽量小
3
# 具体贪心算法: 每次尝试将要考虑的 expert 放到各个 rank 上，选取当前最小化上述目标的放置方案

写成公式，对每一层，找一个把 $N_{phy}$ 个槽位分到 $R$ 张卡上的划分：

其中 $v(e, s)$ 是槽位 $e$ 在第 $s$ 步的负载向量分量。这个目标函数里有一个 max 套在 sum 外面再加 sum，不是可分解的（不是每张卡独立最优就能全局最优），所以精确求解要枚举，代价不可接受。DeepInfer 用的是一个「按当前负载从重到轻逐个安置」的贪心。

循环体逐行读：

1
for job_index in range(num_physical):
2
    experts = expert_order[:, job_index]
3
    vector = vectors[layers, experts]
4
    # candidates[layer][j][step] 描述将当前 expert 放到 rank j 上后，各历史 step 在 rank j 上的负载
5
    # 因此需和 peaks 取 max 得到当前 expert 放到 rank j 上后，各历史 step 的 rank 最大负载
6
    torch.add(loads, vector.unsqueeze(1), out=candidates)
7
    torch.maximum(candidates, peaks.unsqueeze(1), out=candidates)
8
    torch.sum(candidates, dim=-1, out=scores)
9
    # 排除已满的 rank 后选择最低分，argmin 在同分时返回最小 rank ID
10
    scores.masked_fill_(occupancy == capacity, float("inf"))
11
    rank = scores.argmin(dim=-1)
12
    occupied = occupancy[layers, rank]
13
    indices[layers, rank * capacity + occupied] = physical_order[:, job_index]
14
    occupancy[layers, rank] = occupied + 1
15
    updated = loads[layers, rank] + vector
16
    loads[layers, rank] = updated
17
    torch.maximum(peaks, updated, out=peaks)

每一步的状态有三个：

• loads[layers, ranks, steps]：当前已安置的槽位在各卡各步上的负载和。
• peaks[layers, steps]：当前所有卡的峰值（跨 rank 的 max）。
• occupancy[layers, ranks]：各卡已放的槽位数。

对当前这个槽位（按降序处理，所以先安置的是最热的），算法尝试把它放到每一张卡上：

1. candidates = loads + vector——假设放到卡 $r$ 之后，卡 $r$ 的逐步负载（广播到所有 $r$）。
2. max(candidates, peaks)——因为其他卡的负载不变，所以「放完之后各步的全局峰值」= max(卡 $r$ 的新负载, 原峰值)。这就是那个不可分解的 max 被巧妙地以「原峰值 + 单卡增量」的形式算出来的地方：只有被选中的那张卡的负载变了，所以全局 max 只需与原 peaks 比一次。
3. scores = sum over steps——对每张卡 $r$，这就是把它作为落点时的目标函数值。
4. masked_fill_(occupancy == capacity, inf)——已经放满 $capacity$ 个槽位的卡直接排除。
5. argmin 选分数最低的卡；注释点明 argmin 在同分时返回最小 rank ID，所以平局恒定偏向小编号卡。这是确定性放置的第二处保证。
6. 落位：indices[layers, rank*capacity + occupied] = physical_order[:, job_index]——把物理槽位编号写进「第 rank 张卡的第 occupied 个位置」。

最终返回的 indices 就是新的 physical 槽位排列：每连续 capacity 个元素对应一张卡。

为什么「按负载降序、逐个安置、每次选当前代价最小的卡」这个贪心够用？两个原因：

• 大件先放（LPT 启发式，Longest Processing Time first）。这是调度论里的经典结论：对 make-span 最小化问题，先安置大件能显著降低最坏情况，因为大件放错地方后无法补救，而小件可以填缝。
• 目标与贪心的一步代价一致。贪心每步最小化的正是「把当前件加进去之后的峰值之和」，与全局目标同向。虽然不保证全局最优，但在这类「多件、多机、看峰值」的问题上，LPT 的经验表现通常已接近最优，而代价从指数级降到 $O(N \cdot R)$。

单元测试 test_batch_greedy_objective_and_ties 用五个手工构造的布局把这三条性质锁死了：

1
(
2
    [[10, 0, 0, 0, 0, 0], [0, 6, 4, 0, 0, 0]],
3
    [0, 1, 2, 3, 4, 5],
4
    [0, 1, 3, 2, 4, 5],
5
),

这一例的语义注释是「给 expert 1 选卡时两张卡分数相同，rank 0 的累计负载虽然更高，仍然优先 rank 0」。注意这里出现的正是贪心的一个已知弱点：平局偏向小编号卡，不偏向当前更空的卡。这在负载对称时无所谓，在负载不对称时会带来轻微偏差——但换来了确定性，而确定性在 EPLB 里是硬需求：所有 rank 必须独立算出完全相同的映射，否则一张卡以为槽位 3 是 expert 37、另一张卡以为它是 expert 41，权重搬过去就全错了。

重映射 logical_to_physical

放置完成后，logical_to_physical_map 里存的槽位编号已经失效（槽位位置变了），必须一起重映射：

1
indices = _batch_greedy_indices(tokens_per_expert, physical_to_logical_map, logical_count, num_local_physical_experts)
2
physical_to_logical_map = physical_to_logical_map.gather(-1, indices)
3
# indices.argsort 给出旧槽位 -> 新槽位的映射，用它重映射 logical_to_physical_map 里的槽位编号（-1 先临时置 0 再还原）
4
mask = logical_to_physical_map == -1
5
logical_to_physical_map[mask] = 0
6
logical_to_physical_map = (
7
    indices.argsort(-1)
8
    .gather(-1, logical_to_physical_map.view(num_moe_layers, -1).to(torch.int64))
9
    .view_as(logical_to_physical_map)
10
    .to(torch.int)
11
)
12
logical_to_physical_map[mask] = -1

这里有三个技巧：

• physical_to_logical_map.gather(-1, indices)：indices 是「新位置 → 旧槽位」的排列，gather 一次就把映射按新顺序重排好了。
• indices.argsort(-1)：对排列取 argsort 得到逆排列，即「旧槽位 → 新位置」。这是用一次 $O(N \log N)$ 排序替代手工构造逆映射的标准写法。
• -1 的处理：gather 不能带负索引，所以先把 -1 临时置 0、gather 完再还原。用 mask 记住哪些位置无效，最后恢复。

这三步合起来保证了「输出的两张表仍然是互逆的」——也就是测试里断言的那条不变量。重排算法可以近似，但两张表的自洽不能有例外。

通用路径：NCCL batched P2P

expert_model_parallel_shuffle_inplace 对每一层调用 _shuffle_layer_p2p。后者的骨架是三个阶段：卡内拷贝 → 跨卡 P2P → 收尾回填。

阶段一，卡内能满足的搬运：

1
def _plan_intra_copies(origin_local: list[int], current_local: list[int]) -> tuple[list[bool], list[tuple[int, int]]]:
2
    """本卡内部能满足的搬运。received[dst] 表示槽位 dst 无需远端数据；copies 是卡内 (src, dst) 拷贝对。"""
3
    received = [o == c for o, c in zip(origin_local, current_local, strict=True)]
4
    copies = []
5
    for src, expert in enumerate(origin_local):
6
        for dst in range(len(current_local)):
7
            if not received[dst] and expert == current_local[dst]:
8
                received[dst] = True
9
                copies.append((src, dst))
10
    return received, copies

received[dst] = (origin_local[dst] == current_local[dst]) 是「这个槽位压根没变」的快速判断。剩下的逐个配对，配对成功就标记已满足。注意这里的拷贝源是 weight[src]、目标是 buf[dst]（缓冲），不是直接 weight[dst]——因为源和目标是同一张张量的不同行，直接写会覆盖还没读的源。

阶段二，跨卡收发。先算出「谁有、谁要」：

1
def _remote_routes(origin: list[int], current: list[int], num_local: int, expert: int) -> tuple[list[tuple[int, int]], list[int]]:
2
    """expert 的跨卡收发双方。
3

4
    senders: 持有 expert 的 (rank, 第一个本地槽位)，按 rank 升序；
5
    receivers: 需要 expert 且卡内无法满足（origin 里本卡没有 expert）的 rank，升序。
6
    """
7
    senders = []
8
    for i, e in enumerate(origin):
9
        if e == expert and (not senders or senders[-1][0] != i // num_local):
10
            senders.append((i // num_local, i % num_local))
11
    wanters = []
12
    for i, e in enumerate(current):
13
        if e == expert and (not wanters or wanters[-1] != i // num_local):
14
            wanters.append(i // num_local)
15
    sender_ranks = {rank for rank, _ in senders}
16
    receivers = [rank for rank in wanters if rank not in sender_ranks]
17
    return senders, receivers

两处细节值得注意：

• if e == expert and (not senders or senders[-1][0] != i // num_local) ——同一张卡上如果有多个副本，只记第一个。因为同一 expert 在卡内的多个槽位内容相同，跨卡只搬一次就够，卡内再复制。这直接减少了通信量。
• receivers 排除了自己也是 sender 的卡。一张卡同时有该 expert 的旧副本和新副本时，不需要跨卡搬——它在阶段一的卡内拷贝里已经解决。

同一个 expert 可能有多份副本（冗余专家），于是「谁给谁发」需要配对。配对函数很短但很关键：

1
def _paired_sender(j: int, num_senders: int, num_receivers: int) -> int:
2
    """第 j 个 receiver 由第几个 sender 供给：前 n * num_senders 个均分（每个 sender 连续 n 个），其余轮转补齐。
3

4
    收发双方用同一函数配对，保证 P2P 两端一致。
5
    """
6
    n = num_receivers // num_senders
7
    return j // n if j < n * num_senders else j - n * num_senders

「收发双方用同一函数配对」这句注释是整个 P2P 设计的关键：发送端和接收端各自独立计算配对关系，只有保证两侧算出的 $(\text{sender}, \text{receiver})$ 对完全一致，P2P 的 tag 才能对上。用一个纯函数而非通信协商来做这件事，省掉了一轮元数据交换。

发送端的循环：

1
for expert, src in _first_slot_per_expert(origin_local):
2
    senders, receivers = _remote_routes(origin, current, num_local, expert)
3
    my_index = [rank for rank, _ in senders].index(ep_rank)
4
    for j, dst in enumerate(receivers):
5
        if _paired_sender(j, len(senders), len(receivers)) == my_index:
6
            p2p_buckets[(ep_rank - dst) % len(p2p_buckets)].extend(
7
                dist.P2POp(dist.isend, weight[src], dist.get_global_rank(ep, dst), ep) for weight in weights
8
            )

分桶逻辑 p2p_buckets[(ep_rank - dst) % len(p2p_buckets)] 是 DEEPINFER_REARRANGE_P2P_GROUPS（默认 1）控制的分组策略：把「与我 rank 差同余」的对端分到同一个桶，逐桶批量收发。默认 1 个桶就是一次性批量发；设成大于 1 的值可以按拓扑分组（例如按节点差分组，优先机器内通信——代码里留了 TODO：这里可以考虑优先机器内通信和 broadcast）。分桶的意义在于 batch_isend_irecv 的批量大小会影响 NCCL 内部通道的选择，按拓扑分桶有机会让同桶内的 P2P 走更快的链路。

接收端做的是同一件事的镜像，并且把接收缓冲的落点算出来：

1
experts_to_recv = _first_slot_per_expert(current_local, skip=received)
2
for expert, dst in experts_to_recv:
3
    senders, receivers = _remote_routes(origin, current, num_local, expert)
4
    src_rank, _ = senders[_paired_sender(receivers.index(ep_rank), len(senders), len(receivers))]
5
    p2p_buckets[(src_rank - ep_rank) % len(p2p_buckets)].extend(
6
        dist.P2POp(dist.irecv, buf[dst], dist.get_global_rank(ep, src_rank), ep) for buf in recv_buffers
7
    )

skip=received 把阶段一已解决的槽位排除，避免重复接收。注意接收目标是 buf[dst]（缓冲）不是 weight[dst]。

阶段三，收尾回填：

1
recv_slot = dict(experts_to_recv)
2
for i in range(num_local):
3
    if origin_local[i] == current_local[i]:
4
        continue
5
    src = i if received[i] else recv_slot[current_local[i]]
6
    for weight, buf in zip(weights, recv_buffers, strict=True):
7
        weight[i].copy_(buf[src])

src 的两种取值对应两种来源：received[i] 为真说明这个槽位的内容在缓冲里已经有了（来自卡内拷贝或远端接收），直接取 buf[i]；否则说明同一个 expert 在卡内还有另一个槽位已经填好，从那个槽位拷（recv_slot[current_local[i]] 给出的是「该 expert 在本卡的第一个接收槽位」）。这一步是「同一 expert 的多个槽位只搬一次、卡内再复制」策略的收尾，也是这套实现能把通信量压下来的最后一块。

单机路径：symm_mem IPC 镜像

如果 EP 组落在同一个 NVLink 域内（PG.is_within_nvl_domain(ep_group)），DeepInfer 会走一条完全不同的路径：

1
def _initialize_shuffle(self) -> None:
2
    """首次 rearrange：释放缓存显存防 OOM；单机时建 symm_mem 镜像走 IPC shuffle。"""
3
    ep_group = self._ep_group
4
    torch.cuda.empty_cache()
5
    # TODO(NPU): NPU doesn't support symm_mem yet.
6
    self._intranode_shuffle = PG.is_within_nvl_domain(ep_group) and PLATFORM.is_gpu
7
    if not self._intranode_shuffle:
8
        return
9

10
    logger.info("[EPLB] shuffle intra node")
11
    # 镜像只需容纳一层权重，shuffle 时逐层刷新复用。
12
    self._symmetric_weight_buffers = [
13
        symm_mem.empty(weight.shape, dtype=weight.dtype, device="cuda") for weight in self._expert_weights[0]
14
    ]
15
    self._symmetric_memory_handles = [symm_mem.rendezvous(buffer, ep_group) for buffer in self._symmetric_weight_buffers]
16
    self._ipc_weight_buffers = [
17
        [
18
            handle.get_buffer(rank, buffer.shape, buffer.dtype)
19
            for buffer, handle in zip(self._symmetric_weight_buffers, self._symmetric_memory_handles, strict=True)
20
        ]
21
        for rank in range(ep_group.size())
22
    ]

用的是 PyTorch 的对称内存（symmetric memory）接口 torch.distributed._symmetric_memory：

• symm_mem.empty(...) 分配一块所有 rank 都可见的缓冲区。
• symm_mem.rendezvous(buffer, ep_group) 交换句柄，让每张卡都能拿到其他卡的这块内存的映射。
• handle.get_buffer(rank, shape, dtype) 返回一个指向第 rank 张卡的缓冲区的张量。

于是本卡可以直接 weight[dst].copy_(peer[src_slot])——其中 peer 是 ipc_experts_weight[src_rank]，即对端卡上的镜像张量。拷贝完全绕开了 NCCL，因为 GPUDirect 让本卡可以直接读对端显存（同机 NVLink 域内）。

看 _shuffle_layer_intranode 的同步：

1
ep.barrier().wait()  # 等所有 rank 读完上一层的镜像
2
for m, weight in zip(mirror, weights, strict=True):
3
    m.copy_(weight)
4
ep.barrier().wait()  # 等所有 rank 的镜像就绪
5

6
# 卡内 origin -> current 的拷贝（源数据从自己的镜像读，避免覆盖）
7
received, intra_copies = _plan_intra_copies(origin_local, current_local)
8
for src, dst in intra_copies:
9
    for weight, m in zip(weights, mirror, strict=True):
10
        weight[dst].copy_(m[src])
11

12
# 远端读取：对每个缺失的 expert，直接从配对 sender 的镜像拷贝（同一 expert 多个槽位只读一次）
13
experts_to_recv = _first_slot_per_expert(current_local, skip=received)
14
for expert, dst in experts_to_recv:
15
    senders, receivers = _remote_routes(origin, current, num_local, expert)
16
    src_rank, src_slot = senders[_paired_sender(receivers.index(ep_rank), len(senders), len(receivers))]
17
    received[dst] = True
18
    for weight, peer in zip(weights, ipc_experts_weight[src_rank], strict=True):
19
        weight[dst].copy_(peer[src_slot])

两道 barrier 各管一件事：第一道保证「别人已经读完我上一层的镜像，我可以覆写它了」；第二道保证「我的镜像已刷新，别人可以来读」。镜像只有一层的量（注释：镜像只需容纳一层权重，shuffle 时逐层刷新复用），这是显存上的硬约束——如果每层都留一份镜像，等于把专家权重翻倍，在几十 GB 量级上不可接受。

代价是逐层串行的 barrier：每层两次全组同步，层数越多同步开销越大。而 P2P 路径可以跨层用同一个 recv_buffers，但没有 barrier。两条路径的取舍是：P2P 用通信库的批量通道换取可扩展性（跨机也适用），IPC 用显存镜像换取零通信库开销（但必须同机、且要付同步次数）。_batch_greedy_indices 里没有类似的二选一，是因为它只在 CPU/CUDA 张量上做算术，与硬件拓扑无关。

还有个小细节：_initialize_shuffle 里的 torch.cuda.empty_cache() 与注释「释放缓存显存防 OOM」。第一次重排要临时分配镜像缓冲，而此时进程里通常堆着大量已释放但未归还的 CUDA 缓存块，先清一次能显著降低 OOM 概率。这也说明重排的显存水位是这条链路上最容易出事的地方——_rearrange_experts 被包在 @allow_cuda_sync() 装饰器下（该装饰器会临时放开「禁止 CUDA 同步」的守卫，因为重排需要 torch.cuda.synchronize() 来打点计时），并且只在 rank 0 上打日志：

1
should_log_metrics = ep_group.rank() == 0
2
if should_log_metrics:
3
    torch.cuda.synchronize()
4
    start_time = time.perf_counter()

ExpertLoadBalancer 的状态面

类注释点明了它对外的角色：

1
class ExpertLoadBalancer:
2
    """EPLB 运行时状态：累计各 expert 的 token 数，按周期重算放置并 shuffle 权重。
3

4
    所有环境配置均在构造时读取并固定。
5
    """
6

7
    # MoE dispatch、路由与 sampler 直接读写以下状态。
8
    expert_count: torch.Tensor  # [num_moe_layers, num_local_physical_experts]，每 step 清零
9
    physical_to_logical_map: torch.Tensor  # [num_moe_layers, num_physical_experts]
10
    logical_to_physical_map: torch.Tensor  # [num_moe_layers, num_logical_experts, num_redundant_experts + 1]，无效位 -1
11
    logical_count: torch.Tensor  # [num_moe_layers, num_logical_experts]

注意 expert_count 的注释：MoE dispatch、路由与 sampler 直接读写以下状态。这是 EPLB 与模型前向之间唯一的数据接口，而且是热路径上的。MoE dispatch 在每次前向里给对应槽位的计数器加一（按本卡实际处理的 token 数），EPLB 每个 step 读走并清零。

为什么计数器挂在 physical 槽位而不是 logical expert 上？因为本卡只能观察到「发到我这里的 token 落在了我哪个槽位」，它无从知道全局某个 logical expert 的总热度。这个设计有一个直接推论：统计口径是「按副本」的，同一 logical expert 的多个副本在不同卡上各自计数，最终在重排时通过 all_reduce 聚合回 logical 维度。

一个 step 里发生什么

1
def step(self, is_prompt: bool) -> None:
2
    """由 DeepseekForCausalLM.sample 每 step 调用，更新重排与负载统计。"""
3
    if self._ep_group.size() <= 1:
4
        return
5

6
    self._update_rearrangement(is_prompt)
7
    self._update_load_balance_stats()
8
    self.expert_count.zero_()

三件事：

1. _update_rearrangement：维护滑动窗口并判断是否该重排。
2. _update_load_balance_stats：把本 step 的计数存进统计缓冲，够一个周期就 all_gather 出全局比值。
3. expert_count.zero_()：清零，准备下一步。

is_prompt 这个参数说明调用点区分 prefill 与 decode 步。为什么调度器要把这个信息传进来？看 _update_rearrangement：

1
def _update_rearrangement(self, is_prompt: bool) -> None:
2
    if self._rearrangement_interval == 0:
3
        return
4
    if is_prompt and self._skip_rearrangement_on_prefill:
5
        return
6

7
    window_index = self._expert_count_window_index
8
    self._expert_count_window[window_index].copy_(self.expert_count)
9
    self._expert_count_window_index = (window_index + 1) % self._expert_count_window.size(0)
10

11
    if is_prompt and self._schedule_rearrangement_after_prefill:
12
        # 下一个非 prefill step 立即 rearrange。
13
        self._rearrangement_step = self._rearrangement_interval - 1
14
        return
15

16
    self._rearrangement_step += 1
17
    if self._rearrangement_step < self._rearrangement_interval:
18
        return
19
    self._rearrangement_step = 0
20
    self._rearrange_experts()

三种行为由两个 flag 决定，它们的构造来源是：

1
self._schedule_rearrangement_after_prefill = is_normal_engine and envs.DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE
2
self._skip_rearrangement_on_prefill = is_normal_engine and not self._schedule_rearrangement_after_prefill

• _skip_rearrangement_on_prefill：prefill 步完全不参与计数步进。理由是 prefill 阶段的模型执行时间远长于 decode（一次几千 token 的大前向），在 prefill 步里插一段几十 GB 的权重搬运会把这一步的时延再拉长一截，而这段时间整个 EP 组都在等。跳过 prefill 步，把重排留给紧凑的 decode 步，是「不在贵的地方加更贵的东西」。
• _schedule_rearrangement_after_prefill：prefill 步不计数，但立刻把计数器推满（self._rearrangement_step = self._rearrangement_interval - 1），于是下一个非 prefill 步就触发重排。语义是「prefill 之后的流量分布最值得立刻跟进」——prefill 通常意味着新请求涌入、缓存命中率变化、路由分布整体偏移，此时尽快校正专家放置收益最大。注释写的就是 下一个非 prefill step 立即 rearrange。

两个 flag 只在 is_normal_engine 时生效，说明引擎形态也影响 EPLB 的节奏——第一篇提过 DeepInfer 有 NORMAL 与 ALL 两种引擎类型（对应不同的 prefill 调度策略，见 DEEPINFER_PREFILL_ATTEMPT_INTERVAL / DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL 的注释）。

窗口：为什么必须是环形缓冲

1
expert_count_window_size = envs.DEEPINFER_EXPERT_COUNT_WINDOW_SIZE
2
self._rearrangement_interval = envs.DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL
3
if expert_count_window_size < 0:
4
    raise ValueError(f"DEEPINFER_EXPERT_COUNT_WINDOW_SIZE={expert_count_window_size} must be non-negative")
5
if self._rearrangement_interval != 0 and not (0 < expert_count_window_size <= self._rearrangement_interval):
6
    raise ValueError(
7
        f"DEEPINFER_EXPERT_COUNT_WINDOW_SIZE={expert_count_window_size} must be between 1 and "
8
        f"DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL={self._rearrangement_interval} "
9
        "when rearrangement is enabled"
10
    )
11

12
self.expert_count = torch.zeros(num_moe_layers, num_local_physical_experts, dtype=torch.int)
13
self._expert_count_window = torch.zeros(
14
    expert_count_window_size,
15
    num_moe_layers,
16
    num_local_physical_experts,
17
    dtype=torch.int,
18
)
19
self._expert_count_window_index = 0

三件事：

• 窗口大小默认 100（DEEPINFER_EXPERT_COUNT_WINDOW_SIZE），且被强制约束在 $[1, \text{rearrangement\_interval}]$ 之内。为什么窗口不能大于重排周期？因为每次重排都会刷新窗口索引，如果窗口比周期还长，那么这个窗口永远填不满，一部分 slot 永远是零 —— 目标函数里的 sum over steps 会被大量无意义的零稀释。这个断言是在保护目标函数的语义。
• 窗口按 physical 槽位存（num_local_physical_experts），不按 logical。这是前面说的「按副本计数」的存储形态。
• 环形缓冲的写指针 _expert_count_window_index 在 % size 下循环。用固定大小的环形缓冲而不是无限追加，是为了让重排时刻的内存占用与运行时长无关——一个跑几周的推理服务若把每步计数都留着，会积累成 GB 级的张量。

重排时，窗口先聚合回 logical 维度再 all-reduce：

1
# 窗口内各 logical expert 的 token 数：本地 physical 槽位聚合后全局 all-reduce
2
local_physical_to_logical_map = self.physical_to_logical_map[
3
    :, ep_group.rank() * num_local_physical_experts : (ep_group.rank() + 1) * num_local_physical_experts,
4
]
5
logical_expert_token_counts = torch.zeros(window_size, num_moe_layers, num_logical_experts, ...)
6
logical_expert_token_counts.scatter_add_(
7
    -1,
8
    local_physical_to_logical_map.to(torch.int64).expand_as(self._expert_count_window),
9
    self._expert_count_window,
10
)
11
dist.all_reduce(logical_expert_token_counts, group=ep_group)

scatter_add_ 是这里的关键算子：把窗口张量的最后一维（本卡槽位）按 local_physical_to_logical_map 散射累加到 logical 维度上。因为同一 logical expert 的多个副本可能在本卡内存在（冗余专家允许重复），scatter_add_ 会正确地把它们相加。然后一次 all_reduce 把全局统计汇总——注意这里 all_reduce 的是「token 数」而不是「已除以副本数的负载」，除以副本数是在 policy.py 里做的（vectors.div_(logical_count.unsqueeze(-1))）。

为什么这样分工？因为副本数在重排过程中会变，而「历史 token 数」是原始事实。先聚合事实，再按当前副本数折算，比先折算再聚合在语义上更干净，也让 policy.py 能同时用 logical_count 做分母和做贪心的容量约束。

观测指标

重排后 rank 0 打印一行指标（_log_rearrangement_metrics）：

1
physical_expert_loads = logical_expert_token_counts.div(logical_count.unsqueeze(0)).gather(
2
    -1,
3
    physical_to_logical_map.to(torch.int64).unsqueeze(0),
4
)  # [window_size, layers, physical_experts]
5
rank_loads = physical_expert_loads.unflatten(-1, (-1, num_local_physical_experts)).sum(dim=-1)  # [window_size, layers, ranks]
6
mean_peak_rank_load = rank_loads.max(dim=-1).values.mean().item()
7
peak_rank_load = rank_loads.max().item()
8
mean_rank_load = rank_loads.mean().item()
9
...
10
rank_balance = mean_rank_load / (mean_peak_rank_load + 1e-6)
11
logical_balance = mean_logical_expert_load / (max_logical_expert_load + 1e-6)
12
logger.info(
13
    "[EPLB] rearranged: experts/gpu=%d, rank_balance=%.4f (mean=%.1f, peak_mean=%.1f), rank_peak=%.1f, "
14
    "logical_balance=%.4f (mean=%.1f, max=%.1f), physical_peak=%.1f, time=%.3fs",
15
    ...
16
)

指标的设计和算法的目标函数一一对应：

• logical_balance = mean_logical_expert_load / max_logical_expert_load——复制前的天花板。它衡量「专家热度的天然不均衡有多严重」，是 1 就说明所有专家一样热，冗余完全没必要。这个数字低到某个程度，说明该加大冗余数了。
• rank_balance = mean_rank_load / mean_peak_rank_load——复制并放置后的实际均衡度。它的上限是 1；但即使专家热度完美不均，只要副本够多、放得够散，这个数也能接近 1。rank_balance 才是 EPLB 真正的成绩单，logical_balance 是它的问题输入。
• rank_peak（全部 step、全部层、全部卡里的单点峰值）、physical_peak（单个物理槽位的峰值负载）——用来定位「哪一步、哪一层」还在拖后腿。
• time=%.3fs——重排本身的耗时。前面说过重排发生在 decode 步之间，它的耗时直接计入服务时延，所以必须观测。

unflatten(-1, (-1, num_local_physical_experts)).sum(dim=-1) 是「按槽位分组还原到卡」的一步：因为 _batch_greedy_indices 返回的排列是「每连续 capacity 个元素属于同一张卡」，所以这里用同样的分组方式把 physical 维度还原成 rank 维度。分组约定在放置算法与指标计算之间必须一致，这是隐式的接口契约。

如果开了 DEEPINFER_SHOW_EPLB_DETAIL，还会用 tabulate 打印逐层的专家负载表：

1
headers = ["Layer", "MaxIdx", "Max/Avg"] + [f"#{expert_index}" for expert_index in range(num_experts)]

MaxIdx 是该层最热专家的编号，Max/Avg 是它的热度倍数。这两列在排障时最有用：MaxIdx 在连续几次重排里保持不变，说明这个专家一直是热点（该给它更多副本）；如果 MaxIdx 在不断跳变，说明流量分布本身在漂移，重排周期可能需要缩短。 把内部状态以这种形式暴露出来，是生产引擎的常见做法。

负载比值统计：另一条独立的观测通道

除了重排日志，还有一条按 DEEPINFER_MOE_LB_STAT_INTERVAL（默认 128 步）触发的观测：

1
def _update_load_balance_stats(self) -> None:
2
    stat_index = self._load_balance_stat_index
3
    torch.sum(
4
        self.expert_count,
5
        dim=-1,
6
        dtype=torch.long,
7
        out=self._local_tokens_per_step[stat_index],
8
    )
9
    stat_index += 1
10
    if stat_index < self._load_balance_stat_interval:
11
        self._load_balance_stat_index = stat_index
12
        return
13

14
    self._load_balance_stat_index = 0
15
    gathered_tokens = torch.zeros(
16
        self._ep_group.size(),
17
        self._load_balance_stat_interval,
18
        self.expert_count.shape[0],
19
        dtype=torch.long,
20
        device="cuda",
21
    )
22
    dist.all_gather_into_tensor(gathered_tokens, self._local_tokens_per_step, group=self._ep_group)
23
    tokens_per_rank = gathered_tokens.float()  # [ep_size, interval, num_moe_layers]
24
    self.moe_load_balance_ratio = (
25
        tokens_per_rank.mean(dim=0).sum().to("cpu", non_blocking=True),
26
        tokens_per_rank.max(dim=0).values.sum().to("cpu", non_blocking=True),
27
    )

这里刻意不按 expert 细分，而是把本卡的 expert 计数在 dim=-1 上求和，得到「本卡这段区间内处理了多少 token」。然后 all_gather_into_tensor 拿到全部 rank 的轨迹，算出两个标量：各 rank 平均 token 数的总和与各 rank 峰值的总和。它们的比值就是最粗粒度的负载均衡比。

值得注意的实现细节：

• expert_count 是 int，但这里是 torch.sum(..., dtype=torch.long) 显式指定累加类型，避免小整型溢出。
• all_gather 的目标缓冲显式 device="cuda"，用 NCCL 完成，然后 .to("cpu", non_blocking=True) 异步拷回。这是 ZOS 精神的延续（见 第二篇）：统计值只用于观测，绝不能让 CPU 同步阻塞在关键路径上，所以用非阻塞拷贝加上一个「pending」标记，由外部（ExpertLoadBalancerManager.take_moe_load_balance_ratio）在方便的时候取走。

1
def take_moe_load_balance_ratio(self) -> tuple[float, float] | None:
2
    if not self._pending_load_balance_ratios:
3
        return None
4
    ratios = self._pending_load_balance_ratios
5
    self._pending_load_balance_ratios = []
6
    return (
7
        sum(ratio[0].item() for ratio in ratios),
8
        sum(ratio[1].item() for ratio in ratios),
9
    )

.item() 会强制 CPU 同步——但它只在这个「取走」函数里出现，也就是由调用方选择时机。这是一个很典型的「把同步点从热路径移到观测路径」的写法。

还有 steps_until_lb_stat_flush()：返回距下次 flush 还有几步。这个函数存在的理由只有一个——让调度器知道下一步要不要做 all_gather。集合通信会与其他 rank 对齐，若调度器能在已知的 flush 步上做特殊处理（例如避开与 prefill 对齐），就能减少通信冲突。ExpertLoadBalancerManager 里还有一个断言：

1
countdowns = {lb.steps_until_lb_stat_flush() for lb in self._load_balancers} - {None}
2
assert len(countdowns) <= 1, f"各 balancer 的 LB flush 计数不同步：{countdowns}"

同一进程内所有 balancer 的 flush 倒计时必须一致，否则会在不同步上发起集合通信，直接死锁。这个断言把一条隐式的分布式约束写成了可检查的条件。

Manager：把多层接到多个 balancer 上

一篇模型里可能有几十层 MoE（主干 + DFlash），而不同 layout 的 balancer 是分开的。ExpertLoadBalancerManager 做的是层号到 balancer 的映射：

1
class ExpertLoadBalancerManager:
2
    def __init__(self, groups: list[tuple[ExpertLoadBalancer, list[int]]], num_moe_layers: int) -> None:
3
        layer_refs: dict[int, tuple[ExpertLoadBalancer, int]] = {}
4
        self._load_balancers: list[ExpertLoadBalancer] = []
5
        for load_balancer, global_layer_indices in groups:
6
            assert len(global_layer_indices) == load_balancer.expert_count.size(0)
7
            self._load_balancers.append(load_balancer)
8
            for local_layer_index, global_layer_index in enumerate(global_layer_indices):
9
                assert global_layer_index not in layer_refs
10
                layer_refs[global_layer_index] = (load_balancer, local_layer_index)
11
        assert len(layer_refs) == num_moe_layers

三个断言分别保证：层的数量对得上、层号不重复、全部层都被覆盖。

balancer 是怎么被建出来的：按 layout 同质性分组

ExpertLoadBalancerManager 的入参是「一组 balancer，每个 balancer 附带它负责的全局层号列表」。那么谁来决定「哪些层共用一个 balancer」？答案是 DeepseekForCausalLM._init_eplb：

1
grouped_moes: dict[tuple[object, ...], list[tuple[int, DeepseekMoE, list[torch.Tensor]]]] = {}
2
for global_layer_index, moe in enumerate(moes):
3
    weights = expert_comm_weights(moe)
4
    layout = (
5
        moe.num_logical_experts,
6
        moe.num_physical_experts,
7
        moe.num_routed_experts,
8
        moe.num_extra_experts,
9
        moe.ep_group,
10
        tuple((tuple(weight.shape), weight.dtype) for weight in weights),
11
    )
12
    grouped_moes.setdefault(layout, []).append((global_layer_index, moe, weights))

分组键 layout 是一个六元组：logical 数、physical 数、routed 数、extra 数、EP 组、以及每个权重张量的形状与 dtype。凡是这六项完全相同的层，共用一个 balancer；只要有一项不同，就分成两组。

为什么要按「权重形状」也分组？因为 ExpertLoadBalancer 在 shuffle 时按「一层一组张量」的假设复用接收缓冲：

1
# NOTE: 这里假设每一层的 hidden_size_i 相同，各层复用同一组接收缓冲
2
recv_buffers = [torch.empty_like(weight) for weight in experts_weight[0]]

recv_buffers 由第一层的权重形状决定，后续层直接复用。如果某个层的 gate_up_proj_weight 形状不同，复用就会出错。把形状放进分组键，等于把这个假设从「隐式约定」变成「显式保证」——分组之后，每个 balancer 管辖的层在形状上一定一致。

为什么要按 EP 组也分组？因为同一进程内可能存在多个不同的 EP 组（PG 里 ep_in_dp 与 ep 就是两组不同粒度的组），不同组上的集合通信范围不同，不能在同一个 balancer 里混。

再看 expert_comm_weights：

1
def expert_comm_weights(mlp: DeepseekMoE) -> list[torch.Tensor]:
2
    weights = [mlp.gate_up_proj_weight, mlp.down_proj_weight]
3
    if weights[0].element_size() == 1:
4
        weights += [w.sf.mT for w in weights]  # sf 存储时被 .mT 过，通信前转回
5
    return [w.mT if w.stride(-1) != 1 else w for w in weights]

这两行是分块布局与通信布局之间的转换。要点：

• element_size() == 1 判断 FP8（1 字节），此时除了权重本身还要搬 sf（scale factor）。量化专家的重排必须连 scale 一起搬，否则权重与 scale 会错位。w.sf.mT 中的 .mT 是「转置后的视图」，注释说明 sf 在存储时被 .mT 过，通信前要转回来。
• [w.mT if w.stride(-1) != 1 else w for w in weights]：如果最后一维不是连续（stride 不为 1），就转置一下。因为 shuffle 里做的是 weight[src] 这种按专家行的索引，而行内存连续时拷贝效率最高。

也就是说，EPLB 搬的不只是模型权重，还包括量化 scale；而且搬运前要把张量调整成适合按行拷贝的内存布局。这一段在 moe.py/linear.py 的量化实现里能看到对应的存储约定（ParallelRoutedLinear 的 use_fp8_gemm 分支），量化的整体链路可参考 TensorRT-LLM 源码完全拆解（十四）：量化全链路。

最后是构造与注册：

1
groups: list[tuple[eplb.ExpertLoadBalancer, list[int]]] = []
2
with torch.device("cuda"):
3
    for group in grouped_moes.values():
4
        example_moe = group[0][1]
5
        load_balancer = eplb.ExpertLoadBalancer(
6
            [weights for _, _, weights in group],
7
            len(group),
8
            example_moe.num_logical_experts,
9
            example_moe.num_physical_experts,
10
            example_moe.ep_group,
11
            example_moe.num_routed_experts,
12
            example_moe.num_extra_experts,
13
            is_normal_engine=cfg.engine_type == EngineType.NORMAL,
14
        )
15
        eplb.set_current(eplb.ExpertLoadBalancerManager(groups, len(moes)))
16
logger.info(f"EPLB initialized with {len(groups)} homogeneous layout group(s)")

几点：

• for group in grouped_moes.values() 取 group[0][1] 作为 example_moe——同组内的层在布局上完全一致，所以随便取一层的配置就够了。这是分组带来的直接好处。
• with torch.device("cuda") 包住构造，因为 ExpertLoadBalancer.__init__ 里创建了 expert_count、_expert_count_window 等张量，默认要落在 GPU 上（这些张量在热路径上被 dispatch 直接读写）。
• is_normal_engine=cfg.engine_type == EngineType.NORMAL 就是前面那两组 flag 的来源。
• 日志 EPLB initialized with {len(groups)} homogeneous layout group(s) 的措辞值得注意：“homogeneous layout group” 正是这套分组机制的自我描述。如果启动日志里这个数字是 1，说明全部 MoE 层布局一致（最常见的情形）；如果是 2，说明主干与 DFlash 的布局不同。
• __del__ 里调用 eplb.destroy_current()，与 set_current 的「不能重复初始化」断言配对，把单例的生命周期管理收在模型对象上。layer(layer_index) 返回一个 ExpertLoadBalancerLayerState NamedTuple，把「本层视角」的三样东西（logical_to_physical_map[local_layer_index]、logical_count[...]、expert_count[...]）暴露给 dispatch 与路由：

1
class ExpertLoadBalancerLayerState(NamedTuple):
2
    logical_to_physical_map: torch.Tensor
3
    logical_count: torch.Tensor
4
    expert_count: torch.Tensor

这三个张量正是热路径上真正要用到的全部信息：路由要知道「这个 logical expert 有哪些副本」（logical_to_physical_map）与「每个副本现在多热」（expert_count，配合 logical_count 可以做副本级的选择），dispatch 要把计数写回 expert_count。用一个 NamedTuple 打包，避免热路径上反复做字典查找和层号映射。全局只有一份 _current_load_balancer：

1
_current_load_balancer: ExpertLoadBalancerManager | None = None
2

3
def current() -> ExpertLoadBalancerManager | None:
4
    """当前进程的 EPLB manager，未启用（如 dense 模型）时为 None。"""
5
    return _current_load_balancer
6

7
def set_current(load_balancer: ExpertLoadBalancerManager) -> None:
8
    global _current_load_balancer
9
    assert _current_load_balancer is None, "EPLB is already initialized"

set_current 的断言「不能重复初始化」与 current() 允许返回 None（dense 模型没有 MoE 层，自然没有 balancer）合起来说明：这是进程级单例，且是可选组件。MoE 相关代码在访问前必须判空。

六种进程组

1
def create_process_groups(self, *, tag, dp_size, ep_size=1, cp_size=1, enp_size=1) -> None:
2
    self.tag = tag
3
    ranks = torch.arange(max(enp_size, dp_size * cp_size))
4
    ...
5
    # World group
6
    self.world = _split_group([ranks.tolist()], ...)
7
    # Local group
8
    self.local = _split_group(ranks.unsqueeze(-1).tolist(), ...)
9

10
    # DP group
11
    dp_ranks = rearrange(ranks, "(replica dp cp) -> (replica cp) dp", dp=dp_size, cp=cp_size)
12
    self.dp = _split_group(dp_ranks.tolist(), ...)
13
    self.cpu_dp_group = _split_group(dp_ranks.tolist(), backend="gloo")
14

15
    # CP group
16
    cp_ranks = rearrange(ranks, "(dp cp) -> dp cp", cp=cp_size)
17
    self.cp = _split_group(cp_ranks.tolist(), ...)
18
    self.cpu_cp_group = _split_group(cp_ranks.tolist(), backend="gloo")
19

20
    # EP groups: contiguous groups of ep_size ranks.
21
    if ep_size == ranks.numel():
22
        self.ep = self.world
23
    else:
24
        ep_ranks = rearrange(ranks, "(edp ep) -> edp ep", ep=ep_size)
25
        self.ep = _split_group(ep_ranks.tolist(), ...)
26

27
    # Engram groups
28
    enp_ranks = rearrange(ranks, "(replica enp) -> replica enp", enp=enp_size)
29
    self.enp = _split_group(enp_ranks.tolist(), ...)
30

31
    # EP-in-DP groups: contiguous groups of ep_size ranks
32
    ep_in_dp_ranks = rearrange(ranks, "(edp ep_in_dp) -> edp ep_in_dp", ep_in_dp=ep_size)
33
    if ep_size == 1:
34
        self.ep_in_dp = self.local
35
    else:
36
        self.ep_in_dp = _split_group(ep_in_dp_ranks.tolist(), ...)
37
    self.cpu_ep_in_dp = _split_group(ep_in_dp_ranks.tolist(), backend="gloo")

用 einops.rearrange 表达 rank 分组是这套代码的写法特点。把 rank 编号当成一个多维索引，"(replica dp cp) -> (replica cp) dp" 的意思就是：把 rank 向量重排成 (replica, dp, cp) 的坐标，然后按新的维度边界重新分组，得到「同一 (replica, cp)、不同 dp」的组——也就是 DP 组。

几处需要留意：

• EP 组是连续 rank 块（"(edp ep) -> edp ep"）。为什么用连续块而不是交错？因为连续块更容易落在同一个物理节点/NVLink 域内——_intranode_shuffle 的判断 PG.is_within_nvl_domain(ep_group) 正是依赖这一点。EP 组的拓扑形状是 IPC shuffle 能不能用的前提。
• ep_size == ranks.numel() 时直接复用 world 组，不新建 ProcessGroup。NCCL 的 communicator 创建是昂贵的（要建通道、交换句柄），能省则省。这条分支在「整个作业就是一个 EP 组」时命中。
• CPU 组用 gloo 后端（cpu_dp_group、cpu_cp_group、cpu_ep_in_dp）。GPU 组只搬张量，CPU 组搬的是元数据与对象。CP 组还额外建了一个共享内存广播队列：

1
# MessageQueue for cp_broadcast_object, only needed when the CP group is non-trivial
2
if cp_size > 1:
3
    self.mq = MessageQueue.create_from_process_group(self.cpu_cp_group, 1 << 20, 10)

1 << 20 是 1 MiB 缓冲、10 是队列深度。cp_broadcast_object 用它广播任意 Python 对象（不能走 NCCL），实现是「writer enqueue / reader dequeue」：

1
def cp_broadcast_object(obj: Any) -> Any:
2
    assert PG.mq is not None
3
    if PG.mq._is_writer:
4
        PG.mq.enqueue(obj)
5
        return obj
6
    else:
7
        return PG.mq.dequeue()

一个容易踩的坑：lazy 初始化

1
def init_device_groups(self) -> None:
2
    """默认 PG 改为 CPU(gloo) 后端后，NCCL/HCCL 子组失去 eager/split 初始化，全部 lazy：必须先跑一次 collective 才会真正建 comm。
3

4
    单 rank 组也不能跳过：engram deep_ep 经 `backend._comm_ptr()` 复用这里建出的 comm，空指针会直接段错误。
5
    """
6
    seen: set[int] = set()
7
    for value in self.__dict__.values():
8
        try:
9
            if not isinstance(value, ProcessGroup) or dist.get_backend(value) not in ("nccl", "hccl"):
10
                continue
11
        except Exception:
12
            continue
13
        if id(value) not in seen:
14
            seen.add(id(value))
15
            dist.all_reduce(torch.zeros(1).cuda(), group=value)

这段代码值得单独讲，因为它是分布式推理里最经典的坑之一：PyTorch 的 dist.new_group 不保证立刻建立通信资源，NCCL communicator 是懒创建的——第一次在这个组上发起集合通信时才真正创建。如果某个组直到 EPLB 重排时才第一次被使用，那么这次重排会额外付出建 comm 的代价（可能几百毫秒），而且这段延迟正好落在服务的关键路径上。

解决办法是在初始化阶段对所有 NCCL/HCCL 组各跑一次 all_reduce（用 1 个元素的零张量，代价极小），把创建成本前移到启动阶段。注释里的第二句更关键：单 rank 组也不能跳过，因为 DeepEP 会通过 backend._comm_ptr() 复用这里建出的 communicator，如果没有先建，拿到的是空指针，直接段错误。这是一个「宁可多跑一次无用的 all_reduce，也不能留一个未初始化的 comm」的取舍。

遍历用 self.__dict__.values() 加 id() 去重，是因为同一个 ProcessGroup 可能被多个属性引用（例如 ep_size == ranks.numel() 时 self.ep is self.world），重复 all_reduce 虽然无害但没必要。

is_within_nvl_domain

1
@staticmethod
2
def is_within_nvl_domain(group: ProcessGroup) -> bool:
3
    if PLATFORM.is_npu:
4
        return True
5
    return group.size() == 1 or group.size() <= deep_ep.get_physical_domain_size(group)[1]

这个判断决定 EPLB 走 IPC 还是 P2P：

• NPU 直接返回 True。但注意 _initialize_shuffle 里的平台判断是 PG.is_within_nvl_domain(ep_group) and PLATFORM.is_gpu——两个条件与起来，NPU 实际走的是 P2P 路径。is_within_nvl_domain 在 NPU 上返回 True 只是表示「拓扑上是一体的」，而 symm_mem 这条实现路径不支持 NPU（代码里有 TODO(NPU): NPU doesn't support symm_mem yet.）。这是一个「拓扑判断」与「实现可用性判断」分离的例子，读代码时容易误判。
• GPU 上查 deep_ep.get_physical_domain_size(group)，返回的第二项是这个组覆盖的物理域大小（NVLink 域内的 rank 数）。组大小不超过它，说明整组在一个 NVLink 域内，可以走 IPC。

这条判断是运行时拓扑探测，不是配置项——用户不用声明「我们在一台机器上」，由 DeepEP 探测实际硬件的域边界。这也解释了为什么这个函数定义在 parallel_state.py 而不是 EPLB 里：拓扑信息属于并行状态，EPLB 只是消费者。

EP 组上的通信原语

communication_op.py 是 vLLM 派生的文件（文件头有 Derived from vLLM; modified by DeepSeek.），它提供的是CP 组上的广播而非 EP：

1
def cp_broadcast(input_: torch.Tensor, src: int = 0) -> torch.Tensor:
2
    """Broadcast input_ in place across the CP group."""
3
    cp_group = PG.context_parallel_group()
4
    if cp_group.size() > 1:
5
        dist.broadcast(input_, src=dist.get_global_rank(cp_group, src), group=cp_group)
6
    return input_

以及一个张量字典广播：

1
def broadcast_tensor_dict(tensor_dict=None, src: int = 0) -> dict[Any, Any]:
2
    """Broadcast a dict across the CP group: cuda tensors via NCCL, everything else via shm object broadcast."""

实现思路和前面说的一致：CUDA 张量走 NCCL（先用 shm 广播元数据 TensorMetadata，再 cp_broadcast_multi_tensor 批量广播张量），非张量对象走共享内存对象广播。cp_broadcast_multi_tensor 里有一行 TODO：