Improving DeepEP MoE Load Balance in SGLang with Waterfill and LPLB #

1. TL;DR #

SGLang 在 DeepEP/EPLB 之上加了两个 dispatch-time MoE 负载均衡特性：Waterfill 把 dense shared expert 当作可分派的 slot，按各 rank 当前负载"填谷"到轻负载 rank（V3/R1 +1.48%~+4.66%，V4 Flash 最佳 +4.92%）；LPLB 用每层 min–max 线性规划在 redundant 副本间重分流量（+0.84%~+7.34%，无 redundant 时为负）。两者都不改 logical top-k，保持语义。

2. 痛点 / 方法 / 结果 #

Q1 — 痛点：大 MoE（DeepSeek-V3/R1、V4）靠 Expert Parallelism 把 experts 分散到多 GPU 上服务，但 router 产生的 expert traffic 天然不均衡：一个 batch 里某些 expert 拿到远多于其它 expert 的 token，EP group 就被最忙的 rank 拖住（compute 和 communication 都受影响）。静态放置（EPLB）只优化长期分布，无法消除单 batch 的残余不均衡。缺的是一个运行时环节，在不改模型语义的前提下决定"哪个物理 replica / 哪个物理 rank"处理请求。

Q2 — 方法：两个作用于不同 dispatch 选择的运行时特性。

• Waterfill（针对 dense shared expert）：shared expert 被每个 token 都用到，若永远本地计算，重负载 rank 无法卸载。Waterfill 把 shared expert 当作一个额外的可分派 expert slot——统计各 rank 的 routed load 作为负载分数，算出目标 waterline，把 shared-expert slot 按 slack 比例采样分配到较轻的 rank（带轻微 local 偏好）。routed 选择完全不变，只改 shared expert 在哪个物理 rank 执行。
• LPLB（针对 EPLB 复制出的 redundant routed expert）：EPLB 默认把热 expert 的 token 在其物理副本间均分，仅当离线标定分布 == live traffic 时才最优。LPLB 每层每 batch 看真实 per-expert token 计数，解一个 min–max LP，决定每个 replicated expert 的流量在副本间怎么拆，使最大 per-rank load 最小。它是 dynamic（均匀随机选副本）策略的 drop-in 替换。

核心技术壁垒：LPLB 把"per-batch 最优副本分流"做成了可以放进推理关键路径的东西。关键不在于"用 LP 求 min–max"这个想法本身，而在于让它零广播、低开销地在 DP-attention（同 step 各 rank 跑 prefill/decode/idle 不同模式）下运行——所有 rank 参与一次 all-reduce 拿到相同全局分布后各自独立解同一个 LP 得到同解（省掉结果广播），LP 由基于 cuSOLVERDx/cuBLASDx 的融合 IPM kernel 在 GPU 上求解、startup 预编译免 JIT，整条 per-batch 路径压成三次 CUDA kernel launch。

Q3 — 结果：两 Hopper 节点 16 GPU。Waterfill 在所有配置一致提升（V3/R1 +1.48%~+4.66%；V4 Flash +3.28%~+4.92%）；LPLB 仅在存在 redundant replica（red16/red32）时提升（最高 GSM8K red32 +7.34%），red0 时因无副本可平衡而净为负（−0.95%~−1.61%，纯 all-reduce/solve 开销）。两者都验证不损精度。

3. 架构 / 方法图 #

3.1 系统位置：dispatch-time LB 在 DeepEP MoE 层中的位置 #

两条路径都只改"物理落点"，不改 router 的 logical top-k 选择——这是保持语义的结构性原因。

3.2 Waterfill 的 waterline 机制（重建自正文） #

按 token 的 shared-expert slot 分配流程：

1. 统计已落在每个 EP rank 上的 routed expert 负载 $L_r$（dynamic 模式下先做一次 EP-group collective，用全局 routed-load 向量 + 各 rank 当前 local batch size 作为分数）。
2. 每个参与 token 加一个 shared-expert slot，令 $N$ 为待放置的 shared slot 总数、$R$ 为 EP group 大小，计算目标 waterline：

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

1. 低于 waterline 的 rank 有 slack（容纳空间）：

$$S_r = \max(H - L_r, 0)$$

1. 每个 token 从候选 rank 中按 slack 成比例采样目标 rank（带轻微 local-rank 偏好）；若所有候选 slack 为 0，退回明显更轻的候选 rank，仍保留 local 偏好。

