TensorRT-LLM 是 NVIDIA 用于做 LLM 的可扩展推理方案。该方案是基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，然后利用 NCCL 完成设备之间的通讯。考虑到技术的发展和需求的差异，开发者还可以定制算子来满足定制需求，比如基于 cutlass 开发定制 GEMM。TensorRT-LLM 是一款致力于提供高性能并不断完善其实用性的 NVIDIA 官方推理方案。

TensorRT-LLM 已经在 GitHub 上开源，主要分为两个分支，即 Release branch 和 Dev branch。其中 Release branch 每个月更新一次，而在 Dev branch 中则会较为频繁地更新来自官方或社区中的功能，方便开发者体验、评估最新功能。

重要特性：

•丰富的模型支持

•FMHA（fused multi-head attention）kernel 的实现

•量化技术，以更低精度的方式实现推理加速

• In-flight batching

使用流程:

首先需要获得一个预训练好的模型，然后利用 TensorRT-LLM 提供的 API 对模型计算图进行改写和重建，接着用 TensorRT 进行编译优化，然后保存为序列化的 engine 进行推理部署。

以 Llama 为例，首先安装 TensorRT-LLM，然后下载预训练模型，接着利用 TensorRT-LLM 对模型进行编译，最后进行推理。

对于模型推理的调试，TensorRT-LLM 的调试方式与 TensorRT 一致。由于深度学习编译器，即 TensorRT，提供的优化之一是 layer 融合。因此，如果要输出某层的结果，就需要将对应层标记为输出层，以防止被编译器优化掉，然后与 baseline 进行对比分析。同时，每标记一个新的输出层，都要重新编译 TensorRT 的 engine。

对于自定义的层，TensorRT-LLM 提供了许多 Pytorch-like 算子帮助用户实现功能而不必自己编写 kernel。如样例所示，利用 TensorRT-LLM 提供的 API 实现了 rms norm 的逻辑，TensorRT 会自动生成 GPU 上对应的执行代码。

如果用户有更高的性能需求或者 TensorRT-LLM 并未提供实现相应功能的 building blocks，此时需要用户自定义 kernel，并封装为 plugin 供 TensorRT-LLM 使用。示例代码是将 SmoothQuant 定制 GEMM 实现并封装成 plugin 后，供 TensorRT-LLM 调用的示例代码。

未来展望：

NVIDIA 将会持续致力于提升 TensorRT-LLM 的性能。同时通过开源，收集反馈和意见，提高它的易用性。另外，围绕易用性，会开发、开源更多应用工具，如 Model zone 或者量化工具等，完善与主流框架的兼容性，提供从训练到推理和部署端到端的解决方案。

责任编辑：kj005