专栏名称: GiantPandaCV

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

[LLM推理优化][4w字] TensorRT-LLM部署调优-指北

GiantPandaCV · 公众号 · · 2024-07-02 18:18

正文

作者丨DefTruth

来源丨https://zhuanlan.zhihu.com/p/699333691

编辑丨GiantPandaCV

0x00 前言

注意 是“部署”调优，不是“性能”调优 ！因此本文与底层Kernel如果优化等无关， 主要关注应用层面 。本文记录一些使用TensorRT-LLM过程中，对性能有影响的参数的理解以及一些工具的用法。如果理解有误，欢迎指正。本文内容包括：

0x01 入门学习路线推荐(进行中)
0x02 Batch size相关的设置
0x03 影响首Token时延的配置
0x04 是否使用custom_all_reduce
0x05 影响Decode时延的配置
0x06 fp8/int8 KV Cache相关的设置
0x07 In-Flight Batching相关的设置
0x08 bls模式相关的设置
0x09 如何开启debug模式
0x0a tensorrtllm_backend使用问题
0x0b tensorrt_llm离线推理
0x0c 源码编译及benchmark工具使用
0x0d FP8/SQ/AWQ量化校准使用
0x0e 自定义FP8量化校准数据(进行中)
0x0f triton server镜像编译和使用
0x10 总结

好记性，不如烂笔头，本文长期更新，内容随缘（就看最近踩到了什么坑~）。内容比较流水账， 建议直接跳到想看的章节。

0x01 入门学习路线推荐

离线推理推荐先看llama示例，目前是最全面: https:// https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama
服务化也推荐看llama示例：https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/llama.md
推荐看完本文，估计能少踩不少坑~

0x02 Batch size相关的设置

max_batch_size

指允许进入engine的并行跑的最大请求数。对于显存足够的情况下，比如72B模型部署在总显存640G的机器，如果设置地太小，比如8，则会影响服务的吞吐。因为，最多只允许8个请求在engine中并行处理，其他的请求都在排队。另外，启动triton服务时，也需要指定max_batch_size，这个指的是config.pbtxt中的配置；triton server中的max_batch_size和我们在编译engine的时候指定的max_batch_size含义其实是不一样的。

Triton server配置中的max_batch_size: 这个是Triton server本身的dynamic_batching的遗留产物，比如我们在做CV模型的部署时，通常就需要结合triton中的max_batch_size这个参数和dynamic_batching来使用，从而实现动态组batch的功能。这里指的是，triton server的dynamic_batching功能，会把服务请求按照max_batch_size为最大颗粒度组成一个batch，然后再发给TensorRT-LLM处理。也就是triton server的max_batch_size，强调的组batch行为是triton server这个框架自带的特性，和TensorRT-LLM无关。

name: "tensorrt_llm"backend: "${triton_backend}"max_batch_size: ${triton_max_batch_size}

trtllm-build中的max_batch_size: 这个是指trtllm在编译engine的时候，engine支持的最大batch_size。使用过TensorRT的同学们应该对这个参数非常熟悉了。如果太大，可能会导致在编译engine阶段就OOM。

trtllm-build --checkpoint_dir ./tmp --output_dir ./engine --max_batch_size 8 ...

