跳转至

Tackling the Data-Parallel Load Balancing Bottleneck in LLM Serving: Practical Online Routing at Scale

会议: NeurIPS 2026
arXiv: 2605.06113
OpenReview: https://openreview.net/forum?id=6T70E9SD49
领域: LLM 效率
关键词: LLM推理服务、数据并行负载均衡、Prefill-Decode分离、在线调度、KV Cache

一句话总结

针对大模型推理服务中张量/专家并行跨副本集体通信引入的步级同步屏障,本文提出 BalanceRoute 在线路由算法族,通过分段线性 F 分数刻画安全裕度与溢出惩罚的强不对称性,并结合短视界前瞻分类器,在毫秒级开销下大幅降低数据并行负载失衡并提升端到端吞吐。

研究背景与动机

现代大语言模型参数规模庞大,生产部署中通常在单个节点或加速卡组内采用张量并行(TP)或专家并行(EP)进行切分,并在外层通过数据并行(DP)复制多个副本以横向扩展吞吐。在预填充与解码解耦(PD Disaggregation)架构下,请求在完成 Prefill 后被路由至某个 DP 解码副本,且由于搬迁海量 KV Cache 开销极大,请求一旦分配便与该副本强绑定。解码阶段每个自回归 step 都会执行跨卡甚至跨副本的集合通信(如 MoE 中的 EP all-to-all 或注意力中的 TP all-reduce),形成刚性步级同步屏障:每一步的耗时完全由当前负载最高的最慢副本决定,轻载副本必须空转等待。根据极值阶数统计效应,随着 DP 副本数 \(G\) 的扩展,最大负载与平均负载的差距急剧恶化,导致生产集群中超过 40% 的算力被屏障空等吞噬,造成巨大的算力与能耗浪费。

然而,现有开源服务系统(如 vLLM、SGLang)中的路由器仍普遍沿用通用分布式负载均衡启发式策略(Round-Robin、随机、Join-Shortest-Queue 或 Power-of-Two-Choices)。这些经典策略严重忽视了 LLM 解码特有的系统特征:请求任务与 KV Cache 强粘性无法随意抢占迁移、单请求计算与显存开销随生成 token 单调递增、Prefill 突发流量具有非平稳性,且解码调度器必须在每一步 sub-100ms 的硬性时延预算内完成全局决策。端点级负载感知路由器虽然尝试感知工作负载,但它们将每个实例视作独立端点,优化目标仅为单请求端到端响应时延,无法建模或缓解步级集体通信屏障下的系统级木桶短板效应。

本文切入角度在于:直接从步级失衡量的单步与多步边际增量出发,剖析请求分配给不同副本时的数学结构。核心 idea:利用吸收容量与突破上包络线之间的分段线性非对称边际效应构建 F 分数(未满裕度边际增益为 1,一旦溢出则导致全局 \(G-1\) 倍惩罚),设计两阶段贪心与精细子集分配算法,并融合极轻量的常数视界终止分类器,实现亚毫秒级、强鲁棒的在线最优数据并行调度。

方法详解

整体框架

BalanceRoute 部署于 Prefill 与 Decode 实例集群之间的有状态代理层(Stateful Proxy),整体数据流与控制流包含请求准入、状态镜像与两阶段分发调度: 1. Prefill 探测与缓冲:客户端请求到达代理后被克隆并发往 Prefill 节点(设置 stream=False, max_tokens=1),迫使 Prefill 计算完整 prompt 的 KV Cache 并获取实际输入长度 \(s_i\);完成的请求与其 KV 句柄暂存于 Prompt Pool 中。 2. 全局状态维护与预测前向投影:代理通过解析来自各 Decode 副本的 SSE 流实时更新每个请求的解码步数 \(a_i(k)\) 与各副本负载 \(L_g(k)\)。若启用短视界前瞻(BR-H),代理定期调用终止概率分类器,推演未来 \(H\) 步的负载包络线与裕度向量 \(\mathbf{m}_g\)。 3. 两阶段在线调度执行:每个调度 tick,调度器根据总空闲容量与阈值 \(S_{\text{greedy}}\) 分两步将 Prompt Pool 中的就绪请求分配给解码副本,完成 Mooncake 跨节点 KV Cache 拉取。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["客户端请求到达"] --> B["阶段1:Prefill 探测与 Prompt Pool 缓冲<br/>获取 Prompt 长度与 KV 句柄"]
    B --> C["阶段2:全局负载镜像与短视界前瞻投影<br/>解析 SSE 进度并计算安全裕度向量"]
    C --> D{"总空闲容量 Stot > Sgreedy ?"}
    D -->|是| E["阶段3:贪心粗充填 (Stage 1)<br/>最大空余卡分配最大就绪请求"]
    D -->|否| F["阶段4:精细子集背包分配 (Stage 2)<br/>优先队列弹栈 + 候选子集 F 分数最大化"]
    E --> G["阶段5:Mooncake 跨节点 KV 传输与解码执行"]
    F --> G

