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

1. LLM 推理 - 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模型进行推理测试。

（这里为什么没有用最新的Llama3是因为在尝试做部署与推理Llama3-8B-Chinese-Chat模型的过程中遇到了一个暂时未解决的问题，具体报错为：RuntimeError: 【TensorRT-LLM】【ERROR】 Assertion failed: mpiSize == tp * pp (/home/jenkins/agent/workspace/LLM/release-0.10/L0_PostMerge/tensorrt_llm/cpp/tensorrt_llm/runtime/worldConfig.cpp:99)）。）

下面是具体过程：

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个特点：

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

在常规的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。在使用两个引擎的部署过程中，笔者对两者有如下感受：

1. vLLM非常易于使用：vLLM无论是环境准备（pip即可）、模型部署（不需要编译模型）还是上手难度（文档全面、社区活跃），均优于TensorRT-LLM。TensorRT-LLM的学习成本以及故障排查的难度都要高于vLLM
2. vLLM支持的硬件部署选项更多：TensorRT-LLM作为英伟达推出的框架，主要是针对自家的GPU进行的优化，所以仅适用于Nvidia CUDA。而vLLM目前除了支持Nvidia CUDA,外，还支持AMD ROCm, AWS Neuron和CPU的部署
3. 在性能表现上，参考社区在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：

1. No Batching：每次处理1个请求
2. Static Batching：请求放入batch队列，在队列满时一次性执行这个批次
3. Dynamic Batching：收到请求后放入batch队列，在队列满时或达到一定时间阈值，一次性执行这个批次
4. 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个问题：

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

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

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

图片二：Iteration Batching[8]

总结

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

References

[1] TensorRT-LLM Installing on Linux： https://nvidia.github.io/TensorRT-LLM/installation/linux.html

[2] Llama2-7b-chat： https://huggingface.co/meta-llama/Llama-2-7b-chat-hf

[3] Deploying LLMs Into Production Using TensorRT LLM: https://towardsdatascience.com/deploying-llms-into-production-using-tensorrt-llm-ed36e620dac4

[4] TensorRT-LLM 支持的模型： https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm

[5] 比HuggingFace快24倍！伯克利LLM推理系统开源碾压SOTA，GPU砍半： https://baijiahao.baidu.com/s?id=1769290162439240324&wfr=spider&for=pc

[6] 大模型推理优化技术-KV Cache： https://zhuanlan.zhihu.com/p/700197845

[7] BenchMarking LLM Inference Backends: vLLM, LMDeploy, MLC-LLM, TensorRT-LLM, and TGI https://bentoml.com/blog/benchmarking-llm-inference-backends

[8] Iteration batching (a.k.a. continuous batching) to increase LLM inference serving throughput： https://friendli.ai/blog/llm-iteration-batching/