Anyway，为了避免出现各种奇怪的问题，tensorrt_llm/config.pbtxt、tensorrt_llm_bls/config.pbtxt以及trtllm-build中使用的max_batch_size最好保持一致。补充一下，由于tensorrtllm_backend中，还有ensemble(https://github.com/triton-inference-server/tensorrtllm_backend/tree/main/all_models/inflight_batcher_llm/ensemble)、preprocessing和postprocessing，因此需要把里边config.pbtxt的max_batch_size都配置成和tensorrt_llm/config.pbtxt中max_batch_size相同的值，否则无法启动服务（太多配置要改了...）

max_num_tokens

根据官方文档：Best Practices for Tuning the Performance of TensorRT-LLM(https://nvidia.github.io/TensorRT-LLM/performance/perf-best-practices.html) 中的介绍，max_num_tokens表示engine支持并行处理的最大tokens数，TensorRT-LLM需要为此预留部分的显存，此参数与max_batch_size存在相互制约的关系。由于TensorRT-LLM需要根据max_num_tokens预留显存，因此该值越大，留给KV Cache用的显存就越少。（这个理解需要验证~实际上是我自己没有很理解官方文档中的这个解释，但是又不敢承认...）

如何合理地设置max_batch_size和max_num_tokens？

max_num_tokens可以按照以下这个公式进行估计。其中alpha是0.0和1.0之间的浮点值，它代表在推理过程中每次调用前向函数时对context阶段请求数量的粗略估计。建议使用0.05到0.20之间（5%-20%之间）的值，但可能取决于实际场景。

max_batch_size * max_input_len * alpha + max_batch_size * max_beam_width * (1 - alpha)

当max_beam_width=1时，我们知道max_batch_size一般不会很大，比如128，而(1-alpha)小于1，因此，项max_batch_size * max_beam_width * (1 - alpha)可以忽略不计，max_num_tokens的计算公式简化为：

max_batch_size * max_input_len * alpha # alpha建议使用0.05到0.20之间

举个列子，当max_batch_size=64, max_input_len=1024, alpha为0.2，此时max_num_tokens大约为13107。对于没有指定 remove_input_padding 的情况，max_num_tokens不生效。

Anyway，就个人实践而言，不同的max_batch_size和max_num_tokens搭配时，确实对TTFT和TPOT有影响，但是怎么才能设置成最优值，对我来说依然是个巨大的 “迷” ，纯纯手工autotune。所以个人的建议就是，如果你不是为了压榨极致的性能，就使用TensorRT-LLM的默认值吧。

为什么build engine正常到启动服务却报OOM?

实践过程有遇到过一些max_batch_size不合理导致服务崩溃的情况。比如max_batch_size设置太大，在build engine阶段不报错，但是在启动triton server阶段报OOM（由于启动服务时发现，系统剩余显存，已经不足够你存放至少一条请求的KV Cache）。这是为啥呢？build engine不报错，反而跑服务崩了。因为max_batch_size是给TensorRT编译engine用的，编译engine的阶段，无法知道你后续启动服务时需要使用多少KV Cache（其实是可以算的吧，不然启动服务的时候咱会抛出这个信息，感觉是TensorRT-LLM没有在build engine阶段做这个功能），因此就有可能导致，对于某个max_batch_size，虽然编译engine没问题，但是启动服务的时候会报OOM。这个体验不是很好，会导致用户可能需要重复build engine，直到能试探到那个可正常用的max_batch_size（ 在max_batch_size的边缘疯狂试探...... ）

0x03 影响首Token时延的配置

max_queue_delay_microseconds

单位是微秒，1000微秒为1毫秒。表示请求在triton服务队列里为了dynamic_batching最大的等待时间。等到超过这个时间后，才会将请求发给TensorRT-LLM In-Flight Batching的等待队列，接着是IFB进行调度管理，将请求调度给引擎进行推理。如果max_queue_delay_microseconds设置太大，比如100000，则会导致请求被强制等待100ms，表现在LLM时延指标上，大概率就是TTFT慢了100ms。因此，这个值建议根据实际情况设置一个合理的值，避免首Token时延增加。比如：

dynamic_batching {    preferred_batch_size: [ 24 ]    max_queue_delay_microseconds: 100}

并且，从调度的角度来说，dynamic_batching组完batch之后，扔给IFB，实际上只是从一个等待队列扔给IFB的另一个等待队列，并不是直接给推理引擎直接推理了。因此，在开了IFB的情况下，这个max_queue_delay_microseconds的设置，可以设置成0即可，将对请求调度的管理完全交给TensorRT-LLM的IFB即可。不太理解为啥tensorrtllm_backend中的参考config.pbtxt还要保留这个配置。这个配置在使用IFB的情况下对用户具有比较大的迷惑性。首Token时延相关的设置还包括是否使用custom_all_reduce，具体看下一小节。其他的，想到再补充。

enable_kv_cache_reuse

TensorRT-LLM里边的KV Cache Reuse功能，指的就是Automatic Prefix Caching功能，具体的实现方式未知，因为这部分代码闭源。enable_kv_cache_reuse开启后， 主要是影响首Token时延，即TTFT 。对于具有较长system prompt或者多轮对话等场景，可以使用所有请求复用system prompt中的KV Cache，而不需要重新计算，从而降低TTFT的耗时。vLLM中有相同的功能，具体的实现原理可以阅读我写的另一篇文章：

DefTruth：[Prefill优化][万字] 原理&图解vLLM Automatic Prefix Cache(RadixAttention): 首Token时延优化 https://zhuanlan.zhihu.com/p/693556044

在TensorRT-LLM中使用Automatic Prefix Caching功能需要打开enable_kv_cache_reuse开关，比如：

gptManagerBenchmark --enable_kv_cache_reuse enable

并且在build engine阶段，需要开启use_paged_context_fmha，即在context阶段使用fused multihead attention kernel。具体示例如下：

trtllm-build --use_paged_context_fmha enable

对应到triton server，则需要修改all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt中的配置：

parameters: {  key: "enable_kv_cache_reuse"  value: {    string_value: "true"  }}

另外，特别需要注意的是，在使用enable_kv_cache_reuse功能时，有两个比较重要的参数是可以调优的，分别是tokens_per_block以及kv_host_cache_bytes。由于KV Cache是按照block的颗粒度进行reuse的，因此，tokens_per_block的值将决定了KV Cache reuse的边界情况是如果表现的。比如，tokens_per_block=1024，这将意味着，任何<1024 tokens的prompt以及任一prompt的最后不足1024 tokens的KV Cache Block都无法被不同的request复用。也就是说，tokens_per_block越小，能够被prefix caches命中的tokens可能就越多（这和prefix caching的特性有关，推荐阅读：DefTruth：[Prefill优化][万字] 原理&图解vLLM Automatic Prefix Cache(RadixAttention): 首Token时延优化）(https://zhuanlan.zhihu.com/p/693556044)。默认情况下这个值是128（最新的已经修改成64了），但是在使用了prefix caching功能情况下，可以考虑调整成更小的值，比如32/16等。

trtllm-build --tokens_per_block 32 ...

kv_host_cache_bytes则是指定使用多少bytes的CPU主机内存，来保存swap出来的KV Cache。当开启enable_kv_cache_reuse时，如果服务系统达到可支持的QPS上限，则可能会导致大量的KV Cache被逐出，如果此时设置了kv_host_cache_bytes，比如45G，则逐出的KV Cache会被暂存在CPU内存中，以备后续被换入，从而增加KV Cache Reuse的可能性。但是，recompute和swap到底哪个更高效，也是不好说的，得根据实际的机器进行调试。

parameters: {  key: "kv_cache_host_memory_bytes"  value: {    string_value: "45000000000"  }}

更多关于enable_kv_cache_reuse的使用，可以参考TensorRT-LLM的文档：kv_cache_reuse.md(https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/kv_cache_reuse.md)

use_fp8_context_fmha

对于Hopper架构，并且使用FP8量化的模型，可以考虑开启use_fp8_context_fmha，以使用FP8 Context FMHA Kernel，对Prefill阶段进行加速。这个功能目前只支持Hopper架构，期待TensorRT-LLM能够在更多有FP8的架构上支持FP8 Context FMHA。该功能主要影响首Token时延。

0x04 是否使用custom_all_reduce

use_custom_all_reduce

custom_all_reduce目前是自动的开启的（不确定之后会不会改成默认不开启）。custom_all_reduce会受到NVLink、P2P通信、是否跨NUMA通信等情况的影响。以下，是关于NVLink、P2P、NUMA的简单介绍：

[1] NVLink 是英伟达（NVIDIA）开发并推出的一种总线及其通信协议。NVLink采用点对点结构、串列传输，用于中央处理器（CPU）与图形处理器（GPU）之间的连接，也可用于多个图形处理器之间的相互连接。与PCI Express不同，一个设备可以包含多个NVLink，并且设备之间采用网格网络而非中心集线器方式进行通信。该协议于2014年3月首次发布，采用专有的高速信号互连技术（NVHS）。该技术支持同一节点上GPU之间的全互联，并经过多代演进，提高了高性能计算应用中的双向带宽性能。参考：知北游：GPU通信技术

[2] P2P 是指NVIDIA GPUDirect P2P通信，NVIDIA GPUDirect是Magnum IO的一部分，作为一种技术，可以增强GPU中心的数据移动与访问。通过GPUDirect技术，网络适配器和存储驱动可以直接读取GPU内存，避免不必要的内存复制。其优点是：降低CPU负载，减小延迟，带来显著的性能提升。支持功能包括GPUDirect Peer to Peer (P2P)、GPUDirect Remote Direct Memory Access (RDMA)等；参考：一棵红杉树：一文读懂GPU通信互联技术(https://zhuanlan.zhihu.com/p/693806383)

[3] NUMA 服务器的基本特征是具有多个 CPU 模块，每个 CPU 模块由多个 CPU( 如 4 个 ) 组成，并且具有独立的本地内存、 I/O 槽口等。由于其节点之间可以通过互联模块 ( 如称为 Crossbar Switch) 进行连接和信息交互，因此每个人都有 CPU 可以访问整个系统的内存 ( 这是 NUMA 系统与 MPP 系统的重要差别 ) 。显然，访问本地内存的速度将远远高于访问远地内存 ( 系统内其它节点的内存 ) 的速度，这也是非一致存储访问 NUMA 的由来。由于这个特点，为了更好地发挥系统性能，开发应用程序时需要尽量减少不同 CPU 模块之间的信息交互。参考：Linux编程用C：一文掌握CPU的SMP与NUMA架构！(https://zhuanlan.zhihu.com/p/677797033)

[4] NIC 是指NIC网卡，一般安装在计算机或服务器上，通过网络与另一台计算机、服务器或其他网络设备进行通信。如今市场上网卡类型众多，但主要以有线网卡和无线网卡为主，其中无线网卡利用无线技术访问网络，而有线网卡需要使用DAC或AOC或光模块和光纤跳线进行连接。目前局域网基本上都是采用的以太网技术，根据网卡应用领域的不同可分为计算机网卡和服务器网卡。对于客户端计算机而言，一般情况下使用一个网卡即可，但对于服务器而言，则需要使用多个网卡来满足处理更多网络流量的需求。通常，计算机网卡都只有一个网络接口，而服务器网卡拥有多个网络接口，如双端口、四端口。参考：飞速FS：100G网卡NIC详细介绍及其发展趋势分析(https://zhuanlan.zhihu.com/p/370111197)

就个人的实践经验而言，在支持NVLink的机器上，开启custom_all_reduce对TTFT和TPOT均有性能收益。在仅仅支持P2P通信，但是不跨NUMA通信的情况下，也推荐开启custom_all_reduce，或许会有性能收益；在仅支持P2P通信，但是需要跨NUMA通信（比如8卡机器，0-3卡属于一个NUMA，4-7卡属于另一个NUMA），也不推荐使用custom_all_reduce；当连P2P都不支持时，就只能关掉custom_all_reduce了。具体到TTFT和TPOT，TTFT的通信量更大，TPOT的通信量较小，因此是否使用custom_all_reduce对TTFT耗时的影响更为明显。最后，推荐始终使用最新的驱动。

NVLink	P2P	不跨NUMA通信	跨NUMA通信	驱动版本	use_custom_all_reduce
支持				推荐>=550	True
不支持	支持	是		推荐>=550	True
不支持	支持		是	推荐>=550	False
不支持	不支持			推荐>=550	False

通过nvidia-smi topo --matrix可以查看当前的机器的GPU-CPU通信拓扑。以下是一个从vllm issue (https://github.com/vllm-project/vllm/issues/5094)里边捞的日志：

nvidia-smi topo --matrix
GPU Topology:
        GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7    NIC0    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X      PIX     PHB     PHB     SYS     SYS     SYS     SYS     PHB     0-13,28-41      0               N/A
GPU1    PIX      X      PHB     PHB     SYS     SYS     SYS     SYS     PHB     0-13,28-41      0               N/A
GPU2    PHB     PHB      X      PIX     SYS     SYS     SYS     SYS     PHB     0-13,28-41      0               N/A
GPU3    PHB     PHB     PIX      X      SYS     SYS     SYS     SYS     PHB     0-13,28-41      0               N/A
GPU4    SYS     SYS     SYS     SYS      X      PIX     PHB     PHB     SYS     14-27,42-55     1               N/A
GPU5    SYS     SYS     SYS     SYS     PIX      X      PHB     PHB     SYS     14-27,42-55     1               N/A
GPU6    SYS     SYS     SYS     SYS     PHB     PHB      X      PIX     SYS     14-27,42-55     1               N/A
GPU7    SYS     SYS     SYS     SYS     PHB     PHB     PIX      X      SYS     14-27,42-55     1               N/A
NIC0    PHB     PHB     PHB     PHB     SYS     SYS     SYS     SYS      X

Legend:  X    = Self  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinksNIC Legend:

  NIC0: rocep1s0

X : 表示当前卡自己和自己的连接； SYS: 表示通过PCIe进行跨NUMA的通信，比如这个示例中的GPU4和GPU0的topo，被标记为SYS，说明这两个卡的通信是需要跨NUMA的； PHB、PXB和PIX 都是表示通过PCIe进行连接，但连接方式不同，被标记为这些值的GPU之间一般是不跨NUMA的； NIC0 表示这个机器具有一个NIC网卡，可以和所有的卡进行连接； NV# 表示卡间是通过NVLink进行连接，这个示例没有NVLink；再捞一个带NVLink的示例(H20)：

GPU Topology:
  GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7 NIC0 NIC1 NIC2 NIC3 NIC4 CPU Affinity NUMA Affinity GPU NUMA ID
GPU0 X NV18 NV18 NV18 NV18 NV18 NV18 NV18 NODE NODE NODE SYS SYS 0-47,96-143 0 N/A
GPU1 NV18 X NV18 NV18 NV18 NV18 NV18 NV18 NODE PIX NODE SYS SYS 0-47,96-143 0 N/A
GPU2 NV18 NV18 X NV18 NV18 NV18 NV18 NV18 NODE NODE NODE SYS SYS 0-47,96-143 0 N/A
GPU3 NV18 NV18 NV18 X NV18 NV18 NV18 NV18 NODE NODE PIX SYS SYS 0-47,96-143 0 N/A
GPU4 NV18 NV18 NV18 NV18 X NV18 NV18 NV18 SYS SYS SYS PIX NODE 48-95,144-191 1 N/A
GPU5 NV18 NV18 NV18 NV18 NV18 X NV18 NV18 SYS SYS SYS NODE NODE 48-95,144-191 1 N/A
GPU6 NV18 NV18 NV18 NV18 NV18 NV18 X NV18 SYS SYS SYS NODE PIX 48-95,144-191 1 N/A
GPU7 NV18 NV18 NV18 NV18 NV18 NV18 NV18 X SYS SYS SYS NODE NODE 48-95,144-191 1 N/A
NIC0 NODE NODE NODE NODE SYS SYS SYS SYS X NODE NODE SYS SYS
NIC1 NODE PIX NODE NODE SYS SYS SYS SYS NODE X NODE SYS SYS
NIC2 NODE NODE NODE PIX SYS SYS SYS SYS NODE NODE X SYS SYS
NIC3 SYS SYS SYS SYS PIX NODE NODE NODE SYS SYS SYS X NODE
NIC4 SYS SYS SYS SYS NODE NODE PIX NODE SYS SYS SYS NODE X

更多GPU/CPU通信相关，目前有很多文章介绍，推荐阅读以下文章：

[1] 知北游：GPU通信技术（介绍了常见的“数据通信路径”，NVLink/P2P/GDR/GDS等）(https://zhuanlan.zhihu.com/p/693806383)
[2] 一棵红杉树：一文读懂GPU通信互联技术 (https://zhuanlan.zhihu.com/p/684116871)
[3] Linux编程用C：一文掌握CPU的SMP与NUMA架构！ (https://zhuanlan.zhihu.com/p/677797033)
[4] 飞速FS：100G网卡NIC详细介绍及其发展趋势分析 (https://zhuanlan.zhihu.com/p/370111197)
[5] 函谷叨客：【研究综述】浅谈GPU通信和PCIe P2P DMA (https://zhuanlan.zhihu.com/p/430101220)
[6] 神经蛙没头脑：高性能GPU服务器AI网络架构（上篇） (https://zhuanlan.zhihu.com/p/691035405)
[7] 神经蛙没头脑：高性能GPU服务器AI网络架构（下篇） (https://zhuanlan.zhihu.com/p/691061687)
[8] 极致Linux内核：什么是NIC（网络接口卡）？ (https://zhuanlan.zhihu.com/p/618745275)

0x05 影响Decode时延的配置

disable_xqa

TensorRT-LLM针对Decoding阶段以及MQA/GQA，提出了XQA优化技术。按照官方文档的说明，目前依然是一个实验性的feature。XQA主要影响Decode阶段的时延。XQA目前支持以下优化：

1. FP16 / BF16 compute data type.
2. FP16 / BF16 / FP8 / INT8 KV cache data type.
3. Paged KV cache (64 / 128 tokens per block).

XQA是默认开启的，但是可以通过在build engine阶段指定--disable_xqa来关闭。另外，需要注意的是，虽然XQA是默认开启的，但是不代表最后一定会使用XQA，因为TensorRT-LLM内部会通过启发式算法来决定实际使用XQA还是使用Masked MHA Kernel。不过，我们可以通过指定环境变量来强制使用XQA：

export TRTLLM_FORCE_XQA=1 # 在trtllm-build之前设置

enable_chunked_context

chunk context功能和vLLM的chunk-prefills功能类似，在prompt比较长的情况下可以考虑开启，比如>=2048。chunk context主要是对Decode时延进行优化；将prompt进行chunk，并和decode阶段的request组batch进行推理。Chunk Prefills相关论文为《SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills》；但是在TensorRT-LLM的IFB模式下，已经是每个request单独使用一个decode stream进行推理，不同的request是交替运行的，IFB实际上Decode优先的调度策略；而vLLM中的continuos batching是首Token优先的调度策略；因此在IFB模式下，开启enable_chunked_context，应该不会有特别明显的性能提升(TODO: 后续有更详细的实验对比，再更新一下)。

Chunk Prefills也比较有意思，论文看了几遍，之后抽空来补一篇解说。

use_fused_mlp

该选项会把MLP部分融合成一个kernel，开启后对decode阶段的TPOT的性能会有一定提升，个人经验大概是1%~2.5%左右。如果发现对模型效果没有影响，建议开启。

trtllm-build --use_fused_mlp

multi_block_mode

当需要应用的场景是小batch场景时（比如注重时延的Chat场景，服务的吞吐不会很高），并且input_seq_len大于1024时，可以考虑开启multi_block_mode。但是multi_block_mode这个flag只是一个runtime运行时的建议，就算指定了，如果TRT-LLM发现运行时没有性能收益，则不会使用multi_block_mode。不太确定这个multi_block_mode的原理是否和FlashDecoding相似。参考： https:// https://nvidia.github.io/TensorRT-LLM/performance/perf-best-practices.html

0x06 fp8/int8 KV Cache相关的设置

kv_cache_scaling_factor

一般情况下，模型的以INT8 或 FP8 运行时，GPT Attention也可以使用不同的数值精度，比如 FP32、FP16 和 BFloat16 。不过，TensorRT-LLM 支持 INT8 和 FP8 的 KV Cache，即QuantMode.INT8_KV_CACHE 和 QuantMode.FP8_KV_CACHE。GPT Attention运算过程中会保存 KV Cache。当启用 INT8 或 FP8 KV 缓存时，必须使用缩放因子将输入值量化为 8 位。对于量化，缩放因子存储在 kv_cache_scaling_factor 张量中。它的形状是[1]，当前版本仅支持每张量量化。量化使用倒数比例，因为它在plugin中乘以 fp_value * (1.0 / kv_cache_scaling_factor) 。在Decode/Generate期间，从KV Cache中读取的值在 MHA/MQA Kernel中在线反量化，反量化可以描述为 quantized_value * kv_cache_scaling_factor。在covert_checkpoint阶段我们可以指定--int8_kv_cache来使用Int8 KV Cache缓存。

目前，TensorRT-LLM中LLaMA的示例是最为完整的，建议入门先看LLaMA。这个贴一下LLaMA KV Cache INT8 + Weight Only Int8混合使用的例子。原理不展开讲了，不是本文的重点。

# Build model with both INT8 weight-only and INT8 KV cache enabledpython convert_checkpoint.py --model_dir ./llama-models/llama-7b-hf   \                             --output_dir ./tllm_checkpoint_1gpu_int8_kv_wq \                             --dtype float16  \                             --int8_kv_cache \                             --use_weight_only \                             --weight_only_precision int8

trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_int8_kv_wq \            --output_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \            --gemm_plugin auto \            --multi_block_mode enable \

0x07 In-Flight Batching相关的设置

不同的调度策略，会影响对请求的处理方式。一般来说，IFB里边直接使用默认的调度策略即可，但是多了解一些总没坏处。有些情况下，我们可以考虑根据实际的业务场景来选择合适的调度策略。比如，我们经常用吞吐量来衡量服务的好坏，比如 tokens/s，或者 reqs/s；但是站在用户的角度来说，并不是 tokens/s 或 reqs/s越大，用户需要等待的时延就越小。相反，如果你想尽可能提高tokens/s 或 reqs/s，那必然意味着每个请求本身完成的时延会变大。比如，你服务的吞吐很大，1w个请求，在第100秒的那一瞬间一起完成并返回了。这时，reqs/s = 100，意味着平均每秒能处理100个请求。指标上看，服务性能不错。但是呢，这种情况，对于用户来说，其实很糟糕，因为每个用户，实际都等待了100s。那么，这种时候，默认的调度策略不一定是最好的。

batch_scheduler_policy

batch_scheduler_policy指的是IFB的调度策略。目前，包括两种调度策略，即 MAX_UTILIZATION 和 GUARANTEED_NO_EVICT ；当启用IFB时，MAX_UTILIZATION 在表示在每次forward时打包尽可能多的请求。它通过尽可能多地调度服务请求来最大化 GPU 的利用率，如果达到 KV 缓存大小限制时，正在迭代中的一些请求则会被Evict（逐出）。而相对的，GUARANTEED_NO_EVICT策略，则会确保请求不会被逐出。MAX_UTILIZATION 适用在追求高吞吐的场景，GUARANTEED_NO_EVICT则更适合在关注用户体验的场景，因为它会确保请求不会被强制逐出，从而导致额外的重计算耗时。

这里补充一下自己对TensorRT-LLM In-Flight Batching的理解：

0x08 bls模式相关的设置

tensorrtllm_backend支持一种访问服务的模式，即BLS模式（Business Logic Scripting( https://github.com/triton-inference-server/python_backend#business-logic-scripting)）。我们知道，通常情况下，triton sever在服务启动后，输入输出是固定的，根据config.pbtxt来决定。但是在LLM的后处理中，我们经常需要采取各种各样的sampling策略，充斥着大量的条件判断以及采样策略。这些都是动态变化的，采样策略需要根据具体传入的条件值进行不同的处理。因此，tensorrtllm_backend就整了一个BLS模式，来解决这个问题。（BLS这个名称不太好理解，说简单点，就是sampling后处理大杂烩的缝合......）。BLS模式介绍文档： https://github.com/triton-inference-server/python_backend#business-logic-scripting

bls_instance_count

按照tensorrtllm_backend文档的建议，这个值需要设置成和trtllm-build engine时使用max_batch_size相同的值，以确保服务能够对请求进行并发处理。如果设置为1，那可能会导致服务变成串行处理的（当然，如果你不使用BLS模式，应该不会有这个问题）。具体实践中，我尝试过instance_count=1以及instance_count>=max_batch_size，对于前者，会发现每个请求处理的时间变长了，具体表现在TTFT耗时明显变长。当instance_count>=max_batch_size时，能够确保足够的并行性。

accumulate_tokens

BLS 模式中有个叫做accumulate_tokens的参数可以在流式请求中使用。当该参数为True时，BLS在调用postprocessing models时，会使用累计的token ids，而不仅仅是当前step的单个token id；这会影响tokenizer解码的结果，特别是，某些special token实际上是由几个token ids组成的时候，需要设置accumulate_tokens为True。

parameters: {  key: "accumulate_tokens"  value: {    string_value: "True"  }}

0x09 如何开启debug模式

log-verbose

通过指定log-verbose=3, 可以在启动triton server后，看到详细的运行日志。包括每个step的状态和运行信息等。launch_triton_server.py中如果开启的log模式，则使用的log-verbose就是3。

输出日志：

[TensorRT-LLM][INFO] {"Active Request Count":249,"Context Requests":8,"Free KV cache blocks":0,"Generation Requests":231,"Iteration Counter":90,"Max KV cache blocks":2448,"Max Request Count":256,"MicroBatch ID":0,"Runtime CPU Memory Usage":28784,"Runtime GPU Memory Usage":540173600,"Runtime Pinned Memory Usage":0,"Scheduled Requests":239,"Timestamp":"12-13-2023 14:55:14","Tokens per KV cache block":128,"Total Context Tokens":6904,"Used KV cache blocks":2448}

Performance Analysis

TensorRT-LLM性能分析可以参考官方文档：

Performance Analysis nvidia.github.io/TensorRT-LLM/performance/perf-analysis.html

0x0a tensorrtllm_backend使用问题

本小节记录一些tensorrtllm_backend的使用问题，流水账。由于tensorrtllm_backend和TensorRT-LLM是属于两个分离的模块，因此使用上感觉不是很丝滑。也容易引入一些额外的使用问题。目前看到社区所有的LLM推理框架，server和推理引擎都是一个整体的，比如vllm，你不需要额外去理解FastAPI的概念，只需要关注vllm.entrypoint.api_server即可，使用起来也很简单。但是目前，tensorrtllm_backend和TensorRT-LLM是分开的，当用户想要跑个服务时，还必须熟悉Triton Server这一套，不然TensorRT-LLM也无法用起来。这里也记录一下使用tensorrtllm_backend时需要注意的问题。

版本一致性问题

tensorrtllm_backend和TensorRT-LLM的版本目前是严格对应的。或者说，tensorrtllm_backend里边的triton models应该暂时还没考虑向后兼容。你用commit 480的tensorrtllm_backend编译的server镜像来跑某个<480的commit里边的triton models，有可能是跑不通的。因此，当你升级TensorRT-LLM和tensorrtllm_backend时，需要注意把models也都换了。不然，大概率会出现一些奇怪的bugs。

使用体验问题

因为server代码和TensorRT-LLM框架是分离的，导致每次TensorRT-LLM和tensorrtllm_backend升级，用户都得手动拷贝server代码，并且检查到底更新了啥，哪些地方是和现在业务里边用的产生了冲突。从用户体验上来说，这是一个非常割裂和Ugly的过程。举个例子，试想一下，假设你现在用的是vllm，然后vllm的server代码在另外一个repo，现在vllm框架升级了，但是server代码不在框架内，你得手动拷贝server代码来用，并检查有什么需要修改和注意的。这是一件非常痛苦的事情。比较希望TensorRT-LLM能够将server整合到一起，提供一个简单的server使用方式，比如：

python3 -m tensort_llm.entrypoints.api_server \    --model xxx \    --tensor-parallel-size 2 \    --max-model-len 4096 \    --trust-remote-code \    --disable-custom-all-reduce \    --enable-prefix-caching ...

PS: 这不禁让我想起了TensorFlow和PyTorch的静态图vs动态图的故事。目前看来，由于triton server和TensorRT-LLM本来就是两个框架，虽然不可能进行整合，但是在TensorRT-LLM中增加一个server模块，基于triton server封装统一的API，从而将server相关的代码也整合进TensorRT-LLM，这应该是完全可能的。似乎lmdeploy目前就是这个做法，上手用户体验感觉不错。所以，这个用户体验的问题完全是可以解决的。

目前使用下来，最深刻的体验就是，花在非性能问题的排查上的时间是最多的，特别是triton server。

OpenAI协议问题

tensorrtllm_backend目前似乎没有支持OpenAI协议，因此，如果需要使用OpenAI协议，还得用户自己包一层FastAPI，这个过程，有难免会引入新的问题，或者新的bug，排查起来，也不是很高效。大家都支持，没搞懂为啥TensorRT-LLM不支持？为啥要提这个点，还是回到实际的应用，如果已经有应用使用了OpenAI协议，那为了使用TRT-LLM还得做对应的修改。

其他

嗯，是的。本文部分内容存在吐槽的倾向。当然，吐槽归吐槽，该用还是老老实实用着 （你就说他快不快吧）。总之，还是非常感谢TensorRT-LLM团队将这么优秀的工作开源到社区 （虽然是部分开源） 。Anyway，祝TensorRT-LLM越做越好 。最后，再补画一个自己会比较喜欢的使用模式。

trtllm-launcher --model Qwen/Qwen1.5-72B-Chat --tensor-parallel-size 8 --enable-kv-cache-reuse --use-custom-all-reduce --enforce-xqa ...

0x0b tensorrt_llm离线推理

ModelRunner和ModelRunnerCpp的不统一

最近想在多模态场景下将examples中的ModelRunner切换成ModelRunnerCpp，以便可以使用prefix caching的功能。似乎是多模态的模型还不支持使用ModelRunnerCpp。不过为啥会存在两套API和功能都不完全对齐的ModelRunner呢？还不止这个。ModelRunnerCpp支持更多的功能，现在源码编译tensorrt_llm就会默认编译python_binding。ModelRunnerCpp支持Prefix Caching和Chunk Context，ModelRunner不支持。并且对应的SamplingConfig虽然是同名，但是是两个不同的类。期待API完全统一。

ModelRunnerCpp使用

可以参考示例： https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/summarize.py

   # 初始化ModelRunnerCpp
   if test_trt_llm:
        if not PYTHON_BINDINGS and not args.use_py_session:
            logger.warning(
                "Python bindings of C++ session is unavailable, fallback to Python session."
            )
            args.use_py_session = True
        runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp
        runner_kwargs = dict(engine_dir=args.engine_dir,
                             rank=runtime_rank,
                             debug_mode=args.debug_mode,
                             gpu_weights_percent=args.gpu_weights_percent)
        if args.medusa_choices is not None




    
:
            args.medusa_choices = ast.literal_eval(args.medusa_choices)
            assert args.temperature == 1.0, "Medusa should use temperature == 1.0"
            assert args.num_beams == 1, "Medusa should use num_beams == 1"
            runner_kwargs.update(medusa_choices=args.medusa_choices)
        if not args.use_py_session:
            runner_kwargs.update(
                max_batch_size=max_batch_size,
                max_input_len=test_token_num,
                max_output_len=output_len,
                max_beam_width=num_beams,
                max_attention_window_size=max_attention_window_size,
                sink_token_length=sink_token_length,
                max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache,
                kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, # 是否用prefix caching
                kv_cache_free_gpu_memory_fraction=args.
                kv_cache_free_gpu_memory_fraction,
                enable_chunked_context=args.enable_chunked_context,
            )
        runner = runner_cls.from_dir(**runner_kwargs)
     # 调用generate接口
     with torch.no_grad():
            outputs = runner.generate(
                batch_input_ids,
                max_new_tokens=output_len,
                max_attention_window_size=max_attention_window_size,
                sink_token_length=sink_token_length,
                end_id=end_id,
                pad_id=pad_id,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                stop_words_list=stop_words_list,
                bad_words_list=bad_words_list,
                num_beams=num_beams,
                length_penalty=length_penalty,
                early_stopping=early_stopping,
                repetition_penalty=repetition_penalty,
                presence_penalty=presence_penalty,
                frequency_penalty=frequency_penalty,
                lora_uids=args.lora_task_uids,
                output_sequence_lengths=True,
                return_dict=True,
                medusa_choices=args.medusa_choices)
            torch.cuda.synchronize()

提示： 不推荐使用SamplingConfig来指定采样参数，目前SamplingConfig很乱 ，在python实现里边有两个同名的类（最新的更新似乎已经删除了sampling_config这个参数？up to 202406011），一个是binding C++的SamplingConfig，一个纯python实现给ModelRunner（不是ModelRunnerCpp）用的。

PS: 个人觉得这个离线推理的示例还是过于复杂了，我还是喜欢vLLM这种简单明了的风格：

ModelRunner

不推荐使用，目测只是个临时的产物，极大概率是会抛弃的。 它引入了很多使用上的困惑和模糊，功能支持也没有ModelRunnerCpp完善，比如不支持prefix caching、chunk context等。

High-Level API

TensorRT-LLM目前正在开发High-Level API，看着使用方式比较自然，不过目前只支持LLaMA系列的模型。先关注着，看下后续的发展（静待花开）。示例参考： https:// https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/high-level-api/llm_examples.py

High-Level API把LLM离线推理拆解成3个步骤：config -> init -> generate

0x0c 源码编译及benchmark工具使用

源码编译

这里简单贴一下我源码编译的命令，仅做参考，不保证在其他环境能跑通。我的基础镜像和CUDA版本:nvcr.io/nvidia/pytorch:24.02-py3(https://nvcr.io/nvidia/pytorch:24.02-py3)、CUDA 12.3以及TensorRT 10.

apt-get update 
apt-get install sudo 
# prepare for mpi4py buildapt install openmpi-bin libopenmpi-dev 

unset CPLUS_INCLUDE_PATH && unset C_INCLUDE_PATHexport CPLUS_INCLUDE_PATH=//usr/local/mpi/include/:$CPLUS_INCLUDE_PATHexport C_INCLUDE_PATH=/usr/local/mpi/include/:$C_INCLUDE_PATHexport LD_LIBRARY_PATH=/opt/hpcx/ucx/lib:/opt/hpcx/ompi/lib:/usr/lib/x86_64-linux-gnu/:/usr/local/lib/python3.10/dist-packages/torch/lib:/usr/local/cuda/lib64:/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64# 安装TensorRT 10.0.1.6wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.0.1/tars/TensorRT-10.0.1.6.linux.x86_64-gnu.cuda-12.4.tar.gz --no-check-certificate
tar -xvf TensorRT-10.0.1.6.linux.x86_64-gnu.cuda-12.4.tar.gz
rm -f $(find /usr/lib/x86_64-linux-gnu -name "libnvinfer*.so*")rm -f $(find /usr/lib/x86_64-linux-gnu -name "libnv*parser*.so*")cp TensorRT-10.0.1.6/lib/*.so* /usr/lib/x86_64-linux-gnu
cp TensorRT-10.0.1.6/bin/trtexec /opt/tensorrt/trtexec
cp TensorRT-10.0.1.6/bin/trtexec /opt/tensorrt/bin/
python3 -m pip install TensorRT-10.0.1.6/python/tensorrt-10.0.1-cp310-none-linux_x86_64.whl
rm -rf /usr/local/tensorrt
cp -r TensorRT-10.0.1.6 /usr/local/tensorrt# 下载TensorRT-LLM源码git clone https://github.com/NVIDIA/TensorRT-LLM.git
git submodule update --init --recursive --force# 手动安装一些依赖(直接install requirement.txt容易被mpi4py卡主)pip config set global.index-url https://mirrors.cloud.tencent.com/pypi/simple
python3 -m pip uninstall cugraph torch torch-tensorrt tensorrt transformer-engine flash-attn torchvision torchtext torchdata torchaudio dask-cuda cugraph-service-server cuml -y
python3 -m pip install cmake
python3 -m pip install --no-cache-dir --extra-index-url https://pypi.nvidia.com mpi4py
python3 -m pip install -r requirements.txt
python3 -m pip install -r requirements-dev.txt
python3 -m pip install flash-attn --no-build-isolation# 源码编译 全架构 比较耗时python3 ./scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt -j 48# 编译指定架构Ada: sm89, Ampere sm80: A30,A800,A100, sm86: 4090,4080,...python3 ./scripts/build_wheel.py --clean --cuda_architectures "89-real" --trt_root /usr/local/tensorrt

默认编译是不带benchmark功能。需要在源码编译TensorRT-LLM的时候，指定benchmark选项：

python3 ./scripts/build_wheel.py --clean --benchmarks --trt_root /usr/local/tensorrt -j 48# 如果是在conda环境，可能还需要指定对应的python3进行编译python3 ./scripts/build_wheel.py --clean --benchmarks --trt_root /usr/local/tensorrt -j 48 \     --extra-cmake-vars PYTHON_EXECUTABLE=/usr/bin/python3

跑benchmark之前，还要准备好模型engine文件，这里以Internlm2-chat-20b为例（需要main分支最新的代码）。

# InternLM2 TRT-LLM性能测试 FP16 TP2cd TensorRT-LLM/examples/internlm2

python convert_checkpoint.py \        --model_dir ./internlm2-chat-20b/ \        --dtype float16 \        --output_dir ./tmp_fp16_tp2 \        --tp_size 2trtllm-build \        --checkpoint_dir ./tmp_fp16_tp2 \        --output_dir ./engine/internlm2-chat-20b/fp16/2-gpu/ \        --max_batch_size 8 \        --max_input_len 3072 \        --max_output_len 1024 \        --max_num_tokens 16384 \        --max_beam_width 1 \        --workers 2 \        --gemm_plugin float16 \        --gpt_attention_plugin float16 \        --remove_input_padding enable \        --paged_kv_cache enable \        --context_fmha enable \        --logits_dtype float32 \        --context_fmha_fp32_acc enable




    
 \        --use_paged_context_fmha enable

gptSessionBenchmark

gptSessionBenchmark跑的是static batching，可以用来快速看一下模型的性能，做一下性能收益评估之类的。比如：

export TRTLLM_BIN_DIR=/workspace/TensorRT-LLM/cpp/build
mpirun --allow-run-as-root -n 2 $TRTLLM_BIN_DIR/benchmarks/gptSessionBenchmark \    --engine_dir ./engine/internlm2-chat-20b/fp16/2-gpu/ \    --batch_size "1" --warm_up 2 --num_runs 20 \    --input_output_len "3072,128"

又比如，你想快速预估一下prompt为2048 token时的首Token的耗时：

export TRTLLM_BIN_DIR=/workspace/TensorRT-LLM/cpp/build
mpirun --allow-run-as-root -n 2 $TRTLLM_BIN_DIR/benchmarks/gptSessionBenchmark \    --engine_dir ./engine/internlm2-chat-20b/fp16/2-gpu/ \    --batch_size "1" --warm_up 2 --num_runs 20 \    --input_output_len "2048,1"

gptManagerBenchmark

gptSessionBenchmark跑的是static batching，没有In-Flight Batching，如果需要跑IFB，则可以使用gptManagerBenchmark。gptManagerBenchmark是模拟请求发送和调度的，因此在使用前，需要先准备好数据集。TensorRT-LLM提供了一个prepare_dataset.py脚本可以帮助用户快速准备测试数据集（只关注性能）。

在性能评估中，我们经常需要固定输入和输出的长度，那么可以这样：

# 准备定长2048 token输入，定长128输出cd TensorRT-LLM/benchmarks/cpp
python3 prepare_dataset.py \  --output ./tokens-fixed-lengths.json \  --tokenizer PATH-TO/internlm2-chat-20b/ \   token-norm-dist \   --num-requests 512 \   --input-mean 2048 --input-stdev 0 \   --output-mean 128 --output-stdev 0

产出的数据会保存在json文件中，大概长这样：

{"metadata": {"workload_type": "token-norm-dist", "input_mean": 2048, "input_stdev": 0, "output_mean": 128, "output_stdev": 0, "num_requests": 512, "tokenize_vocabsize": 92544, "max_input_len": 2048, "max_output_len": 128, "workload_name": "workload_type:token-norm-dist__input_mean:2048__input_stdev:0__output_mean:128__output_stdev:0__num_requests:512__tokenize_vocabsize:92544__max_input_len:2048__max_output_len:128"}, "samples": [{"input_len": 2048, "input_ids": [3452, 88226, 47775, 35731, 52781, 12529, 86990, 88852, 86033, 14974, 3960, 68403, 42664, 63581, 85238, 2417, 81796, 16710, 71676, 43421, 56550, 35152, 16223, 44050, 35639, 19196, 89462, 13952, 34254, 64423, 24180, 63111, 87473, 13318, 32424, 65016, 1218, 51691, 79986, 10

运行gptManagerBenchmark，并指定调度方式为IFB，request_rate表示需要模拟的QPS：

export CUDA_VISIBLE_DEVICES=0,1export TRTLLM_BIN_DIR=/workspace/TensorRT-LLM/cpp/build
mpirun --allow-run-as-root -n 2 $TRTLLM_BIN_DIR/benchmarks/gptManagerBenchmark \    --engine_dir PATH-TO/engine/internlm2-chat-20b/fp16/2-gpu/ \    --type IFB --request_rate 10 --max_num_samples 128 --warm_up 2 \    --enable_kv_cache_reuse false --dataset ./tokens-fixed-lengths.json

需要注意的是max_batch_size这个参数，可能会对gptManagerBenchmark的吞吐有影响，太小的max_batch_size会限制TensorRT-LLM能够并行处理的请求数，太大的又容易导致OOM（哭，能不能搞成自动算好的，伸手党拒绝手工autotune...）。比如对于Intermlm2-chat-20b这个模型，在max_batch_size=8时你会发现高并发情况，trtllm不如lmdeploy，但是将max_batch_size改成>=16，结论可能就反过来了。这也是目前使用TensorRT-LLM比较难受的地方，总是要反复尝试才能得到一个最优配置。

模拟static batching一次发出所有请求

gptManagerBenchmark可以模拟static batching一次发出所有请求，但是需要注意的是，这只是模拟static batching一次发出所有请求的方式，但是内部应该还是走的IFB对到达的请求进行调度。

# 模拟static batching一次性发出所有请求(实际内部会按照IFB调度到达的请求)export CUDA_VISIBLE_DEVICES=0,1export TRTLLM_BIN_DIR=/workspace/TensorRT-LLM/cpp/build
mpirun --allow-run-as-root -n 2 $TRTLLM_BIN_DIR/benchmarks/gptManagerBenchmark \    --engine_dir PATH-TO/engine/internlm2-chat-20b/fp16/2-gpu/ \    --type IFB --request_rate -1 static_emulated_batch_size 16 \    --static_emulated_timeout 100 --max_num_samples 16 \    --enable_kv_cache_reuse false \    --dataset ./tokens-fixed-lengths.json

static_emulated_batch_size表示一次发送要组batch的大小，示例这里是16；static_emulated_timeout表示等待组batch的时长，单位是ms，如果超过这个时长没有组满batch，也会直接送给推理引擎。

W8A18/W4A16性能测试

如果要测试weight only int8/int4量化，目前还是比较方便的，这个要点赞一下。通过转换脚本直接转换，并重新编译engine文件即可。TensorRT-LLM中weight only的实现原理和FasterTransformers中的应该是一样的，推荐看我之前写的几篇文章：

DefTruth：[LLM推理优化] WINT8/4-(00): 通俗易懂讲解-快速反量化算法 https://zhuanlan.zhihu.com/p/657072856

DefTruth：[LLM推理优化] WINT8/4-(01): PRMT指令详解及FasterTransformer源码解析 https://zhuanlan.zhihu.com/p/657070837

DefTruth：[LLM推理优化] WINT8/4-(02): 快速反量化之INT8转BF16 https://zhuanlan.zhihu.com/p/657073159

DefTruth：[LLM推理优化] WINT8/4-(03): LOP3指令详解及INT4转FP16/BF16分析 https://zhuanlan.zhihu.com/p/657073857

这里提供两个TensorRT-LLM使用W4A16和W8A16的参考脚本：

# W8A16 TP2python3 convert_checkpoint.py \        --model_dir $HF_MODELS/internlm2-chat-20b/ \        --dtype float16 \        --output_dir ./tmp_int8_tp2 \        --use_weight_only \        --weight_only_precision int8 \        --tp_size 2trtllm-build \        --checkpoint_dir ./tmp_int8_tp2 \        --output_dir $HF_MODELS/engine/internlm2-chat-20b/int8/2-gpu/ \        --max_batch_size 16 \        --max_input_len 3072 \        --max_output_len 1024 \        --max_num_tokens 16384 \        --max_beam_width 1 \        --workers 2 \        --gemm_plugin float16 \        --gpt_attention_plugin float16 \        --remove_input_padding enable \        --paged_kv_cache enable \        --context_fmha enable \        --logits_dtype float32 \        --context_fmha_fp32_acc enable \        --use_paged_context_fmha enable \        --weight_only_precision int8# W4A16 TP2python3 convert_checkpoint.py \        --model_dir $HF_MODELS/internlm2-chat-20b/ \        --dtype float16 \        --output_dir ./tmp_int4_tp2 \        --use_weight_only \        --weight_only_precision int4 \        --tp_size 2trtllm-build \        --checkpoint_dir ./tmp_int4_tp2 \        --output_dir $HF_MODELS/engine/internlm2-chat-20b/int4/2-gpu/ \        --max_batch_size 16 \        --max_input_len 3072 \        --max_output_len 1024 \        --max_num_tokens 16384 \        --max_beam_width 1 \        --workers 2 \        --gemm_plugin float16 \        --gpt_attention_plugin float16 \        --remove_input_padding enable \        --paged_kv_cache enable \        --context_fmha enable \        --logits_dtype float32 \        --context_fmha_fp32_acc enable \        --use_paged_context_fmha enable \        --weight_only_precision int4

跑gptManagerBenchmark替换为新编译的int8/int4的engine即可。这对于快速预估int8/int4的性能收益还是很方便的。点赞。

0x0d FP8/SQ/AWQ量化校准使用

TensorRT-LLM的量化工具在examples/quantization目录下，需要安装额外的依赖：

cd TensorRT/examples/quantization
python3 -m pip install -r requirements.txt

FP8量化-W8A8

如果只是为了快速评估FP8量化后的性能，可以不用额外准备数据。quantize.py脚本会使用默认的数据集cnn_dailymail进行FP8校准。比如，量化Qwen1.5-7B-Chat模型：

# FP8量化-W8A8python3 quantize.py \    --model_dir $HF_MODELS/Qwen1.5-7B-Chat/ \    --dtype float16 \    --qformat fp8 \    --output_dir ./tmp_fp8_tp2 \    --tp_size 2 \    --calib_size 256 \    --calib_max_seq_length 4096# build enginetrtllm-build \        --checkpoint_dir ./tmp_fp8_tp2 \        --output_dir $HF_MODELS/engine/Qwen1.5-7B-Chat/fp8/2-gpu/ \        --max_batch_size 16 \        --max_input_len 3072 \        --max_output_len 1024 \        --max_num_tokens 16384 \        --max_beam_width 1