关键设计

1. 非对称边际失衡 F 分数:刻画木桶效应的剧烈跃变

定义单步解码总失衡度为 \(I(k) = G \cdot M(k) - \sum_{g=1}^G L_g(k)\),其中 \(M(k) = \max_g L_g(k)\) 为最慢副本负载。对于任意副本 \(g\),其安全裕度定义为当前它距离最慢副本的差距 \(m_g = M(k) - L_g(k) \ge 0\)。当向副本 \(g\) 准入一组增量负载为 \(\Delta_s(Q) = \sum_{i \in Q} s_i\) 的请求集合 \(Q\) 时,系统失衡度的变化量为 \(\Delta I_g(Q) = G \cdot (\Delta_s(Q) - m_g)_+ - \Delta_s(Q)\)。BalanceRoute 据此定义增益指标 F 分数: $\(F_g(Q) := -\Delta I_g(Q) = \Delta_s(Q) - G \cdot (\Delta_s(Q) - m_g)_+\)$ 该函数呈现极强的非对称分段线性特征:当 \(\Delta_s(Q) \le m_g\)(安全区)时,\(F_g(Q) = \Delta_s(Q)\),将负载填入未饱和副本能以 1:1 的斜率直接减小全局空等时间;而一旦 \(\Delta_s(Q) > m_g\)(溢出区),\(F_g(Q) = G m_g - (G-1)\Delta_s(Q)\),该副本一跃成为新的最慢 straggler,超出的每一单位负载都迫使其余 \(G-1\) 个副本同时多等待一单位时间,导致失衡以 \(G-1\) 倍的斜率恶化。这种单步向右折角的非对称代价表明:调度器必须优先把请求填满各副本的安全裕度;若不可避免溢出,也必须选择溢出边际成本最低的副本。

2. 产能自适应两阶段分配:平衡毫秒决策时延与尾部精准防溢

生产环境要求调度决策在毫秒级内完成,无法在海量等待队列与副本间进行复杂的组合搜索。算法设置了容量切换阈值 \(S_{\text{greedy}}\),将分配过程分为两阶段: - Stage 1(充裕容量贪心填入):当集群总可用空闲槽位 \(S_{\text{tot}} > S_{\text{greedy}}\) 时,大量副本均有充裕容量且处于安全区,算法直接挑选剩余容量最多的副本 \(g\),并将等待池中 \(F_g(\{i\})\) 最大(即 prefill 最长)的单请求贪心分配给它。 - Stage 2(紧缺容量精选子集):当 \(S_{\text{tot}} \le S_{\text{greedy}}\) 时,剩余槽位的分配极易触发跨越安全区边界的跃变。算法将有空闲槽位的副本构建为以 \((\text{cap}[g], m_g)\) 为键的优先队列,每次弹出容量和裕度最大的副本,在其前 \(R_{\max}\)(如 \(R_{\max}=4\))个候选请求构成的子集中,求解最优组合 \(Q^\star = \arg\max_{|Q| \le \min(\text{cap}[g], R_{\max})} F_g(Q)\)。若最佳子集得分仍非正,算法采用防饥饿兜底(Starvation Guard)分发最优单请求,确保系统不发生请求搁浅。

3. 短视界折扣 F 分数与终止分类接口(BR-H):前瞻平抑未来负载漂移

真实解码中单请求生命周期未知且 KV 负载动态上升。研究发现,长周期的精确解码长度预测因非平稳到达和重尾分布极其嘈杂且无必要,只需常数步短视界 \(H\)(如 \(H=80\))便能主导累积失衡。BR-H 设计了一个双阶段轻量预测契约: - 阶段一(终止分类器):校准模型输出请求在未来 \(H\) 步内是否终止的概率 \(\widehat{p}_{\text{fin}}^{(i)} = \Pr(r_i(k) \le H \mid x_i, a_i(k))\)。 - 阶段二(条件均值回归器):仅在确定于窗口内终止的正样本子集上,估计剩余步数均值 \(\widehat{\mu}_{\text{rem}}^{(i)}\)。 两阶段组合为未来 \(H\) 步内的预期负载贡献 \(\widehat{c}_i(k) = (1 - \widehat{p}_{\text{fin}}^{(i)}) \cdot H + \widehat{p}_{\text{fin}}^{(i)} \cdot \widehat{\mu}_{\text{rem}}^{(i)}\)。引入时间衰减向量 \(\mathbf{d} = (1, \gamma, \dots, \gamma^H)^\top\)(\(\gamma \in (0, 1]\))与惩罚校准系数 \(\beta\),视界 F 分数拓展为: $\(F_g(Q) = \alpha (\mathbf{1}^\top \mathbf{d}) \Delta_s(Q) - \beta \left( \Delta_s(Q)\mathbf{1} - \mathbf{m}_g \right)_+^\top \mathbf{d}\)$ 通过置信度门控(当 \(\widehat{p}_{\text{fin}} < 0.5\) 时强制回退到保守的 \(H\) 锚点),BR-H 避免了长尾幻觉,在预测完全不可知时(\(H=0, \alpha=1, \beta=G\))严格退化为免预测的 BR-0,兼具理论自洽性与极高容错性。

