释放算力潜能：TensorRT LLM ADP 平衡策略让推理吞吐量再提升 33%

动机与背景

在 DeepSeek MLA + MoE 架构下，在最大吞吐量场景中，通常采用注意力数据并行 (Attention Data Parallel, ADP) + MoE 专家并行 (Expert Parallel, EP) 策略，以消除冗余的 KV cache 缓存存储，并使用分离服务来防止 ADP 不平衡。

然而，某些部署场景中动态批处理 (In-Flight Batching, IFB) 推理更具优势，包括：

• 系统复杂性降低：避免分离式架构所带来的运营开销与维护成本。
• 特定工作负载模式：输入序列长度 (input sequence lengths, ISL) 较短而输出序列长度 (output sequence lengths, OSL) 较长的场景。
• 离线推理：在批处理环境中对首 token 时延 (Time-To-First-Token, TTFT) 和每个输出 token 的延迟 (Time-To-Output-Token, TPOT) 的要求更为宽松。

然而，IFB 在注意力模块中带来了负载不平衡的挑战，严重影响了系统性能。核心问题在于不同的 rank 在同一迭代中同时处理异构工作负载。例如，一些 rank 可能在处理计算密集型的 context（上下文）阶段，而其他 rank 则在执行 generation（生成）阶段，从而形成 token 处理负载的巨大差异。这会成为系统整体吞吐量的瓶颈，因为迭代时间取决于速度最慢的 rank。

为了解决这一关键性能限制，我们提出了 ADP 平衡策略——一种新颖的调度优化方案，旨在实现 DP rank 之间的最佳负载分配，并最大化系统利用率。

理论分析与建模

优化目标：最小化不同 GPU rank 之间的负载不平衡，以最大化整体系统吞吐量。

数学建模

我们对 ADP 中负载不平衡带来的性能影响进行了建模与量化。由于各个 rank 间的工作负载可能是异构的，因此在给定迭代中的注意力模块的执行时间由负载最大的 rank 所限制。

其中 time_{i, m} 代表第 i 次迭代中 rank m 的执行时间，N 是数据并行的规模。

为了量化负载平衡和理论性能界限，我们定义了两个关键指标：

1. 平衡比率

平衡比率用来衡量每次迭代中注意力模块内各 rank 之间的负载分布情况：

其中：

• tokens_avg 表示所有 rank 中平均的 token 数量
• tokens_max表示所有 rank 中最大的 token 数量
• tokens_i 表示 rank i 处理的 token 数量

2. 光速吞吐量 (Speed-of-Light Throughput, SOL TPS)

光速吞吐量代表在理想负载均衡情况下可实现的理论最大吞吐量：

其中：

• time_i: 第 i 次迭代的测量执行时间
• time_elapsed: 端到端执行时间的实测总值
• tps_actual: 每秒观察到的吞吐量（token / 秒）
• tps_sol: 在理想负载均衡下的理论最大吞吐量

这一理论框架使我们能够量化当前与最佳系统利用率之间的性能差距，从而为优化提供明确的目标。

负载均衡的调度策略

ADP 中的基本挑战在于，同一迭代中不同 rank 在处理的 token 负载可能差异很大，这使得整体执行时间会受限于负载最重的 rank。

基线：轮询式 token 分配

传统方法采用全局负载均衡策略，根据 num_tokens 对输入请求进行排序，并通过轮询调度将它们分配到各个 rank，如图 1 所示。该方法从累积的角度实现了合理的 token 分配，且当所有 rank 同时处理 context 请求时有效减少了 token 数量的差异。

图 1. 基线轮询式策略通过排序和循环分配，平衡各 rank 之间的 context 请求 token

局限性：尽管该方法在全局上有效，但并不能保证每次迭代的负载均衡。当一些 rank 处理 context 阶段，而其他 rank 处理 generation（解码）时，会产生严重的负载不平衡，从而影响整体执行时间。

ADP 平衡策略：协调等待机制

为了解决每次迭代的负载不平衡问题，我们提出了 ADP 平衡策略，它采用复杂的等待机制来同步各个 rank 之间的 context 处理。其核心原则是战略性延迟：系统不会立即将 context 请求调度到可用的 rank，而是进行战略性等待，确保多个 rank 拥有相似的工作负载后再继续。

