LLM 推理：Nvidia TensorRT-LLM 与 Triton Inference Server

随着LLM越来越热门，LLM的推理服务也得到越来越多的关注与探索。在推理框架方面，tensorrt-llm是非常主流的开源框架，在Nvidia GPU上提供了多种优化，加速大语言模型的推理。但是，tensorrt-llm仅是一个推理框架，可以帮助我们完成模型的加载与推理。若是要应用在生产上，仍需要与推理服务例如Triton Inference Server进行配合。

本文会介绍TensorRT-LLM与Triton的基本原理，并展示使用这两个组件部署LLM推理服务的流程。同时也会介绍其中使用到的推理优化技术。

2. 基础环境

硬件：1 x A10 GPU（24GB显存）

操作系统：Ubuntu 22.04

3. TensorRT-LLM

TensorRT-LLM是一个易于使用的Python API,用于定义大型语言模型(LLM),并构建包含最先进优化的TensorRT引擎,以在NVIDIA GPU上高效执行推理。TensorRT-LLM包含用于创建执行这些TensorRT引擎的Python和C++运行时的组件。它还包括与NVIDIA Triton推理服务器集成的后端。使用TensorRT-LLM构建的模型可以在从单个GPU到多个节点(使用张量并行或/和管道并行)的多个GPU的广泛配置上执行。

为了最大化性能并减小内存占用,TensorRT-LLM允许使用不同的量化模式执行模型(参考支持矩阵)。TensorRT-LLM支持INT4或INT8权重(和FP16激活;又称INT4/INT8仅权重),以及SmoothQuant技术的完整实现。

下面我们先使用TensorRT-LLM来进行模型的推理，然后介绍TensorRT-LLM在提升推理性能上做的部分优化。

3.1. 设置TensorRT-LLM环境

下面我们参考TensorRT-LLM的官网[1]进行设置。

3.2. 模型推理

在设置好TensorRT-LLM的环境后，下面对llama2模型进行推理测试。
4. TensorRT-LLM原理

与其他部分推理引擎（例如vLLM）不一样的是：TensorRT-LLM并不是使用原始的weights进行模型推理。而是先对模型进行编译，并优化内核，实现在英伟达GPU上的高效运行。运行编译后模型的性能比直接运行原始模型的性能要高得多，这也是为什么TensorRT-LLM速度快的原因之一。

图片一：原始模型被编译成优化过的二进制文件[3]

原始模型的权重，与设定的优化选项（例如量化级别、tensor并行度、pipeline并行度等），一起传递给编译器。然后编译器根据这些信息输出特定针对GPU优化过的模型二进制文件。

需要注意的一个点是：整个模型编译的过程必须在GPU上进行。生成的编译模型是专门针对运行它的GPU进行优化的。例如，如果是在A40 GPU上编译的模型，就无法在A100上运行。所以在推理时，必须使用和编译时的GPU同一个类型。

TensorRT-LLM并不直接支持所有的LLM，因为每个模型的架构都不一样，TensorRT会做深层图级别优化，所以这就需要对不同模型进行适配。不过目前大部分模型例如Mistral、Llama、chatGLM、Baichuan、Qwen等都是支持的[4]。

TensorRT-LLM的python包使得开发者可以在不需要了解C++或CUDA的情况下以最高效的方式运行LLM。除此之外，它还提供了例如token streaming、paged attention以及KV cache等使用功能，下面我们具体进行介绍。

4.1. Paged Attention

Paged Attention无论是在TensorRT-LLM还是vLLM中，都是一个核心的加速功能。通过Paged Attention高效地管理attention中缓存的张量，实现了比HuggingFace Transformers高14-24倍的吞吐量[5]。

4.1.1. LLM的内存管理限制性能

在自回归解码过程中，LLM的所有输入token都会生成attention机制的keys和values的张量，并且这些张量被保留在GPU内存中，用来生成下一个token。大语言模型需要大量内存来存储每个token的keys和values，并且随着输入序列变长时，所需的内存也会扩展到非常大的规模。

这些缓存的key和value的张量通常称为KV缓存。KV缓存有2个特点：

内存占用大：在LLaMA-13B中，单个序列的KV缓存占用高达1.7GB的内存
动态化：其大小取决于序列的长度，而序列长度高度易变，且不可预测

