Improving DeepEP MoE Load Balance in SGLang with Waterfill and LPLB

发表时间: 2026-06 · Blog post by NVIDIA Team (lmsys.org)

原文: https://lmsys.org/blog/2026-06-26-waterfill-lplb

作者/机构：NVIDIA Team

速读

一句话结论 本文在 SGLang 中引入了 Waterfill 和 LPLB 两种运行时的负载均衡算法，分别针对 MoE 模型的共享专家和冗余路由专家进行动态调度，在不改变模型语义的前提下，将 DeepSeek 系列模型的推理吞吐量提升了最高 7.34%。

要解决什么问题 混合专家（MoE）模型在推理时依赖专家并行（EP）将专家分布到多张 GPU 上，但原有的静态专家放置策略（如 EPLB）卡在无法处理运行时的动态负载不均衡。具体机制在于，路由器对 token 的分发天然是不均匀的，即使离线校准阶段做好了专家的均匀放置，实际线上单个 batch 的 token 依然可能高度集中于少数几个逻辑专家。这导致拥有这些专家的 GPU 节点负载暴增，而整个 EP 通信组必须等待最慢的节点完成计算。此外，MoE 模型中既有稀疏的路由专家，又有每个 token 都必须经过的密集型共享专家，如果强制在本地计算共享专家或仅对冗余专家做平均流量拆分，会进一步加剧节点间的算力闲忙差异，导致严重的计算与通信长尾效应。

怎么做的 作者提出了两种互补的调度期负载均衡算法，核心思路是在不改变模型逻辑路由选择的前提下，动态决定哪个物理 GPU 来执行具体的专家计算。第一部分是针对共享专家的 Waterfill 算法。由于所有 token 都要算共享专家，如果都在本地算，重负载节点会雪上加霜。Waterfill 将共享专家视为一个可调度的槽位，像注水填谷一样，把共享专家的计算任务分配给当前负载较低的节点。具体而言，算法先统计各节点已有的路由专家负载 $L_r$，然后加上共享专家槽位总数 $N$，计算出目标水位线：

$$H = \left\lceil\frac{\sum_r L_r + N}{R}\right\rceil$$

其中 $R$ 是 EP 组大小。低于此水位线的节点拥有空闲余量 $S_r = \max(H - L_r, 0)$。随后，算法按空闲余量的比例，将 token 的共享专家任务概率采样发送给候选节点。为了避免增加额外的 all-to-all 通信开销，候选节点通常限制在 token 已经需要访问的路由专家节点或本地节点。此方法依赖于将共享专家与路由专家在 DeepEP 底层融合为统一的调度布局。第二部分是针对冗余路由专家的 LPLB 算法。EPLB 会为热门逻辑专家创建多个物理副本，默认做法是将 token 平均分给这些副本，但这在真实 batch 分布偏离离线校准时依然会不均衡。LPLB 将其转化为一个每层求解的线性规划（LP）问题：目标是最小化所有节点中的最大负载 $M$。约束条件包括：每个节点的负载（冗余专家分配量加单副本专家固定量加距离峰值的松弛变量）等于 $M$；且每个逻辑专家的分配量之和必须等于其真实接收到的 token 总数。在实现上，所有节点先进行一次 all-reduce 汇总全局 token 分布，然后各自独立在 GPU 上利用预编译的内点法 kernel 求解该 LP，得出最优的物理副本流量拆分比例，最后按此概率分布进行 token 调度。