算法设计：该策略引入两个互补的控制参数：

1. Context 同步 (timeout_iters)

• 目的：确保各个 rank 之间 context 处理的时间一致性
• 机制：当某个 rank 已准备好处理 context 任务，而其他 rank 仍处于 generation 阶段时，该 rank 会等待最多 timeout_iters 次迭代，直至所有的其他 rank 都接收到 context 任务
• 优势：避免出现单个 rank 处理 context 任务，其余 rank 处理 generation 任务的情况

2. 批次均衡（batching_wait_iters）

• 目的：平衡各 rank 中累积的 context 批次数量
• 机制：在初始同步后，累积 context 批次较少的 rank 会等待最多 batching_wait_iters 次额外迭代，以累积更多批次
• 优势：避免因 context 批次累积不均导致的负载失衡，防止部分 rank 拥有多个批次而其他 rank 仅拥有一个批次的情况。

性能分析：基线 vs. ADP 平衡

为了说明我们方法的有效性，考虑一个简单的场景，其中：

• 所有 rank 具有相同长度的 context，并且有 M 个正在进行的请求
• N 个新请求在 N 次迭代中按顺序到达
• Context 处理时间：time(ctx) >> generation 处理时间：time(gen)

基线行为：在传统方法中，context 会在各 rank 间顺序处理，导致严重的负载失衡：

iter_i:     [*C0*, g01, ..., g0M], [g10, g11, ..., g1M], ..., [gN0, gN1, ..., gNM]

iter_i+1:   [g00, g01, ..., g0M], [*C1*, g11, ..., g1M], ..., [gN0, gN1, ..., gNM]

...

iter_i+N-1: [g00, g01, ..., g0M], [g10, g11, ..., g1M], ..., [*CN*, gN1, ..., gNM]

图例：*Ci* = context 请求 i，gij = 在 rank i 上的 generation 请求 j

• 每次迭代时间：time(ctx)（由 context 处理主导）
• 总执行时间：time(ctx) × N
• 平衡比率：(ctx_len + (M-1) + M × (N-1)) / (N × ctx_len)（平衡性差）

ADP 平衡策略：我们的方法是通过战略性等待来同步 context 处理：

iter_i:     [g00, g01, ..., g0M], [g10, g11, ..., g1M], ..., [gN0, gN1, ..., gNM]

iter_i+1:   [g00, g01, ..., g0M], [g10, g11, ..., g1M], ..., [gN0, gN1, ..., gNM]

...

iter_i+N-1: [*C0*, g01, ..., g0M], [*C1*, g11, ..., g1M], ..., [*CN*, gN1, ..., gNM]

• 每次迭代时间：前 N-1 次迭代为 time(gen)，最后一次迭代为 time(ctx)
• 总执行时间：time(gen) × (N-1) + time(ctx)
• 平衡比率：1.0（理想平衡）
• 时间节省：(time(ctx) – time(gen)) × (N-1)

权衡因素：

✅ 由于最佳负载均衡而提高吞吐量

✅ 最大化所有 rank 的 GPU 利用率

⚠️ 由于战略性等待机制而增加 TTFT

📋 最适合吞吐量导向的场景，在这些场景中 TTFT 不是关键因素

实验

设置

数据集配置

我们使用一个综合数据集评估我们的方法，该数据集包含 16,000 个推理请求，具有以下特征：

• 请求量：总共 16,000 个请求
• 平均输入长度：803 个 token
• 平均输出长度：3,653 个 token
• Token 分布：图 2 展示了输入和输出序列的分布模式

图 2. 输入和输出 token 长度的分布

数据集特征：该数据集在序列长度上表现出显著的多样性，输出 token 呈现出明显的长尾分布。这种异质性给负载均衡带来了严峻的挑战，因为很难在同一迭代中协同调度多个 context 请求并最小化计算空泡。这一特性使其成为评估我们调度策略的理想测试场景。

性能结果

我们评估了三种不同的配置，以展示我们 ADP 平衡策略的渐进式优势：

1. 基线：轮询调度
2. ADP 平衡（context 等待）：仅实现 timeout_iters 参数
3. ADP 平衡（完整策略）：完整实现，包括 timeout_iters 和 batching_wait_iters 两个参数