"填谷"直觉：把等量的水（shared 工作）倒进高低不平的容器（各 rank 负载），先填满最低处。候选集默认通信保守——只用 token 已为 routed experts 访问过的 rank（source rank 兜底），避免新增 all-to-all 目的地；all-rank 模式给更大自由度但可能增加每 token 一个新 dispatch 目的地。

3.3 LPLB 的线性规划（重建自正文） #

决策变量：各 replicated expert 的 per-copy 负载 $x_i$、各 rank 的 slack、标量峰值 $M$。single-copy expert 不是变量（只贡献固定项），故 LP 规模只随 redundant expert 数 和 rank 数 增长，不随全 expert 数。

目标——最小化峰值：

$$\min\ M$$

每 rank 一条负载约束（redundant 副本 load + single-copy 固定 load + 到峰值的非负 slack = $M$），逼 $M \ge$ 每 rank 真实 load：

$$\sum_{i \in \text{redundant}(r)} x_i + C_r + s_r = M,\quad s_r \ge 0$$

每个 replicated logical expert 一条守恒约束（副本 load 之和 = 观测总 load，保证只重分配、不增删 token）：

$$x_1 + x_2 + \dots + x_n = L$$

约束矩阵拆两半：离线结构块（copy→logical 映射、per-rank 副本归属、slack/$-M$ 列）只依赖 expert-to-GPU 放置，startup 与每次 EPLB rebalance 后预算一次；在线部分每 batch 只换 RHS（观测到的 redundant load + per-rank single-copy load）。Big-M 辅助列保证求解可行、被重罚趋零。

4. 作者证明 #

无形式化作者证明 — 仅实证（blog 未给收敛/最优性定理，仅给出 LP 的构造正确性论证 + 吞吐实测）。下面重建 notation 与方程物理意义，并给出可核验的一致性检查。

符号表

符号	含义
$L_r$	rank $r$ 的 routed 负载分数（Waterfill）
$N$	待放置的 shared-expert slot 总数
$R$	EP group 大小（本文 = 16）
$H$	Waterfill 目标 waterline
$S_r$	rank $r$ 的 slack（低于 waterline 的空间）
$M$	所有 rank 的最大负载（LPLB 目标标量）
$x_i$	replicated expert 第 $i$ 个副本承担的负载（LPLB 决策变量）
$C_r$	rank $r$ 的 single-copy 固定负载（LPLB 输入常量）
$s_r$	rank $r$ 到峰值的 slack（LPLB，非负）
$L$	某 replicated logical expert 的观测总负载

方程物理意义

• $H=\lceil(\sum_r L_r + N)/R\rceil$：把总负载（现有 routed + 新增 shared slot）均摊到 $R$ 个 rank 的理想水位；向上取整保证整数 slot 可容纳。
• $S_r=\max(H-L_r,0)$：只有低于水位的 rank 才有接收 shared 工作的空间；已超水位的 rank slack=0，不再被分派。
• $\min M$ s.t. rank 约束：把最忙 rank 拉向均值，直接缩短 EP imbalance 造成的 grouped-GEMM 尾。
• 守恒 $x_1+\dots+x_n=L$：LP 只在同一 logical expert 的合法副本间搬运既有流量。

一致性检查（6 项）

1. 量纲一致：$H$、$L_r$、$M$、$x_i$、$C_r$ 全为"token/负载计数"，方程两边同量纲。✓
2. 守恒：$\sum_i x_i = L$ 使全 batch token 总数在重分配前后不变，不凭空增删。✓
3. 可行性下界：rank 约束 + $s_r\ge0$ 强制 $M\ge C_r+\sum x_i$，即 $M$ 不低于任一 rank 真实负载——min–max 的正确下界。✓
4. 退化行为：red0（无 redundant 副本）时 LP 无决策变量可调，解 = 均分，理论无收益；实测转为净开销（负增益），与公式预期一致。✓
5. 边界（Waterfill 全零 slack）：所有候选 rank 均达/超水位时，probability-proportional-to-slack 分母为 0，退回"明显更轻的候选 rank"——保证算法总能产出合法目标。✓
6. 语义不变量：两算法都不改 router logical top-k，且 LPLB 各副本权重相同，token 结果与副本选择无关——与 EPLB/dynamic 依赖的精度保证同源。✓