效果如何 实验在两台 Hopper 架构 GPU 节点（共 16 张卡）上搭建，采用 TP16、DP16、EP16 的并行配置。测试模型为 DeepSeek-V3/R1（FP8）以及 DeepSeek V4 Flash（FP8），评测数据源自 MMLU、GPQA 和 GSM8K 的 prompt 池。对比基线包含两种路线：一是无 EPLB 的纯动态路由基线；二是代表静态放置路线的 Static EPLB 基线，并细分为无冗余专家（red0）、16 个冗余专家（red16）和 32 个冗余专家（red32）三种设置。量化结果显示，在 DeepSeek-V3/R1 负载下，Waterfill 能够稳定地将共享专家任务转移到轻载节点，在所有基线对比中使总吞吐量提升了 1.48% 到 4.66%。在 DeepSeek V4 上，最优配置下的吞吐量从 49,253 tok/s 提升至 51,677 tok/s（提升 4.92%）。LPLB 则在配置了冗余专家的场景下（red16 和 red32）展现出优势，吞吐量提升了 0.84% 到 7.34%。作者也明确指出了方法的局限与代价：Waterfill 为了控制通信成本，在候选节点的选择上做出了妥协，优先复用已有的通信路径；而 LPLB 的代价是每层需要一次 all-reduce 通信和一次 LP 求解。更重要的是，LPLB 存在明确的失效场景，当系统没有配置冗余专家（如 red0 基线）时，LPLB 没有任何可供调度的物理副本，此时算法不仅无法提升吞吐，其求解路径反而会变成纯粹的额外开销，导致性能出现 0.95% 到 1.61% 的微小下降。

主要贡献

混合专家（Mixture-of-Experts, MoE）模型依赖于专家并行（Expert Parallelism, EP）在多个GPU上扩展推理能力。在SGLang中，DeepEP和EPLB在EP下提供了高性能的服务，但由于token并没有均匀地路由到各个专家，每个rank所面临的工作负载仍然可能存在不平衡。

本文介绍了SGLang中两个在调度时（dispatch-time）进行负载均衡的特性，旨在解决静态放置方法无法消除的残余不平衡问题：

背景知识与设计原则

DeepEP通过为专家并行提供优化的token调度和合并内核（combine kernels）来加速MoE推理。在典型的DeepSeek风格的MoE层中，每个token被路由到由模型路由器选择的几个路由专家（routed experts）。一些模型还包含一个共享专家（shared expert），该专家应用于每个token。

从服务系统的角度来看，路由专家和共享专家产生了不同的负载模式：

静态专家放置虽然有用，但无法消除所有的运行时不平衡。一个批次中的实际token可能仍然集中在某些专家或rank的子集上。在DeepEP中，这可能导致一些rank等待超载的对等节点。Waterfill和LPLB的设计原则都是在保留模型语义的前提下，减少这种调度时的不平衡。

方法细节

Waterfill：针对共享专家调度的轻量级负载均衡

传统共享专家计算的局限性
如果共享专家总是在每个rank上本地计算，那么无论该rank是否已经被路由专家（routed experts）超载，它都必须承担共享专家的计算开销。这导致超载的rank继续超载，而负载较轻的rank无法帮助分担共享专家的工作。

Waterfill的调度策略核心思想
Waterfill通过将共享专家视为一个可调度的专家槽位（dispatchable expert slot）来改变这一现状。在选择完路由专家后，Waterfill会估计每个EP rank上当前的路由负载，然后将共享专家的工作分配给负载较低的rank。从概念上讲，它就像向不平的容器中注水一样，填补rank负载分布中的“低谷”。

保持模型语义不变的分配机制
对于每个token，Waterfill为共享专家添加一个额外的专家槽位。它不再总是将该槽位分配给token的本地rank，而是根据当前的负载分布选择一个rank。这保持了路由专家的选择不变，因此模型仍然计算相同的逻辑路由专家和相同的共享专家。唯一改变的是由哪个物理rank执行共享专家的工作。

Waterfill算法执行步骤
从高层次来看，该算法的过程如下：

  1. 统计已经分配到每个EP rank上的路由专家负载。
  2. 将该统计值作为每个rank的负载分数（load score）。在动态模式下，SGLang首先运行一个EP组的集合通信（collective），因此该分数可以使用全局路由负载向量加上每个rank当前的本地批处理大小（batch size）。

  3. 为每个参与的token添加一个共享专家槽位，并计算目标水位线（target waterline）：
    $H = \left\lceil\frac{\sum_r L_r + N}{R}\right\rceil$
    其中 $L_r$ 是rank $r$ 的负载分数，$N$ 是需要放置的共享专家槽位数量，$R$ 是EP组的大小（EP group size）。

  4. 低于此水位线的rank具有空闲容量（slack）：
    $S_r = \max(H - L_r, 0)$

  5. 对于每个token，Waterfill根据与空闲容量成正比的概率，从候选rank中采样共享专家的目标rank，并带有微小的本地rank偏好。如果所有候选rank的空闲容量均为零，则回退到明显较轻的候选rank，同样保持本地rank偏好。