性能总结

配置	实际 TPS	平均平衡比例	SOL TPS	加速比
基线	25,664	54.11%	39,552	1.00×
ADP 平衡 （context 等待）	33,499	84.33%	38,312	1.31×
ADP 平衡（完整策略）	34,140	87.70%	37,912	1.33×

核心观察：

• 仅 context 等待就显著实现了 31% 吞吐量提升
• 完整策略实现了 33% 的总加速比，并且接近最优平衡（87.70%）
• 平衡比率的改善：从 54% 提升至 87%，大幅降低了负载不平衡问题

注意：使用等待机制时，SOL TPS 的下降是因为 context 调度中的策略性延迟增加了完成所有请求所需的总迭代次数。由于 SOL TPS 的计算仅考虑每次迭代中的负载不平衡影响，并未反映因 context 延迟进入而导致的迭代次数增加，因此即便整体系统效率有所提升，SOL TPS 仍会呈现出表面上的下降。

基线性能

图 3 全面展示了基线系统的运行情况，该图同时呈现了迭代中各 rank 的平均 token 数量（上图）以及对应的平衡比（下图）。平衡比是一项关键指标：其数值越接近 1.0，代表负载越均衡；数值越接近 0.0 则意味着负载严重失衡。

图 3. 基线性能概述，显示所有迭代的 token 分布和平衡比率

关键洞察：

• 性能落差：SOL TPS 数值为 39,552，与实际的 25,664 TPS 之间存在 54% 的效率损耗
• 系统表现：迭代超过 12,000 次后，所有请求均过渡到 generation 阶段，这自然地减少了系统负载不平衡现象

图 4 聚焦关键失衡区间 [100-12,000]，清晰揭示了负载分布巨大的不稳定性：

图 4. 针对 100-12,000 次迭代的详细基线分析，显示存在严重的负载波动

性能瓶颈：

• 平衡比频繁降至 0.4 或以下，表明负载失衡程度超过 60%
• 在关键窗口期内，理论上有 70.23% 的改进潜力
• 负载分布的剧烈波动导致系统性能呈现高度不可预测性

采用 context 等待机制的 ADP 平衡

Context 等待机制 (timeout_iters=50)，作为我们的首个优化组件展现出显著成效：该方案通过实现 context 同步使性能获得显著提升。

性能提升成果：

• 吞吐能力：达到 33,499 TPS（实现 1.31 倍加速比）
• 均衡改善：平均达 84.33%（基线水平为 54.11%）
• 运行效率：实际 TPS 大幅接近理论 SOL TPS (38,312)

图 5. 等待机制在 100-12,000 次迭代期间的性能表现，展现改善后的均衡稳定性

尚未解决的挑战：尽管取得显著改进，但由于以下因素，残余失衡现象依然存在：

• 超时情况：当 context 请求非均匀到达时，部分 rank 会超出等待阈值
• 批次累积差异：等待时间较长的 rank 会累积多个 context 批次，而刚释放的 rank 仅处理单个批次
• 平衡不彻底：尽管取得了一定平衡效果，但不同 rank 上 context 数量不同仍造成负载波动

基于此分析结果，我们开发出第二个优化组件：批次均衡机制。

采用完整策略的 ADP 平衡

完整的 ADP 平衡策略融合了 context 同步与批次均衡两大机制，从而实现极佳的负载均衡效果。

配置参数：timeout_iters=50 + batching_wait_iters=10

性能优化成果：

• 峰值吞吐量：达 34,140 TPS（实现 1.33 倍加速比）
• 最优均衡度：平均负载均衡率达到 87.70%
• 接近理论效率：实际 TPS (34,140) 更接近 SOL TPS (37,912)
• 系统稳定性：各迭代周期内的负载波动显著降低

图 6 清晰印证了完整 ADP 平衡方案的有效性。该图通过可视化数据展示了通过 context 同步与批次均衡机制的协同运作，如何在关键执行窗口期实现近乎最优的负载均衡效果。

图 6. 完整 ADP 平衡策略在 100-12,000 次迭代区间展现卓越的均衡稳定性

相较于等待机制的改进亮点：

