调整最大批大小和最大 Token 数#

TensorRT-LLM 的关键功能之一是其 inflight 批处理调度器和运行时,它可以同时调度和执行上下文和生成阶段的请求。 最大 Token 数与最大批大小一起决定了调度器如何以及何时调度新的和当前的请求,了解它们的作用以及如何调整它们可以提供显着的性能优势。

免责声明: 虽然此处显示的性能数字是真实的,但仅用于演示目的。 环境、SKU、互连和工作负载的差异都会显着影响性能,并导致您的结果与此处显示的结果不同。

了解 TensorRT-LLM 调度器#

本节可视化 TensorRT-LLM 如何根据最大批大小和最大 Token 数调度请求。 该示例从新初始化的引擎以及一些已传入但未调度的请求开始。 为了便于演示,将玩具值设置为 最大 批大小 = 4最大 token = 12。 每个方块代表一个 token,其颜色代表它所属的请求。

TRT-LLM Scheduler Visualization 1

现在,调度器获取前两个请求(请求 1 和请求 2),并调度它们以执行上下文阶段。 但是,它无法调度更多请求,因为前两个请求的提示各有 5 个 token,由于最大 token 数的限制,剩余预算为 2 个 token。 由于所有剩余请求都具有超过 2 个提示 token,因此无法调度它们(在这种情况下,上下文分块可以提供帮助,请参阅下面的分页上下文注意力部分)。 这些 token 标有“C”,表示它们是在上下文阶段处理的提示 token。

注意:不同请求的 token 显示在不同的行上,仅用于可视化目的,并不代表实际的内存布局

TRT-LLM Scheduler Visualization 2

现在,引擎运行一次执行迭代,完成两个已调度请求的上下文阶段。 完成后,已创建这两个请求的提示的 kv-cache,并且已生成第一个 token。 生成的 token 标有“G(n)”,例如,标有“G1”的 token 表示它是为其请求生成的第一个 token。

TRT-LLM 首先优先调度生成阶段的请求,因此两个生成的 token 排队在下一次迭代中处理。 现在,由于先前调度的两个请求已进入生成阶段,并且仅占用最大 token 数预算 12 中的两个 token,因此调度程序能够调度两个额外的请求(请求 3 和请求 4)。 由于最大批大小限制为 4,因此即使最大 token 数预算中有空间,它也无法调度最后一个请求(请求 5)。

TRT-LLM Scheduler Visualization 3

在下一次执行迭代之后,请求 1 和 2 的第二个 token 已生成,并且请求 3 和 4 的第一个 token 已生成。 假设为请求 1 生成的 G2 是停止 token,表示请求 1 已完成。 在这种情况下,调度程序会在执行另一次迭代之前驱逐请求 1,并准备将其返回给用户。 此驱逐将引擎的状态置于低于最大批大小限制之下,并允许调度请求 5。

另一件事要注意的是,为请求 2 生成的 G1 已添加到请求 2 的 kv-cache 中,表示随着生成越来越多的 token,请求的 kv-cache 如何增长。

TRT-LLM Scheduler Visualization 4

总体而言,最大批大小和最大 token 数限制在决定何时实际执行请求方面起着关键作用,并且调整它们可能会对吞吐量数字以及引擎如何在生成阶段平衡先前调度的请求与新请求的上下文阶段产生重大影响

注意: 这提供了调度程序的简化可视化,以突出显示最大批大小和最大 token 数如何影响它。 调度程序还会考虑可用于 kv-cache 的可用空闲内存量,并且具有其他可配置选项,这些选项会影响其行为。 有关更多信息,请参阅其他选项页面的运行时标志。

调整最大批大小#

重要的是将最大批大小设置得足够大,以使其不会成为新请求调度的瓶颈。 因此,建议测试工作负载的最大批大小的附加值。 默认值为 2048。2 的幂是很好的初始值,可以进行扫描。

如何更改最大批大小#

您可以在构建配置中指定最大批大小。

build_config = BuildConfig(
    max_batch_size=512
)

如果您使用 CLI 流构建引擎,请将 --max_batch_size <int> 传递给 trtllm-build 以调整最大批大小。

调整案例研究#

继续我们关于 4 个 H100 上的 Llama-3.3-70B 的案例研究,我们从上一节(构建具有多个配置文件、gemm 插件、分页上下文注意力和启用 reduce 融合的引擎)中停止的地方继续。 扫描 64、512 和默认 2048 的最大批大小会产生以下结果

指标

最大批大小 64

最大批大小 512

最大批大小 2048

Token 吞吐量 (token/秒)

1944.3031

2466.7933

2044.2628

请求吞吐量 (req/秒)

0.9494

1.2045

0.9982

第一个 Token 的平均时间 (ms)

145.7607