通信权衡与模式选择
这里存在一个重要的通信权衡。如果每个token都能将其共享专家的工作发送到任何EP rank，Waterfill将获得更大的平衡自由度，但这也可能增加all-to-all的通信流量。对于GPU MoE服务而言，通信通常比额外的共享专家计算更昂贵。因此，保守通信的候选集（communication-conservative candidate set）将共享专家保留在token已经为路由专家访问过的rank上，并将源rank作为回退选项。SGLang还支持全rank模式（all-rank mode），这为Waterfill提供了更多的平衡自由度，但可能会增加一个新的每token调度目的地。这是一个经过深思熟虑的通信权衡，而不是模型语义的改变。

Waterfill的最终效果
通过将共享专家的工作从已经繁重的rank转移到较轻的rank，Waterfill平衡了每个rank的工作量并提高了端到端的吞吐量（详见Fig 1）。

Fig 1: 共享专家负载均衡前后的Waterfill时间线。Waterfill在保持路由专家选择不变的情况下，将共享专家的工作从超载的rank转移到较轻的rank，从而在不改变模型语义的情况下缩短了最慢的MoE层路径。
Fig 1: 共享专家负载均衡前后的Waterfill时间线。Waterfill在保持路由专家选择不变的情况下，将共享专家的工作从超载的rank转移到较轻的rank，从而在不改变模型语义的情况下缩短了最慢的MoE层路径。

共享专家融合作为启用机制
Waterfill可以通过融合共享专家和路由专家得到进一步加速。
在EP下，共享专家曾使用与路由专家独立的执行路径。在Waterfill选择非本地的共享专家rank后，该设计需要从已调度的路由专家布局中提取共享专家token，并启动独立的共享专家计算，这增加了额外的布局转换和启动开销。
共享专家融合（Shared expert fusion）通过将共享专家表示为同一DeepEP MoE布局中的另一个专家槽位来避免上述独立路径。在DeepSeek V3/R1中，路由器仍然选择原始的路由top-k专家，并且TopK输出获得一个额外的列用于共享专家。在DeepEP物理专家ID布局中，每个rank在其路由专家旁边保留一个额外的共享专家槽位。这使得路由专家和共享专家能够共享相同的DeepEP调度、分组GEMM（grouped-GEMM）和合并流程（combine flow）。
融合本身并不是最终的负载平衡算法。它是使共享专家调度对DeepEP可见并因此可由Waterfill控制的必要机制。

LPLB：针对冗余专家副本的基于LP的负载均衡

LPLB解决的核心问题
EPLB放置热点逻辑专家的冗余副本，然后默认情况下，将其热点专家的token均匀分配到其物理副本上。只有当用于构建放置的离线分布与实时流量匹配时，均匀分割才是最佳的。在实践中，情况往往并非如此：单个批次集中的专家可能与校准集不同，服务的数据集偏离记录的数据集，并且重新平衡周期足够长，以至于放置在许多批次中实际上是静态的。当发生这种情况时，均匀划分热点专家的负载仍然会使拥有其副本的rank相对于EP组的其余部分负载不均，整个组都要等待最繁忙的rank。
LPLB在调度时（dispatch time）弥补了这一差距。对于每个MoE层，在每个批次上，它查看实际的每个专家的token计数，并决定如何在复制的专家的物理副本之间分配token，以便最小化最大单rank负载。它不移动权重，也不改变路由器的逻辑top-k选择——它只在逻辑专家的有效物理副本中选择每个副本接收多少流量。其结果是针对当前批次的最佳最小-最大分配（min-max assignment），而不是EPLB离线固化的静态均匀分割。

线性规划（LP）公式化构建
LPLB将其转化为每层求解的小型线性规划问题。直觉直接映射到以下约束条件上：