在常规的attention机制中，keys和values的值必须连续存储。因此，即使我们能在给序列分配的内存的中间部分做了内存释放，也无法使用这部分空间供给其他序列，这会导致内存碎片和浪费。因此，有效管理KV缓存是一个重大挑战。

4.1.2. Paged Attention机制

为了解决上述的内存管理问题，Paged Attention的解决办法非常直接：它允许在非连续的内存空间中存储连续的keys和values。

具体来说，Paged Attention将每个序列的KV缓存分为若干个block，每个block包含固定数量token的key和value张量。在注意力计算过程中，Paged Attention内核能够高效地识别和提取这些block。由于这些block在内存中不需要连续，因此也就可以像操作系统的虚拟内存一样，以更灵活的方式管理key和value张量——将block看作page，token看作bytes，序列看作process。

序列的连续逻辑块通过块表映射到非连续的物理块。随着生成新的token，物理块会按需进行分配。这种方式可以防止内存碎片化，提高内存利用率。在生成输出序列时，page即可根据需要动态分配和释放。因此,如果我们中间释放了一些page，那这些空间就可以被重新用于存储其他序列的KV。

通过这种机制，可以使得极大减少模型在执行复杂采样算法（例如parallel sampling和beam search）的内存开销，提升推理吞吐性能。

4.2. KV Cache

KV Cache是LLM推理优化里的一个常用技术，可以在不影响计算精度的情况下，通过空间换时间的办法，提高推理性能。KV Cache发生在多个token生成的步骤中，并且只发生在Decoder-only模型中（例如GPT。BERT这样的encoder模型不是生成式模型，而是判别性模型）。

我们知道在生成式模型中，是基于现有的序列生成下一个token。我们给一个输入文本，模型会输出一个回答（长度为 N），其实该过程中执行了 N 次推理过程。模型一次推理只输出一个token，输出的 token 会与输入序列拼接在一起，然后作为下一次推理的输入，这样不断反复直到遇到终止符。

而由于每个token都需要计算其Key和Value（Attention机制），所会存在一个问题：每次生成新的token时，都需要计算之前token的KV，存在大量冗余计算。而在整个计算过程中，每个token的K、V计算方式是一样，所以可以把每个token（也就是生成过程中每个token推理时）的K、V放入内存进行缓存。这样在进行下一次token预测时，就可以不需要对之前的token KV再进行一次计算，从而实现推理加速。这便是KV Cache的基本原理，如果各位希望了解更多的细节，可以继续阅读这篇文章[6]。

5. TensorRT-LLM与vLLM对比

vLLM也是开源社区一个非常热门的推理引擎，上面提到的Page Attention这一重要就是源自vLLM。在使用两个引擎的部署过程中，笔者对两者有如下感受：

vLLM非常易于使用：vLLM无论是环境准备（pip即可）、模型部署（不需要编译模型）还是上手难度（文档全面、社区活跃），均优于TensorRT-LLM。TensorRT-LLM的学习成本以及故障排查的难度都要高于vLLM
vLLM支持的硬件部署选项更多：TensorRT-LLM作为英伟达推出的框架，主要是针对自家的GPU进行的优化，所以仅适用于Nvidia CUDA。而vLLM目前除了支持Nvidia CUDA,外，还支持AMD ROCm, AWS Neuron和CPU的部署
在性能表现上，参考社区在2024年6月使用LLAMA3-8B做的benchmark[7]测试，可以看到vLLM在TTFT（Time to First Token，即从发送请求到生成第一个token的时间，在交互式聊天场景较为重要）上表现不俗，优于TensorRT-LLM。但是在Token生成速率上TensorRT-LLM表现更好（特别是在做了4-bit量化后，两者表现差异更大）

总的来说，如果是团队有一定的技术能力、仅使用英伟达GPU，以及使用量化后的模型做推理的场景，可以优先考虑TensorRT-LLM，可以获得更高效的推理性能。

6. Nvidia Triton Inference Server