147.7876

146.6628

平均 Token 间延迟 (ms)

14.6475

14.6554

14.4493

由此可见,最大批大小为 64 会导致瓶颈,而最大批大小为 512 则是最佳选择,吞吐量提高了近 20%,而对延迟的影响可以忽略不计。

调整最大 Token 数#

如果最大 token 数太小,则可能会成为请求调度的瓶颈。 但是,如果它太大,则可能会导致提示 token 占用太多内存(尤其是在长上下文工作负载中),而没有为 kv-cache 留下足够的空间,从而导致性能下降,甚至出现内存不足错误。 最大 token 数的默认值为 8192。 建议您扫描几个最大 token 数值,以确定最适合您的工作负载的数字。 要尝试的良好值是 2 的幂 >= 1024。 由于最大 token 数和最大批大小都会影响调度,因此如果可能,最好在它们的组合上进行网格搜索。

如何更改最大 Token 数#

与最大批处理大小类似,最大 token 数量也在构建配置中指定。

build_config = BuildConfig(
    max_batch_size=512
    max_num_tokens=2048
)

如果您使用通过 CLI 构建引擎的流程,请将 --max_num_tokens <int> 传递给 trtllm-build 以调整 max_num_tokens。

调优案例研究#

在最大批处理大小设置为 512 的情况下,扫描 2048、8192 和 16384 的最大 token 数量值,得到了以下性能数据。

指标

最大 Token 数量 2048

最大 Token 数量 8192

最大 Token 数量 16384

Token 吞吐量 (token/秒)

2474.2581

2466.7933

2461.0165

请求吞吐量 (req/秒)

1.2081

1.2045

1.2017

第一个 Token 的平均时间 (ms)

147.5742

147.7876

147.9623

平均 Token 间延迟 (ms)

14.6852

14.6554

14.6769

对于这个特定的工作负载,最大 token 数量为 2048 提供了最佳性能,但优势不是非常大。这反映了这样一个现实:对于任何给定的工作负载,调整各种标志可能会产生不同的影响。但是,检查不同的值以确保您不会因为调度不平衡而放弃大量收益非常重要。

重新审视分页上下文注意力机制和上下文分块#

之前 我们建议启用分页上下文注意力机制,即使在我们的案例研究中它对性能没有显著影响。现在我们了解了 TensorRT-LLM 调度器,我们可以解释为什么这有益。简而言之,我们建议启用它,因为它启用了上下文分块,这允许将请求的上下文阶段分解成多个部分并在多个执行迭代中处理,从而使引擎能够提供更稳定的上下文和生成阶段执行平衡。

TensorRT-LLM 调度器的可视化显示,最初无法调度请求 3,因为它会使调度器超过最大 token 数量限制。但是,通过上下文分块,情况不再如此,并且可以调度请求 3 的第一个块。

TRT-LLM Scheduler Visualization Chunked Context 1

这非常有益,原因有几个。首先,它消除了相对于最大 token 数量具有较大提示的请求,由于已在运行的其他请求而无法调度的可能性。在生产工作负载中,这有助于改善最坏情况下的 TTFT 数字。其次,它允许设置较小的最大 token 数量值,因为您不再需要最大 token 数量至少与您要支持的最长提示一样大。对于长上下文情况,这非常重要,因为设置非常大的最大 token 数量值会占用可用于作为 kv-cache 的内存。鉴于在最坏的情况下,分块上下文对性能的影响最小,但在许多情况下可以显着提高性能,因此建议您始终启用它。

结论#

TensorRT-LLM 调度器在性能方面起着重要作用,正确调整它可以显着提高性能。在本案例研究示例中,与上一页的结果相比,调整最大批处理大小和最大 token 数量可以显著提高性能

指标

启用构建时标志

调整了最大批处理大小和最大 Token 数量

% 提升

Token 吞吐量 (token/秒)

2044.2628

2474.2581

21.03

请求吞吐量 (req/秒)

0.9982

1.2081

21.03

第一个 Token 的平均时间 (ms)

146.6628

147.5742

-0.62

平均 Token 间延迟 (ms)

14.4493

14.6852

-1.63

解释这些结果,调整最大批处理大小和最大 token 数量显着提高了吞吐量,并在延迟方面保持了相当的水平(轻微下降在运行到运行的方差范围内)。与初始基线相比,调整后的引擎实现了以下提升

指标

基线

启用构建时标志并调整了最大批处理大小和最大 Token 数量

% 提升

Token 吞吐量 (token/秒)

1564.3040

2474.2581

58.17

请求吞吐量 (req/秒)

0.7638

1.2081

58.17

第一个 Token 的平均时间 (ms)

147.6976

147.5742

0.08

平均 Token 间延迟 (ms)

31.3276

14.6852

53.12