决策变量与矩阵规模优化
决策变量是复制专家的每副本负载，加上每rank的空闲容量和 M。单副本专家不是变量——它们只贡献固定项——这使得LP保持较小规模：其大小随冗余专家的数量和rank的数量扩展，而不是随全部专家数量扩展。
约束矩阵分为离线部分和在线部分。结构块——副本到逻辑专家的映射、冗余副本的每rank所有权以及空闲容量/-M列——仅依赖于专家到GPU的放置，因此它们在启动时和每次EPLB重新平衡后预先计算一次。只有等式右侧随批次变化：观察到的冗余专家负载和每rank的单副本负载。一个Big-M辅助列在求解期间保持系统可行，并在目标函数中受到重罚，以便求解器将其驱动为零。

从全局计数到LP求解的工程实现
DP-attention的一个微妙之处在于，不同的EP rank在同一步骤中运行不同的前向模式——预填充（prefill）、解码（decode）或空闲（idle）——因此没有单个rank能看到全局token分布。LPLB通过一个刻意简单的集合通信设计来处理这个问题：

  1. 每个rank统计其本地每个逻辑专家的token数。
  2. 所有EP rank参与一次这些计数的all-reduce操作——空闲rank贡献零——因此每个rank最终获得相同的全局每专家分布。
  3. 每个rank然后从这些相同的输入独立求解相同的LP，并获得相同的解决方案，因此不需要广播结果。

LP本身在GPU上由基于 cuSOLVERDx/cuBLASDx 的融合内点法（IPM）内核求解，在启动时针对该层的矩阵形状进行预编译，因此第一个真实请求不需要支付JIT开销。整个每批次路径——构建右侧、求解和提取每副本分割——折叠为三个CUDA内核启动（kernel launches），它们写入预先分配的缓冲区，使启动开销和主机同步脱离关键路径。

从LP解决方案到Token调度
LP返回对于每个复制的逻辑专家，其负载应该如何在其物理副本中划分。LPLB将其归一化为有效物理副本上的每专家概率分布（log2phy_prob）。在调度时，路由到复制的逻辑专家的每个token从该分布中采样一个物理副本；单副本专家像以前一样映射到它们唯一的物理位置。这是现有 dynamic 策略的直接替代品，后者均匀随机地选择一个副本——LPLB保持相同的概率性、每token调度形状，但用为该批次计算的负载最优分布替换了均匀抽取（详见Fig 2）。

Fig 2: LPLB在冗余专家副本间的感知负载拆分。LPLB在不改变逻辑专家路由的情况下，将复制专家的流量推向较轻的rank，因此当提供冗余副本时，相同的选定专家可以在负载较轻的物理副本上完成计算。
Fig 2: LPLB在冗余专家副本间的感知负载拆分。LPLB在不改变逻辑专家路由的情况下，将复制专家的流量推向较轻的rank，因此当提供冗余副本时，相同的选定专家可以在负载较轻的物理副本上完成计算。

Waterfill与LPLB的区别与互补性

Waterfill和LPLB共享最终目标——在DeepEP下展平每rank负载——但作用于不同的调度选择，并具有不同的机制：

它们是互补的而不是竞争的：Waterfill消除了由密集共享专家贡献的不平衡，而LPLB消除了稀疏路由副本之间的不平衡。因为LPLB只在相同逻辑专家的有效副本之间重新分配流量，并且从不改变路由器的逻辑top-k，它与Waterfill一样保留了模型语义。

LPLB的最佳适用场景
LPLB的收益取决于实时批次与EPLB校准分布的偏差程度。当流量均衡且批次巨大（大规模、高度多样化的服务）时，LP几乎没有剩余的不平衡可以消除。当流量基本不变且狭窄（少数几乎相同的问题）时，静态EPLB已经捕获了分布，均匀分割接近最佳。LPLB在两者之间提供最强的信号——中等规模服务（专注于中等数量相关主题），其中每个批次的不平衡方式是离线放置未预料到的，但仍然具有足够的结构，使得最佳的每批次分割能够有意义地降低峰值rank负载。

实验环境

实验结果

1. DeepSeek V3/R1上的Waterfill和LPLB评估

本组实验在MMLU, GPQA, 和 GSM8K数据集上对比了无EPLB、以及静态EPLB（包含red0, red16, red32不同冗余度）情况下的基线吞吐量、Waterfill吞吐量及其增益、LPLB吞吐量及其增益。

实验结果与分析：