在TensorRT-LLM里我们可以看到它实现了模型的加载以及推理的调用。但是在生产上，我们不可能仅有这套API工具就足够，而还需要对应一套Serving服务，也就是一套Server框架，帮助我们部署各种模型，并以API的方式（例如HTTP/REST等）提供模型的推理服务。同时，还要能提供各种监控指标（例如GPU利用率、服务器吞吐、延迟等），可以帮助用户了解负载情况，并为后续扩缩容的推理服务的实现提供指标基础。

所以，要创建一个生产可用的LLM部署，可使用例如 Triton Inference Server的Serving服务，结合TensorRT-LLM的推理引擎，利用TensorRT-LLM C++ runtime实现快速推理执行。同时它还提供了in-flight batching 和paged KV caching等推理优化。

这里还需要再强调的是：Triton Inference Server，如其名字所示，它只是一个Server服务，本身不带模型推理功能，而是提供的Serving的服务。模型推理功能的实现，在Triton里是通过一个backend的抽象来实现的。TensorRT-LLM就是其中一种backend，可以对接到Triton Inference Server里，提供最终的模型推理功能。所以，Triton不仅仅是只能和TensorRT-LLM集成使用，还可以和其他推理引擎集成，例如vLLM。

在对Triton Inference Server有了简单了解后，下面我们介绍如何实现部署。

6.1. 部署推理服务

创建docker容器来部署推理服务

而后即可正常启动Triton Server推理服务：

6.2. 使用推理服务

在Triton Inference Server启动后，即可使用HTTP的方式进行模型推理，例如：

6.3. In-flight Batching

在部署过程中，大家可能有注意到：在指定tensorrt_llm的backend时，我们指定了一个batching_strategy:inflight_fused_batching的配置项。这里In-flight Batching也是模型推理场景里的一个优化项，一般也称为Continuous Batching。

首先，Batching（批处理）的基本原理是：将多个独立的推理请求组合为一个批次，然后一次性提交给模型进行推理，从而提高计算资源的利用率。在GPU上的推理请求可以以4种方式进行Batching：

No Batching：每次处理1个请求
Static Batching：请求放入batch队列，在队列满时一次性执行这个批次
Dynamic Batching：收到请求后放入batch队列，在队列满时或达到一定时间阈值，一次性执行这个批次
In-flight Batching/Continuous Batching：请求按照token逐个处理，如果有请求已经完成并释放了其占据的GPU资源，则新的请求可以直接进行处理，不需要等待整个批次完成

为什么需要In-flight Batching？我们可以想象一个例子，在Static Batching的场景下，会有多个请求提交给模型去处理。假设一个批次有2个请求，其中request_a的推理需要生成3个token（假设耗时3s），request_b的推理要生成100个token（假设耗时100s），所以当前这个批次需要耗时100s才能结束。同时又来了request_c的推理请求，此时request_c智能等待。这样会造成2个问题：

request_a只耗时3s即结束，但是由于它与request_b是同一个批次，所以无法在完成后立即返回结果给客户端，而是需要等待request_b也完成之后才能一起返回结果
新的请求request_c也必须等待前一个批次完成后（例如当前这个批次耗时100s），才能被模型进行推理

在这种方式下，可以看到，会有GPU资源空闲，得不到完全利用。并且会导致整体推理延迟上升。

而In-flight Batching的方式是：它可以动态修改构成当前批次的请求，即使是这个批次正在运行当中。还是以上面的场景举例，假设当前过去了3s，request_a已经完成，request_b仍需97s完成。这时候request_a由于已经完成，所以可以直接返回结果给客户端并结束。而由于request_a释放了其占据的资源，request_c的推理请求可以立即被处理。这种方式就提升了整个系统的GPU使用率，并降低了整体推理延迟。对应过程如下图所示：

总结

从TensorRT-LLM+Triton这套方法的部署过程来看，还是较为复杂的，需要先对模型进行编译，还要特别注意环境、参数的设置，稍不注意就无法正常部署。除此之外，英伟达的文档也不够清晰。总体来说，学习成本比较高，但也带来了相对应的推理效果提升的收益（相对于vLLM在量化模型推理、以及token生成速度的场景）。但是，在模型推理场景下，推理效果也并不是唯一选择因素，其他例如框架易用性、不同底层硬件的支持，也需要根据实际情况进行考量，选择最合适推理引擎。

如何学习AI大模型？

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述