5. 实验与数据 #

通用设置：两 Hopper GPU 节点、16 GPU；DeepSeek-V3 FP8 作 V3/R1-style workload；TP16 / DP16 / EP16 / DP attention / DeepEP normal mode。run dsv3_ep16_three_dataset_lplb_matrix_20260605_101821，SGLang commit a462e0f864103785fd3e64327104103f1356f220；benchmark shape batch_size=1000, concurrency=256, request_rate=inf, max_tokens=1。每次比较都在同一 placement 配置内；LPLB 只有在 EPLB placement 启用时才有意义。

5.1 DeepSeek V3/R1：Waterfill 与 LPLB 吞吐矩阵 #

Dataset	Baseline setting	Baseline	Waterfill	Waterfill gain	LPLB	LPLB gain
MMLU	No EPLB	28,968	29,697	+2.52%	-	-
MMLU	Static EPLB, red0	30,392	31,424	+3.40%	29,938	-1.50%
MMLU	Static EPLB, red16	30,638	31,483	+2.76%	31,104	+1.52%
MMLU	Static EPLB, red32	30,714	31,169	+1.48%	31,547	+2.72%
GPQA	No EPLB	23,201	24,283	+4.66%	-	-
GPQA	Static EPLB, red0	26,322	26,970	+2.46%	25,899	-1.61%
GPQA	Static EPLB, red16	26,124	26,683	+2.14%	26,350	+0.86%
GPQA	Static EPLB, red32	25,975	26,655	+2.62%	26,193	+0.84%
GSM8K	No EPLB	29,649	30,892	+4.19%	-	-
GSM8K	Static EPLB, red0	33,058	34,529	+4.45%	32,744	-0.95%
GSM8K	Static EPLB, red16	34,026	35,226	+3.53%	35,474	+4.26%
GSM8K	Static EPLB, red32	33,988	35,070	+3.19%	36,482	+7.34%

（单位均为 tok/s）读法要点：

• Waterfill 在每一行都为正（+1.48%~+4.66%），因为它只改 shared expert 的物理落点，与是否有 redundant 副本无关。
• LPLB 的收益随 redundant 副本数上升：red0 全为负（无副本可分流，只剩 all-reduce+solve 开销），red16 转正，red32 最强（GSM8K +7.34%）。红字负值行是关键诚实信号：它精确对应 §4 检查 4 的退化预期。

5.2 DeepSeek V4 Flash：Waterfill 验证 #

V4 用 HashTopK 路由路径，Waterfill 需在 HashTopK 输出路径也 append/remap shared-expert slot（#25391 扩展）。V4 Flash FP8、两 Hopper 节点、14,042-prompt MMLU pool、batch=512, concurrency=128, max_tokens=1、2 warmup + 4 measured round、trimmed mean。

Configuration	Baseline	Waterfill	Gain
No EPLB	45,951	47,876	+4.19%
Static EPLB, red0	49,253	51,677	+4.92%
Static EPLB, red16	50,006	51,655	+3.30%
Static EPLB, red32	50,167	51,813	+3.28%

（单位 tok/s）因 batch/concurrency 更小，这组应读作 V4 专属验证，而非与 §5.1 直接横向对比。

5.3 收益的适用区间 #

LPLB 收益 ∝ live batch 偏离 EPLB 标定分布的程度，呈"中间最大"：

流量特征	LPLB 收益	原因
均衡且 batch 巨大（大规模高多样性）	低	残余不均衡本就少
近乎不变且窄（少量几乎相同问题）	低	静态 EPLB 已捕获，均分近最优
中等规模、适度相关主题	最强	每 batch 以离线未预料的方式不均衡，但仍结构化，最优 per-batch 拆分显著降峰

6. 论证链 #

