NVIDIA TensorRT-LLM

TensorRT-LLM 是 NVIDIA 推出的开源库，用于在 NVIDIA GPU 上加速和优化大语言模型（large language model）的推理（inference），也就是让训练好的模型真正回答问题的那一步。NVIDIA 于 2023 年 10 月 19 日宣布公开发布，用户可以免费从 GitHub 上的 NVIDIA/TensorRT-LLM 仓库获取，也可以作为 NVIDIA NeMo 的一部分使用。^[1]

由来

GitHub 仓库本身创建于 2023 年 8 月 16 日，比公开发布早约两个月。^[2] NVIDIA 表示，这次发布建立在过去两年与大语言模型公司的推理合作之上，合作方包括 Anyscale、Cohere、Meta、Mistral AI、Perplexity AI 和 Together.ai 等。^[1] GitHub 记录的该仓库许可证为 NOASSERTION，而非标准的 SPDX 许可证标识，也就是说它没有挂上常见的简短许可证标签。^[2]

怎么工作

生成文本并不是一次大计算。模型是一个词元（token）一个词元地输出答案的，每生成一个新词元，都要对之前的所有内容做注意力计算，因此瓶颈往往在显存读写，以及请求如何分组、缓存如何管理，而不只是算术量。TensorRT-LLM 把几类针对这些瓶颈的技术打包在一起：算子融合（kernel fusion），把多个小的 GPU 运算合并成一个；量化（quantization），用更低的精度存放权重和激活值；连续在途批处理（continuous in-flight batching），让请求随时加入或离开正在运行的批次，而不必等整批跑完；分页注意力（paged attention）；以及多 GPU、多节点推理。在 Hopper 架构的 GPU 上，它通过 Hopper transformer engine 支持 FP8 精度。^[1]

开发者通过 Python API 定义模型，库再去封装 NVIDIA 的深度学习编译器 TensorRT，并提供针对 FlashAttention 和掩码多头注意力（masked multi-head attention）的优化算子。所谓掩码注意力，是指每个位置只能看到它之前的位置，这正是预测下一个词元所需要的。^[1]

发布时支持的模型与硬件

发布之初，TensorRT-LLM 支持 Llama 1 和 Llama 2、ChatGLM、Falcon、MPT、Baichuan 与 Starcoder，覆盖 Ampere、Ada Lovelace 和 Hopper 三种 GPU 架构；另有一个面向 RTX 显卡的 Windows 测试版。^[1] 这份名单横跨多家机构的开放权重模型，从 Meta 的 Llama，到中文的 ChatGLM 和 Baichuan，再到面向代码的 Starcoder。

在 NVIDIA 技术栈中的位置

TensorRT-LLM 是 NVIDIA NeMo 中大语言模型推理的优化底座；NeMo 提供的容器里同时包含 TensorRT-LLM 和 NVIDIA Triton，用于生成式 AI 的部署。^[1] 在这样的分工里，TensorRT-LLM 负责针对具体 GPU 编译和调优模型的底层工作，NeMo 提供外围框架，Triton 则是接收请求的服务层。

对已经在 NVIDIA 硬件上运行的一方来说，它的实际吸引力在于把优化工作打包成了一个库：批处理、精度、注意力缓存的内存管理、跨多块 GPU 的切分，不必每个项目自己重写一遍。代价是它绑定 NVIDIA GPU，而且发布时的仓库元数据没有标明标准许可证。^[2]