Improving DeepEP MoE Load Balance in SGLang with Waterfill and LPLB
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)进行负载均衡的特性,旨在解决静态放置方法无法消除的残余不平衡问题:
- Waterfill:一种针对DeepEP的轻量级共享专家(shared-expert)负载均衡方法。它通过DeepEP调度共享专家,并将其分配给负载较轻的rank。在两个Hopper GPU节点上运行DeepSeek-V3/R1风格的服务工作负载时,Waterfill在MMLU、GPQA和GSM8K数据集上将总吞吐量提高了 +1.48% 到 +4.66%。在DeepSeek V4上,测得的最佳性能点从 49,253 tok/s 提升至 51,677 tok/s(+4.92%)。
- LPLB:一种基于线性规划(linear-programming)的负载均衡器,专为冗余专家副本设计。它解决了一个针对冗余专家的逐层调度优化问题。在相同的两个Hopper GPU节点上结合冗余EPLB放置策略,LPLB在MMLU、GPQA和GSM8K数据集上将总吞吐量提高了 +0.84% 到 +7.34%。
背景知识与设计原则
DeepEP通过为专家并行提供优化的token调度和合并内核(combine kernels)来加速MoE推理。在典型的DeepSeek风格的MoE层中,每个token被路由到由模型路由器选择的几个路由专家(routed experts)。一些模型还包含一个共享专家(shared expert),该专家应用于每个token。
从服务系统的角度来看,路由专家和共享专家产生了不同的负载模式:
- 路由专家是稀疏的:不同的token选择不同的专家,因此它们的负载取决于路由器的分布。
- 共享专家是密集的:每个token都需要共享专家,因此共享专家的工作负载在整个批次中都存在。
- 冗余专家(Redundant experts):由EPLB风格的放置策略引入,为某些逻辑专家提供多个物理副本。它们创造了调度时平衡的机会,因为系统可以选择由哪个物理副本处理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算法执行步骤
从高层次来看,该算法的过程如下:
- 统计已经分配到每个EP rank上的路由专家负载。
- 将该统计值作为每个rank的负载分数(load score)。在动态模式下,SGLang首先运行一个EP组的集合通信(collective),因此该分数可以使用全局路由负载向量加上每个rank当前的本地批处理大小(batch size)。
-
为每个参与的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)。 -
低于此水位线的rank具有空闲容量(slack):
$S_r = \max(H - L_r, 0)$ -
对于每个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)。
共享专家融合作为启用机制
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将其转化为每层求解的小型线性规划问题。直觉直接映射到以下约束条件上:
- 目标——最小化峰值:引入一个标量
M,代表所有rank上的最大负载,并最小化它。降低M会将最繁忙的rank拉向平均值,这正是缩短由EP不平衡引起的分组GEMM尾部延迟的方法。 - Rank负载约束:对于每个rank,其约束为:(来自其冗余专家副本的负载) + (来自其单副本专家的负载) + (到峰值的空闲容量) = M。每个rank上的单副本负载是固定的输入——这些专家没有调度选择。每个rank得到一个这样的方程;空闲容量(slack)是非负的,因此
M被强制至少等于每个rank的真实负载。 - 冗余专家守恒:对于每个复制的逻辑专家,分配给其副本的负载总和必须等于该专家观察到的总负载:$x_1 + x_2 + ... + x_n = L$,其中 $x_i$ 是放置在副本 $i$ 上的负载,$L$ 是该专家观察到的总负载。这保证了LPLB只重新分配现有的流量,绝不凭空创造或丢弃token。
决策变量与矩阵规模优化
决策变量是复制专家的每副本负载,加上每rank的空闲容量和 M。单副本专家不是变量——它们只贡献固定项——这使得LP保持较小规模:其大小随冗余专家的数量和rank的数量扩展,而不是随全部专家数量扩展。
约束矩阵分为离线部分和在线部分。结构块——副本到逻辑专家的映射、冗余副本的每rank所有权以及空闲容量/-M列——仅依赖于专家到GPU的放置,因此它们在启动时和每次EPLB重新平衡后预先计算一次。只有等式右侧随批次变化:观察到的冗余专家负载和每rank的单副本负载。一个Big-M辅助列在求解期间保持系统可行,并在目标函数中受到重罚,以便求解器将其驱动为零。
从全局计数到LP求解的工程实现
DP-attention的一个微妙之处在于,不同的EP rank在同一步骤中运行不同的前向模式——预填充(prefill)、解码(decode)或空闲(idle)——因此没有单个rank能看到全局token分布。LPLB通过一个刻意简单的集合通信设计来处理这个问题:
- 每个rank统计其本地每个逻辑专家的token数。
- 所有EP rank参与一次这些计数的all-reduce操作——空闲rank贡献零——因此每个rank最终获得相同的全局每专家分布。
- 每个rank然后从这些相同的输入独立求解相同的LP,并获得相同的解决方案,因此不需要广播结果。
LP本身在GPU上由基于 cuSOLVERDx/cuBLASDx 的融合内点法(IPM)内核求解,在启动时针对该层的矩阵形状进行预编译,因此第一个真实请求不需要支付JIT开销。整个每批次路径——构建右侧、求解和提取每副本分割——折叠为三个CUDA内核启动(kernel launches),它们写入预先分配的缓冲区,使启动开销和主机同步脱离关键路径。
从LP解决方案到Token调度
LP返回对于每个复制的逻辑专家,其负载应该如何在其物理副本中划分。LPLB将其归一化为有效物理副本上的每专家概率分布(log2phy_prob)。在调度时,路由到复制的逻辑专家的每个token从该分布中采样一个物理副本;单副本专家像以前一样映射到它们唯一的物理位置。这是现有 dynamic 策略的直接替代品,后者均匀随机地选择一个副本——LPLB保持相同的概率性、每token调度形状,但用为该批次计算的负载最优分布替换了均匀抽取(详见Fig 2)。
Waterfill与LPLB的区别与互补性
Waterfill和LPLB共享最终目标——在DeepEP下展平每rank负载——但作用于不同的调度选择,并具有不同的机制:
- 目标:Waterfill针对应用于每个token的共享(密集)专家;LPLB针对EPLB已复制的路由专家。
- 决策:Waterfill决定哪个rank执行每个token的共享专家槽位;LPLB决定如何将每个复制专家的token拆分到其物理副本上。
- 方法:Waterfill是基于当前rank负载的轻量级填谷启发式算法;LPLB是每层在GPU上求解的最小-最大线性规划。
- 前置条件:Waterfill需要共享专家融合;LPLB需要存在EPLB冗余副本。
- 开销:Waterfill几乎零开销;LPLB每层需要一次all-reduce加上一次LP求解。
它们是互补的而不是竞争的:Waterfill消除了由密集共享专家贡献的不平衡,而LPLB消除了稀疏路由副本之间的不平衡。因为LPLB只在相同逻辑专家的有效副本之间重新分配流量,并且从不改变路由器的逻辑top-k,它与Waterfill一样保留了模型语义。
LPLB的最佳适用场景
LPLB的收益取决于实时批次与EPLB校准分布的偏差程度。当流量均衡且批次巨大(大规模、高度多样化的服务)时,LP几乎没有剩余的不平衡可以消除。当流量基本不变且狭窄(少数几乎相同的问题)时,静态EPLB已经捕获了分布,均匀分割接近最佳。LPLB在两者之间提供最强的信号——中等规模服务(专注于中等数量相关主题),其中每个批次的不平衡方式是离线放置未预料到的,但仍然具有足够的结构,使得最佳的每批次分割能够有意义地降低峰值rank负载。
实验环境
- 模型架构:DeepSeek-V3 FP8(作为DeepSeek-V3/R1风格的服务工作负载),以及DeepSeek V4 Flash FP8。
- 硬件配置:2个Hopper GPU节点,共计16个GPU。
- 并行与后端配置:TP16(张量并行), DP16(数据并行), EP16(专家并行), DP attention, DeepEP normal mode。
- 软件配置:SGLang直接集成运行(commit
a462e0f864103785fd3e64327104103f1356f220)。 - 数据集:MMLU, GPQA, 和 GSM8K 提示池(prompt pools)。
- 基准测试参数(DeepSeek V3/R1):
batch_size=1000,concurrency=256,request_rate=inf,max_tokens=1。 - 基准测试参数(DeepSeek V4 Flash):使用14,042-prompt MMLU池,
batch_size=512,concurrency=128,max_tokens=1,2轮预热,4轮测量。
实验结果
1. DeepSeek V3/R1上的Waterfill和LPLB评估
本组实验在MMLU, GPQA, 和 GSM8K数据集上对比了无EPLB、以及静态EPLB(包含red0, red16, red32不同冗余度)情况下的基线吞吐量、Waterfill吞吐量及其增益、LPLB吞吐量及其增益。
实验结果与分析:
- Waterfill表现:在所有数据集和配置下,Waterfill一致地提高了总吞吐量(增益范围从+1.48%到+4.66%)。这是因为它有效地将共享专家的工作转移到了负载较轻的EP rank上(参考Fig 3)。
- LPLB表现:LPLB在存在冗余专家副本(
red16和red32)时显著提高了吞吐量(GSM8K上最高达+7.34%),因为LP有实际的物理副本可供调度选择。相反,在没有冗余专家提供(red0)的情况下,LPLB没有平衡负载的空间,因此仅显示出算法的开销(吞吐量下降约0.95%至1.61%)(参考Fig 4)。
2. DeepSeek V4上的Waterfill评估
DeepSeek V4可以使用 HashTopK 路由路径,Waterfill必须在该输出路径中追加并重新映射共享专家槽位(通过PR #25391扩展)。本实验在MMLU风格的服务工作负载上验证V4 Flash FP8的性能。
实验结果与分析:
- 在无EPLB和静态EPLB(red0, red16, red32)设置下,Waterfill在DeepSeek V4 Flash上表现出一致的吞吐量提升,增益范围从+3.28%到+4.92%。
- 在静态EPLB red0配置下,吞吐量从基线的 49,253 tok/s 提升至 51,677 tok/s(+4.92%)。这些结果验证了共享专家平衡的理念不仅限于特定路由,在DeepSeek V4 Flash上同样有效(参考Fig 5)。
结论
本文介绍的Waterfill和LPLB两种调度时负载均衡方法,有效解决了MoE模型在推理服务阶段由于token路由不均匀导致的残余负载不平衡问题。实验证明,这两种方法在DeepSeek-V3/R1及V4模型上均能显著提升系统吞吐量。
精度与模型语义验证:
- Waterfill保留了模型语义,因为它不改变路由器的逻辑top-k决策,仅改变执行共享专家槽位的物理EP rank。
- LPLB同样保留了模型语义。它从不改变逻辑top-k;仅选择由选定逻辑专家的哪个物理副本执行每个token。由于逻辑专家的所有副本拥有相同的权重,token接收的结果独立于处理它的副本。这与EPLB和动态策略所依赖的精度保证完全一致。
补充细节:如何使用
启用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
关键标志说明:
--moe-a2a-backend deepep:使用DeepEP进行MoE all-to-all调度。--enable-deepep-waterfill:启用共享专家融合和Waterfill路径。--init-expert-location:可选项,从收集的专家分布统计数据中初始化专家放置和rank负载元数据。- DeepSeek V4支持使用在PR #25391中添加的
HashTopK路径。
启用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
关键标志说明:
--ep-dispatch-algorithm lp:选择LPLB线性规划调度器,以取代默认的static或均匀随机的dynamic策略。--ep-num-redundant-experts:为热点逻辑专家创建冗余物理副本。如果没有它们,LPLB就没有什么可平衡的(这就是为什么上述red0行没有显示LPLB增益的原因)。--init-expert-location:加载静态EPLB放置(物理到逻辑的映射,包括冗余槽位),该放置从专家分布记录运行中收集。此处的副本数必须与--ep-num-redundant-experts保持一致。
附录:致谢与相关PR
这项工作建立在SGLang DeepEP和MoE服务栈的基础之上,并通过SGLang项目中的社区协作开发完成。
相关代码合并请求(PRs):
- [#20089]:在EP下将共享专家融合到MoE调度中。
- [#19290]:为共享专家调度添加Waterfill负载均衡。
- [#25391]:支持DeepSeek V4 DeepEP Waterfill。
- [#24515]:LPLB:用于MoE专家并行的线性规划负载均衡器。
贡献者:
- NVIDIA团队:Xuting Zhou, Fei Liang, 和 Aichen Feng。
- SGLang团队:Cheng Wan。
特别感谢DeepSeek在 deepseek-ai/LPLB 开源了他们的LPLB工作,其用于在冗余专家副本之间平衡token的线性规划公式启发了本文描述的SGLang LPLB集成。
💬 评论讨论
欢迎在这里分享您的想法和见解!