损失函数 / 训练策略

BR-H 的预测模块并非复杂的端到端大模型,而是解耦的高效轻量级组件: 1. 终止分类损失:阶段一二分类器在历史轨迹上以步长 \(\Delta T = H/2\) 采样数据元组,使用标准二元交叉熵损失(Binary Cross-Entropy)优化,重点追求预测概率的严格校准度(Calibration)。 2. 条件均值回归损失:阶段二回归器仅在真实终止的样本集(\(r_i(k) \le H\))上训练,使用均方误差(MSE)损失函数,确保输出逼近条件均值 \(\mathbb{E}[r_i \mid r_i \le H]\)。 3. 部署实现:实际系统支持无参数经验生存分析(Empirical-Survival,基于输出长度经验 CDF,单次调用开销 \(O(\log n) \approx 1\mu\text{s}\))与基于 Prompt 哈希前缀的记忆化匹配(ExactMatch),推理决策开销低于调度心跳两个数量级。

实验关键数据

主实验

在 144 张华为昇腾 Atlas 910C NPU 构成的真实生产集群上部署 DeepSeek-V3 671B(W8A8,4P1D 拓扑:4 Prefill 节点 TP=8,1 Decode 节点 DP=8, TP=4,共 \(G=8\) 个解码副本),分别在重尾长文本生产私有数据集(Proprietary Data)和微软 Azure-2024 公开推理数据集(筛选输出长度 \(>1000\) 饱和负载)上进行全面评测。

负载数据集 调度算法 平均负载失衡度 (Tokens) ↓ TPOT P95 (ms) ↓ 吞吐量 (Tokens/s) ↑ 吞吐相较最优基线提升
Proprietary Data Random 438,225 85.2 811 基线对比
Round-Robin 254,798 83.2 847 +4.4%
Power-of-Two (P2C) 262,940 84.0 836 +3.1%
Join-Shortest-Queue (JSQ) 215,110 83.2 843 基线(最强)
BR-0 (免预测版) 51,927 79.3 943 +11.8%
BR-H (部署版, Survival) 40,757 78.7 979 +16.1%
BR-H (部署版, ExactMatch) 36,277 78.8 977 +15.9%
BR-H (Oracle 完美视界) 23,577 78.0 1,029 +22.1%
Azure-2024 Random 146,564 90.3 706 基线对比
Round-Robin 110,049 89.1 771 +9.2%
Power-of-Two (P2C) 109,033 89.3 794 +12.5%
Join-Shortest-Queue (JSQ) 104,737 89.3 785 基线(最强)
BR-0 (免预测版) 54,051 83.3 854 +8.8%
BR-H (部署版, Survival) 44,010 82.1 893 +13.8%
BR-H (部署版, ExactMatch) 38,496 82.3 909 +15.8%
BR-H (Oracle 完美视界) 35,301 77.3 957 +21.9%

消融与扩展性实验

集群扩展性测试(在 Atlas 910C 上评估 \(G=4\) [2P1D, 48 NPUs]、\(G=8\) [4P1D, 96 NPUs] 和 \(G=16\) [5P1D, 144 NPUs])及组件消融:

配置 / 系统规模 评估算法 平均失衡度 (Tokens) ↓ TPOT P95 (ms) ↓ 吞吐量 (tok/s) ↑ 吞吐增益变化 (vs JSQ)
\(G=4\) (48 NPUs) JSQ (基线) 75,326 66.45 365.4 0.0%
BR-0 9,711 63.19 407.1 +11.4%
BR-H (Oracle, \(H=80\)) 6,576 62.52 415.3 +13.7%
\(G=8\) (96 NPUs) JSQ (基线) 215,110 83.20 842.8 0.0%
BR-0 51,927 79.32 942.5 +11.8%
BR-H (Oracle, \(H=80\)) 23,577 78.04 1,028.9 +22.1%
\(G=16\) (144 NPUs) JSQ (基线) 675,607 77.40 926.6 0.0%
BR-0 215,170 74.92 1,117.8 +20.6%
BR-H (Oracle, \(H=80\)) 117,067 72.04 1,247.4 +34.6%
调度开销分析 (P50) Stage 1 (贪心填充) < 1.2 μs - - 几乎零开销
Stage 2 (背包搜索) ~ 924 μs - - 占比 69%,仍 < 1ms
总调度 Tick 耗时 1.20 ms (P50) / 2.77 ms (P99) - - 远低于 60ms 解码心跳 (~50×)

关键发现

  • 屏障损失随副本规模超线性放大:随着 DP 副本数由 \(G=4\) 扩展到 \(G=16\),传统基线(如 JSQ)的失衡度从 7.5 万暴涨至 67.5 万 Tokens,屏障空等急剧恶化;而 BR-H 在 \(G=16\) 时吞吐提升幅度达到惊人的 +34.6%,证明理论预期的极值统计效应完全吻合。
  • 免预测 BR-0 即足以斩获显著收益:无需任何生成长度预测模块,BR-0 仅凭单步分段线性非对称 F 分数和两阶段分发,便将生产环境失衡降低了 4.1 倍,吞吐提升超 11%,为无法引入预测工程的系统提供了极其坚实的即插即用底座。
  • 置信度门控赋予极佳鲁棒性:在长尾且完全无 Prompt 重复的极端工况下,基于生存曲线的 Survival 预测器虽然 AUC 仅为 0.70,但因置信度门控自动使 \(\widehat{c}_i\) 回落至保守值 \(H\),其实际端到端吞吐依然超越 BR-0 并逼近 ExactMatch,避免了任何噪声劣化。

亮点与洞察

  • 将集群木桶效应精准形式化为分段线性非对称控制:洞察到超载一步对全局集群造成 \(G-1\) 倍放大损耗的核心数学结构,将原本复杂的非凸 NP-hard 调度规约为简洁高效的边际 F 分数优化。
  • 预测契约解耦与短视界截断哲学:打破了学术界追求长文本“精确剩余 token 预测”的迷思,证明只需捕获常数步短视界(\(H \approx 80\))的二分类终止信号,便能收敛绝大部分负载均衡收益。
  • 非侵入式工业级全栈落地:通过有状态 Proxy 对标准 OpenAI API 进行 payload 变异与 SSE 逆向解析,无需修改 vLLM 底层 C++/CUDA 引擎即可无缝实现端到端闭环调度。

局限与展望

  • 跨集群搬迁重叠调度尚未完全解耦:当前 BR-H 的 Pool 机制为了获得全局调度视界,KV 传输必须等待槽位完全释放后发起(Late-binding),未来可探索将 Cross-tier KV 搬迁与前序批次的残余解码流水线深度重叠。
  • 异构加速卡适配:论文基于同构 Atlas 910C 集群展开,在不同算力、不同内存带宽(BW-to-FLOPs 比值不同)的异构副本场景下,单步线性假设 \(ax+b\) 的系数需引入动态自适应标定。

相关工作与启发

  • vs vLLM / SGLang 官方路由器 (Round-Robin, JSQ, P2C):传统开源路由仅考虑请求计数或连接数,对动态膨胀的 KV 显存开销及同步屏障完全盲目;本文构建了首个针对步级集合通信同步特性的专用在线调度理论与实现。
  • vs 工作负载感知端点路由 (Intelligent Router, Jain et al., 2024/2025):后者将各 LLM 实例视为独立的 Web 服务节点,优化目标为端到端延迟且各节点无同步关系;本文面向切分模型在自回归每一步所受到的刚性跨卡步级屏障,解决的是全局极值木桶短板问题。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次揭示 LLM 数据并行解码同步屏障的非对称边际特性,提出视界 F 分数与两阶段调度理论]
  • 实验充分度: ⭐⭐⭐⭐⭐ [在 144 卡昇腾集群上运行 DeepSeek-V3 671B 生产真实负载与公开 Azure 轨迹,涵盖 G=4/8/16 跨规模全方位评测]
  • 写作质量: ⭐⭐⭐⭐⭐ [数学建模严密优雅,系统设计契约清晰,工程落地与理论推导环环相扣]
  • 价值: ⭐⭐⭐⭐⭐ [直击超大规模 MoE/稠密模型集群中超 40% 算力被屏障吞噬的工业级第一痛点,具有极强的落地与推广价值]