数据中心/云端

借助 NVIDIA TensorRT-LLM 分块预填充提高 AI 推理效率和简化部署

在本 博文 中,我们详细介绍了分块预填充,这是 NVIDIA TensorRT-LLM 的一项功能,可提高 GPU 利用率并简化开发者的部署体验。本文基于之前的博文,讨论了 TensorRT-LLM 中的高级 KV 缓存优化功能如何在需要系统预填充的用例中将性能提升高达 5 倍。

传统预填充和解码推理方法面临的挑战

当用户向模型提交请求时,它将经历两个不同的计算阶段:预填充和解码。每个阶段使用 GPU 计算资源的方式各不相同。

在预填充阶段,系统会处理所有输入令牌以计算 KV 缓存,然后使用该缓存生成输出的第一个令牌。此阶段的计算要求很高,可以有效利用 GPU 庞大的并行计算资源。

在解码阶段,系统会单独生成输出 tokens,使用每个新 tokens 更新预填充阶段的中间状态。由于中间状态计算的繁重计算工作是在预填充阶段完成的,因此解码阶段主要涉及仅处理新生成的 tokens。因此,此阶段的计算密集度较低。

TensorRT-LLM 支持动态批处理,允许并行处理多个请求,使得预填充和解码阶段可以同时进行。虽然这种方法比暂停解码请求以处理预填充请求的分阶段批处理更高效,但由于解码阶段延迟到预填充请求完成,因此可能会导致延迟。图 1 的顶部说明了这一概念。

A visual showing a batch of three queries being processed with and without TensorRT-LLM chunked prefill. Without chunking, the decode phase has to wait until the full prefill phase is completed before the next decode can start. With chunking, the latency is shorter and the decode can be completed faster while at the same time increasing GPU utilization.
图 1. 分块预填充过程的简化表示,展示了如何提高 GPU 利用率并加速查询完成。

使用分块预填充平衡预填充和解码阶段

借助 TensorRT-LLM 分块预填充,令牌被分成较小的单元或分块,以加快处理速度。这可以防止预填充阶段成为瓶颈,使解码阶段的令牌实现更多的并行化,并提高 GPU 的利用率。图 1 的底部部分说明了这一概念。

此外,使用 TensorRT-LLM 分块预填充使 GPU 系统能够处理更长的上下文和更高的并发级别。由于内存使用量取决于每次迭代期间处理的令牌数量,因此使用预填充块可将内存消耗量与传入请求的上下文长度解耦。然后,系统可以在不增加内存需求的情况下处理更大的上下文长度。

通过动态数据块大小调整简化 TensorRT-LLM 引擎创建

图 1 显示了块大小百分比对用户交互性和整体系统吞吐量的影响,这需要根据所需的用户体验和可用 GPU 资源进行权衡。使用大块大小可以减少处理预填充序列所需的迭代次数,从而缩短第一个令牌(TTFT)的时间。

但是,这也会增加完成正在进行的请求的解码阶段所需的时间,从而增加查询完成时间并降低每秒输出令牌数(TPS)。对于开发者而言,找到合适的平衡有时非常耗时。TensorRT-LLM 部署动态块大小,根据 GPU 利用率指标为用户提供理想建议。

动态预填充块大小调整的另一个优势是简化了 TensorRT-LLM 引擎构建流程。过去,开发者必须提供最大输入序列长度,以便在引擎构建期间计算激活缓冲区。为确保激活缓冲区能够处理最复杂的用户查询,开发者需要为其用例设置预期的最大序列长度。这种方法针对异常和最差情况优化了 GPU 资源,导致内存使用效率低下。

借助动态预填充块大小调整,开发者无需再手动指定最大输入长度。相反,激活缓冲区大小由块大小(由 TensorRT-LLM 配置)自动确定。这消除了手动配置,并提高了执行期间的内存使用效率。

TensorRT-LLM 分块预填充入门

使用 TensorRT-LLM 分块预填充可显著提高系统性能和利用率。将传统的连续预填充阶段分解为更小、更易于管理的块,可通过解码阶段实现更好的并行化,从而减少瓶颈并加速查询完成。

块大小的动态调整还简化了 TensorRT-LLM 引擎配置流程,无需开发者手动设置激活缓冲区大小,而这在之前会导致内存使用效率低下。动态块大小优化内存使用,并确保系统能够无缝适应不同 LLM 用户的需求。

查看 我们的 GitHub 文档 ,开始使用 TensorRT-LLM 分块预填充。详细了解 NVIDIA AI 推理解决方案 ,并随时了解最新的 AI 推理性能更新

标签