#	论点	依据（篇内）
1	EP 让大 MoE 可服务，但 router 产生不均衡 traffic，EP group 被最忙 rank 拖住	Introduction / Background：routed 稀疏、shared 稠密、redundant 多副本三类 load pattern
2	静态放置（EPLB）无法消除单 batch 残余不均衡 → 需 dispatch-time LB	Background：静态 placement 优化长期分布，实际 batch 仍集中
3	shared expert 若永远本地算，重负载 rank 无法卸载 → Waterfill 按 waterline 填谷到轻 rank	Waterfill 节：$H$/$S_r$ 机制 + 通信保守候选集
4	Waterfill 要低开销，需让 shared dispatch 对 DeepEP 可见 → shared expert fusion 前置机制	#20089 fusion（固定分配）→ #19290 换 load-aware
5	EPLB 对 redundant 副本的均分仅在分布匹配时最优，实际常漂移 → LPLB 每 batch 解 min–max LP	LPLB 节：$\min M$ + 守恒约束
6	LP 要进关键路径，需零广播 + 低开销 → all-rank all-reduce 后各自解同一 LP + on-GPU IPM kernel + 3 次 launch	From Global Counts / From LP Solution 节
7	两法都不改 logical top-k 且副本权重相同 → 保持语义	Accuracy Validation 节
8	实测：Waterfill 全正、LPLB 随 redundant 数增益、red0 净负 → 与机制预期自洽	Evaluation 两表

7. 实现 cross-reference #

SGLang 上游 PR（blog 致谢与正文明确列出）：

• #20089 — Fuse shared expert into MoE dispatch under EP。把 shared expert 表示为同一 DeepEP MoE layout 里的一个额外 expert slot（TopK 输出加一列，每 rank 在 routed experts 旁保留一个 shared slot），共用 dispatch / grouped-GEMM / combine 流程；此 PR 用固定 local 分配。
• #19290 — Add Waterfill load balancing for shared expert dispatch。把 #20089 的固定分配替换为 load-aware 的 waterline 分派（static/dynamic 两种行为）。
• #25391 — Support DeepSeek V4 DeepEP Waterfill。在 V4 的 HashTopK 输出路径 append/remap shared-expert slot。
• #24515 — LPLB: linear-programming load balancer for MoE expert parallelism。引入 --ep-dispatch-algorithm lp，把 LP 解归一化为 log2phy_prob 概率分布，per-token 从中采样物理副本。

启用方式（正文给出的代表性命令行）：

• Waterfill：--moe-a2a-backend deepep --deepep-mode normal --enable-dp-attention --enable-deepep-waterfill（--enable-deepep-waterfill 同时打开 shared expert fusion + Waterfill 路径）；--init-expert-location 可选加载 expert 分布统计。
• LPLB：--ep-dispatch-algorithm lp + --ep-num-redundant-experts 16（无 redundant 则无可平衡，对应 red0 无增益）+ --init-expert-location（加载含 redundant slot 的 physical-to-logical map，replica 数须与 --ep-num-redundant-experts 一致）。

核心技术壁垒（展开）：把 per-batch 最优 min–max 分流做成关键路径可承受的组件。三个使其可行的设计缺一不可——(a) DP-attention 下同 step 各 rank 跑 prefill/decode/idle，无单 rank 见全局分布，靠"所有 rank（含 idle 贡献 0）一次 all-reduce"补齐全局计数；(b) 各 rank 用相同输入独立解同一个 LP 得同解，从而省掉结果广播；(c) LP 由 cuSOLVERDx/cuBLASDx 融合 IPM kernel 在 GPU 上求解、startup 按层矩阵形状预编译免首请求 JIT，整条 build-RHS→solve→extract-split 路径压成三次写预分配 buffer 的 CUDA kernel launch，把 launch 开销与 host sync 移出关键路径。

关键实现细节（易漏）：

1. Waterfill 默认候选集是"通信保守"的——只允许 token 已为 routed experts 访问过的 rank（source rank 兜底），因为在 GPU MoE serving 中通信通常比额外 shared 计算更贵；all-rank 模式换来更大平衡自由度但可能新增每 token 一个 dispatch 目的地。
2. LPLB 约束矩阵离线/在线分离：结构块（copy→logical 映射、per-rank 副本归属、slack/$-M$ 列）只依赖放置，startup 与每次 EPLB rebalance 后预算一次；每 batch 只更新 RHS。single-copy expert 不进决策变量，使 LP 规模只随 redundant 数与 rank 数增长。
3. LPLB 是 dynamic（均匀随机选副本）的 drop-in 替换：保持同样的概率化 per-token dispatch 形态，仅把均匀抽样换成该 batch 的 load-optimal 分布 log2phy_prob。

blog 未开源本文所述 SGLang 集成的独立代码文件行号，实现以上述 PR 编号为准；DeepSeek 侧灵感来源为 deepseek-ai/LPLB。