• 稳定性增强：负载均衡率持续保持更高水平，同时波动性显著降低
• 残余失衡缓解：批次均衡机制有效处理了剩余的负载差异
• 系统可预测性：各迭代周期展现出更一致的性能特征

实施方案的权衡因素：

✅ 最大吞吐量提升：较基线实现 33% 的增益

✅ 近乎最佳的负载均衡：达到 87.70% 的平均均衡比率

⚠️ 迭代开销：等待机制增加总迭代次数

⚠️ TTFT 影响：策略性延迟会影响首 token 时延指标

生产环境配置：用户可通过添加以下配置，启用完整的 ADP 平衡策略：

 attention_dp_config:

    enable_balance: true

    batching_wait_iters: 10

    timeout_iters: 50  

帕累托分析：吞吐量与延迟的权衡优化

了解 ADP 平衡策略固有的性能权衡对生产环境部署决策至关重要。图 7 通过全面的帕累托前沿 (Pareto frontier) 分析，呈现了在不同工作负载强度与参数配置下，系统吞吐量（每 GPU 的 TPS）与 TTFT 的关联关系。

实验设计：通过在不同系统负载条件下对 timeout_iters (TO) 与 batching_wait_iters (BW) 进行多组配置验证，揭示了参数调节如何影响吞吐量与延迟的核心权衡关系。

图 7. 帕累托前沿分析呈现各 ADP 平衡配置方案中的吞吐量与延迟权衡关系

核心发现：

1. 全域吞吐量提升：在从延迟敏感型到吞吐最大化部署的整个操作频谱范围内，ADP 平衡方案均能持续提供更优异的 TPS / GPU 性能
2. 可扩展性优势：在系统负载较高、负载失衡代价最严重的场景下，性能改进效果尤为显著
3. TTFT 权衡因素：受策略性等待机制影响，吞吐量提升会导致首 token 时延的增加，参数值越高，带来的吞吐量提升越明显，但代价是响应启动耗时延长

配置指南：

• 低负载场景：batching_wait_iters 收效甚微，反而会增加延迟开销
• 高吞吐场景：两项参数共同作用可显著提升性能表现
• 均衡部署场景：适中的参数值能实现吞吐量与延迟的最优平衡

生产影响：以上分析能帮助系统操作人员，无论是追求最低响应延迟还是最大系统吞吐量，都可以根据具体部署需求做出基于数据的配置决策。

总结

ADP 中的负载不均衡已成为大语言模型推理系统的核心瓶颈，尤其在 IFB 场景下，异构工作负载会引发严重的性能损耗。本研究提出的 ADP 平衡策略通过协同等待机制实现的精细化调度优化方案，成功攻克了这一关键难题。

技术贡献：本方案采用了两项互补的优化组件：context 同步机制 (timeout_iters) 与批次均衡机制 (batching_wait_iters)。这些机制协同运作，实现了数据并行 rank 间计算密集型 context 处理的时间对齐，从而有效消除了因 rank 级负载不均导致的性能瓶颈。

实验验证：在 DeepSeek V3 架构上的全面评估证明了显著的性能提升：

• 吞吐量提升 33%：从 25,664 TPS 提高至 34,140 TPS
• 实现 87% 负载均衡率：较 54% 的基线实现巨大飞跃
• 接近理论效率：实际性能已接近 SOL TPS 极限

生产准备情况：帕累托前沿分析为实际部署提供了关键洞见。该策略虽会引入 TTFT 权衡，但在多样化运行场景中始终能提供更卓越的吞吐量。此外，可配置参数框架支持运维人员根据特定性能需求进行精准优化，无论是优先保障响应延迟还是追求系统吞吐量。

致谢

从系统性能分析到优化，ADP 平衡策略的成果离不开卓越的团队协作。虽无法逐一致谢每位贡献者，但我们由衷感谢工程师团队的集体奉献——正是他们的专业知识将 TensorRT LLM 的性能推向了全新高度。通过这次协同，我们在提升大语言模型推理的 GPU 利用率方面获得了宝贵洞见。期待本文分享的技术方案与实践经验，能助力开发者群体在关键业务的大模型推理应用中更充分地发挥 NVIDIA GPU 的性能优势。