Fig 3: MMLU、GPQA和GSM8K数据集上的Waterfill吞吐量。通过将共享专家的工作转移到负载较轻的EP rank，Waterfill在匹配的基线上一致地提高了总吞吐量。
Fig 3: MMLU、GPQA和GSM8K数据集上的Waterfill吞吐量。通过将共享专家的工作转移到负载较轻的EP rank，Waterfill在匹配的基线上一致地提高了总吞吐量。
Fig 4: MMLU、GPQA和GSM8K数据集上的LPLB吞吐量。当存在冗余专家副本（red16/red32）时，LPLB提高了吞吐量，因为LP可以在物理副本之间进行选择。在red0下，没有冗余副本可供重新平衡，因此算法无法改善调度，其all-reduce/求解路径仅作为开销出现。
Fig 4: MMLU、GPQA和GSM8K数据集上的LPLB吞吐量。当存在冗余专家副本（red16/red32）时，LPLB提高了吞吐量，因为LP可以在物理副本之间进行选择。在red0下，没有冗余副本可供重新平衡，因此算法无法改善调度，其all-reduce/求解路径仅作为开销出现。

2. DeepSeek V4上的Waterfill评估

DeepSeek V4可以使用 HashTopK 路由路径，Waterfill必须在该输出路径中追加并重新映射共享专家槽位（通过PR #25391扩展）。本实验在MMLU风格的服务工作负载上验证V4 Flash FP8的性能。

实验结果与分析：

Fig 5: DeepSeek V4 Flash上的Waterfill吞吐量。Waterfill在DeepSeek V4 Flash上同样有效，在无EPLB和静态EPLB设置下提高了吞吐量，增益从+3.28%到+4.92%不等。
Fig 5: DeepSeek V4 Flash上的Waterfill吞吐量。Waterfill在DeepSeek V4 Flash上同样有效，在无EPLB和静态EPLB设置下提高了吞吐量，增益从+3.28%到+4.92%不等。

结论

本文介绍的Waterfill和LPLB两种调度时负载均衡方法，有效解决了MoE模型在推理服务阶段由于token路由不均匀导致的残余负载不平衡问题。实验证明，这两种方法在DeepSeek-V3/R1及V4模型上均能显著提升系统吞吐量。

精度与模型语义验证：

补充细节：如何使用

启用Waterfill

Waterfill通过DeepEP MoE路径启用。一个典型的启动命令如下：

python3 -m sglang.launch_server \
    --model-path /path/to/DeepSeek-V3 \
    --tp 16 \
    --dp-size 16 \
    --nnodes 2 \
    --node-rank ${NODE_RANK} \
    --dist-init-addr ${HEAD_NODE_IP}:${PORT} \
    --host 0.0.0.0 \
    --port 30000 \
    --trust-remote-code \
    --moe-a2a-backend deepep \
    --deepep-mode normal \
    --enable-dp-attention \
    --enable-deepep-waterfill \
    --init-expert-location /path/to/expert_distribution.pt

关键标志说明：

启用LPLB

LPLB通过DeepEP MoE路径上的EP调度算法选择。由于LPLB在冗余副本之间平衡token，它需要一个实际包含冗余专家的EPLB放置。一个典型的双节点启动命令如下：

python3 -m sglang.launch_server \
    --model-path /path/to/DeepSeek-R1 \
    --tp 16 \
    --dp-size 16 \
    --ep-size 16 \
    --nnodes 2 \
    --node-rank ${NODE_RANK} \
    --dist-init-addr ${HEAD_NODE_IP}:${PORT} \
    --host 0.0.0.0 \
    --port 30000 \
    --trust-remote-code \
    --moe-a2a-backend deepep \
    --deepep-mode normal \
    --enable-dp-attention \
    --ep-num-redundant-experts 16 \
    --ep-dispatch-algorithm lp \
    --init-expert-location /path/to/expert_stats.pt

关键标志说明：

附录：致谢与相关PR

这项工作建立在SGLang DeepEP和MoE服务栈的基础之上，并通过SGLang项目中的社区协作开发完成。

相关代码合并请求（PRs）：

贡献者：

特别感谢DeepSeek在 deepseek-ai/LPLB 开源了他们的LPLB工作，其用于在冗余专家副本之间平衡token的线性规划公式启发了本文描述的SGLang LPLB集成。