TensorRT-LLM：大语言模型的高效推理部署

大语言模型（LLM）的推理部署是 AI 基础设施的核心挑战之一。随着模型参数量从数十亿增长到数千亿，如何在保证低延迟和高吞吐的前提下高效运行这些模型，成为生产环境的必答题。NVIDIA 推出的 TensorRT-LLM 正是为了解决这一痛点而诞生的专用推理框架。本文将系统性地介绍 TensorRT-LLM 的核心机制、优化手段和落地实践。

1. TensorRT-LLM 是什么？

TensorRT-LLM 是 NVIDIA 基于 TensorRT 构建的大语言模型推理框架，于 2023 年底开源（Apache 2.0 协议）。它继承了 TensorRT 在 GPU 推理优化上的深厚积累，同时针对 LLM 的生成式特点引入了多项专用优化。

核心定位：

• 面向 NVIDIA GPU（Ampere、Hopper 及更新架构）的 LLM 推理引擎
• 支持主流模型：GPT、LLaMA/LLaMA-2、ChatGLM、Baichuan、Falcon、Mistral 等
• 提供 Python 高层 API 定义模型，C++ 运行时执行推理
• 与 Triton Inference Server 深度集成，支持生产级部署

相比通用推理框架，TensorRT-LLM 的优势在于对 NVIDIA 硬件的极致利用——通过 Tensor Core 的 FP8/INT8 加速、定制的 CUDA Kernel 和显存管理策略，在 A100/H100 上往往能达到数倍于其他框架的吞吐。

2. 核心优化技术

2.1 In-flight Batching（连续批处理）

传统推理采用静态批处理（Static Batching）：一批请求必须全部完成生成后，才能处理下一批。这导致 GPU 在部分短请求完成后出现空闲，产生严重的气泡效应。

In-flight Batching 允许新请求在任何时候加入当前正在执行的批次，无需等待已有批次全部完成。其工作原理如下：

• 每个时序 step 中，调度器从等待队列中动态选择可加入的请求
• 已完成生成的请求被移除，其占用的 KV Cache 被回收
• 新请求在下一个 step 立即参与计算，无需等待

这种机制显著提升了 GPU 利用率，尤其在请求长度差异较大的场景中，吞吐量可提升 2-5 倍。

2.2 PagedAttention for KV Cache

PagedAttention 最初由 vLLM 提出，TensorRT-LLM 也将其纳入核心设计。其核心思想是将 KV Cache 分页管理：

• 将 KV Cache 划分为固定大小的块（Block），类似操作系统的页表
• 每个请求的 KV Cache 不必连续存储，而是通过块表（Block Table）间接寻址
• 请求完成后，其占用的块立即释放回全局池，供其他请求复用

这有效解决了传统分配方式中因预分配最大长度导致的显存浪费问题。在 TensorRT-LLM 中，PagedAttention 与 In-flight Batching 协同工作，使得显存利用率接近理论上限。

2.3 量化与低精度推理

TensorRT-LLM 支持多种量化策略，充分利用 NVIDIA Tensor Core 的硬件加速能力：

FP8 量化 是 Hopper 架构的核心亮点。H100/H200 的 Tensor Core 原生支持 FP8 计算，在多数 LLM 任务上几乎不损失精度的情况下，可将推理吞吐提升 1.5-2 倍。

SmoothQuant 通过将激活的离群值"平滑"迁移到权重，缓解了激活量化难度大于权重量化的问题。TensorRT-LLM 内置了 SmoothQuant 的校准与转换流程，用户只需在构建引擎时指定即可启用。

2.4 CUDA Graph

LLM 推理中，单步计算量虽小，但 Kernel 启动开销（CPU launch overhead）在短序列场景下不可忽视。TensorRT-LLM 支持 CUDA Graph，将重复的推理步骤捕获为静态图，消除 Host 侧的 Kernel 调度开销，在低 batch 场景下可降低 10-20% 的延迟。

3. 张量并行（TP）与流水线并行（PP）

当单 GPU 显存无法容纳整个模型时，需要通过分布式并行策略进行多卡扩展。TensorRT-LLM 支持两种主要并行方式：

3.1 张量并行（Tensor Parallelism, TP）

TP 将单个层的计算切分到同一节点内的多张 GPU。以 Attention 层为例：

• 多头注意力（MHA）的 Q/K/V 投影矩阵按头数切分到各卡
• 各卡独立计算局部注意力，通过 NCCL AllReduce 聚合结果
• FFN 层的两个线性变换也采用类似切分策略

TP 的特点：

• 通信量较大（每 layer 一次 AllReduce），依赖高带宽互联（NVLink）
• 仅适合同一节点内的 GPU（通常 2/4/8 卡）
• 对延迟敏感型任务更友好

3.2 流水线并行（Pipeline Parallelism, PP）

PP 将模型的不同层切分到不同 GPU 或不同节点。例如 80 层的 LLaMA-70B 可以按 10 层一组，分配到 8 张卡上。

PP 的特点：

• 通信量小（仅传递层间激活值），适合跨节点部署
• 存在流水线气泡（Bubble），吞吐量略低于 TP
• 扩展性更好，可以跨机扩展

3.3 混合并行策略

实际生产中常采用 TP + PP 的组合：

# 推荐配置示例：16 张 H100 部署 70B 模型
# 每节点 8 卡，共 2 节点
# TP=8（节点内），PP=2（跨节点）

选择策略的决策树：

1. 单节点可放下模型：优先 TP 最大化吞吐
2. 需要跨节点：引入 PP，TP 设为每节点卡数
3. 超长上下文：TP 不宜过大，避免 AllReduce 通信瓶颈

4. TensorRT-LLM 架构解析

TensorRT-LLM 采用分层架构：

+---------------------------+
| Python High-Level API     |  <- 模型定义、量化配置、引擎构建
| (Model Definition)        |
+---------------------------+
| TensorRT Builder          |  <- 图优化、Kernel 融合、引擎序列化
| (Optimization + Plan)     |
+---------------------------+
| C++ Runtime               |  <- 推理执行、KV Cache 管理、调度
| (Executor)                |
+---------------------------+
| CUDA Kernels / Plugins    |  <- 自定义注意力、量化算子
+---------------------------+

4.1 Python API 定义模型

用户通过 Python API 描述模型结构，而非直接操作 ONNX 或原始 TensorRT API。以 GPT 为例：

from tensorrt_llm.layers import Attention, LayerNorm, MLP
from tensorrt_llm.models import GPTLMHeadModel

# 使用预置模型类
model = GPTLMHeadModel(config)
# 或自定义 layer 组装

TensorRT-LLM 内置了主流模型的预置实现，用户通常只需加载 Hugging Face 权重并调用转换脚本。

4.2 Builder 配置

构建 TensorRT Engine 时，需配置多项优化参数：

from tensorrt_llm import Builder

builder = Builder()
network = builder.create_network()

# 优化配置示例
builder_config = builder.create_builder_config(
    name="llama_70b",
    precision="fp16",       # 或 fp8 / int8
    timing_cache="cache.bin",
    tensor_parallel=8,
    pipeline_parallel=2,
    max_batch_size=32,
    max_input_len=4096,
    max_output_len=2048,
)

关键配置项：

• Optimization Profile：定义动态 shape 范围（batch、input_len、output_len）
• Plugin System：对自定义算子（如特定 Attention 变体）使用 Plugin 注册
• Precision：逐层精度配置，对敏感层保留 FP16

4.3 Plugin 系统

当内置算子无法满足需求时，可通过 Plugin 机制注入自定义 CUDA Kernel。Plugin 以动态库形式加载，与主运行时解耦，便于集成特定硬件特性或实验性算法。

5. 构建第一个 TensorRT-LLM 引擎

以下以 LLaMA-2-7B 为例，演示完整的引擎构建流程。

5.1 环境准备

# 安装 TensorRT-LLM (推荐通过 pip 或 Docker)
pip install tensorrt_llm==0.11.0 --extra-index-url https://pypi.nvidia.com

# 拉取官方示例
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM/examples/llama

5.2 从 Hugging Face 转换权重

# 下载 Llama-2-7b-hf 权重并转换
python convert_checkpoint.py \
    --model_dir ./Llama-2-7b-hf \
    --output_dir ./tllm_checkpoint_1gpu_fp16 \
    --dtype float16

5.3 构建引擎

# FP16 引擎
trtllm-build \
    --checkpoint_dir ./tllm_checkpoint_1gpu_fp16 \
    --output_dir ./llama_7b_trt_engine_fp16 \
    --gemm_plugin float16 \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_output_len 512

5.4 量化引擎构建（SmoothQuant W8A8）

# 第一步：INT8 权重量化 + SmoothQuant 校准
python convert_checkpoint.py \
    --model_dir ./Llama-2-7b-hf \
    --output_dir ./tllm_checkpoint_sq \
    --dtype float16 \
    --use_weight_only \
    --weight_only_precision int8 \
    --use_smooth_quant

# 第二步：构建量化引擎
trtllm-build \
    --checkpoint_dir ./tllm_checkpoint_sq \
    --output_dir ./llama_7b_trt_engine_sq \
    --quant_ckpt_path ./tllm_checkpoint_sq \
    --gemm_plugin float16 \
    --use_custom_all_reduce enable \
    --max_batch_size 16 \
    --max_input_len 2048 \
    --max_output_len 512

5.5 推理验证

from tensorrt_llm.runtime import ModelRunnerCpp
from transformers import AutoTokenizer

engine_dir = "./llama_7b_trt_engine_fp16"
runner = ModelRunnerCpp.from_dir(engine_dir)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

input_text = "TensorRT-LLM 的优化技术包括："
input_ids = tokenizer.encode(input_text, return_tensors="pt")

outputs = runner.generate(
    batch_input_ids=input_ids,
    max_new_tokens=128,
    temperature=0.7,
    top_p=0.9,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

5.6 序列化引擎的部署

构建好的引擎（.engine 文件）可直接序列化到磁盘，运行时无需重新编译 DAG，实现秒级冷启动：

# 引擎文件通常位于输出目录的 rank*.engine
ls ./llama_7b_trt_engine_fp16/
# config.json  rank0.engine  ...

6. Triton Inference Server 部署

生产环境中，TensorRT-LLM 通常以 Triton Inference Server 的 Backend 形式提供服务。

6.1 TensorRT-LLM Backend 架构

Triton 提供了 Python 和 C++ 两种 Backend 实现：

• Python Backend：开发灵活，适合快速迭代和自定义前后处理
• C++ Backend：延迟更低，适合对尾延迟（P99）敏感的场景

6.2 启用 In-flight Batching

在 Triton 的 config.pbtxt 中配置调度策略：

# config.pbtxt
parameters: {
  key: "enable_trt_overlap"
  value: { string_value: "true" }
}
parameters: {
  key: "batching_strategy"
  value: { string_value: "inflight_fused_batching" }
}
parameters: {
  key: "max_batch_size"
  value: { string_value: "64" }
}

6.3 模型 Ensemble

LLM 服务通常需要分词（Preprocessing）→ 推理 → 反分词（Postprocessing）三段流程。Triton 的 Ensemble 功能可以将这些模型串联为一个端点：

ensemble_scheduling {
  step [
    { model_name: "preprocessing" model_version: -1 },
    { model_name: "tensorrt_llm" model_version: -1 },
    { model_name: "postprocessing" model_version: -1 }
  ]
}

客户端只需调用一个端点，即可完成从文本输入到文本输出的完整链路。

7. 与替代方案的对比

7.1 TensorRT-LLM vs vLLM

选型建议：

• 若已基于 NVIDIA 基础设施且追求极限性能，选择 TensorRT-LLM
• 若需要多硬件支持或快速原型验证，选择 vLLM

7.2 TensorRT-LLM vs DeepSpeed-Inference

DeepSpeed-Inference 是微软推出的推理框架，同样支持 TP 和量化。相比之下：

• TensorRT-LLM：编译时优化更激进，推理延迟更低，但层数稍深
• DeepSpeed-Inference：与 PyTorch 模型定义结合更紧密，适合已有 DeepSpeed 训练代码的团队

7.3 性能参考数据

以 LLaMA-2-70B 为例，在 8xH100 SXM 上的对比（input 2048 / output 512 tokens）：

总结

TensorRT-LLM 是 NVIDIA 生态中 LLM 推理的旗舰方案。它通过 In-flight Batching、PagedAttention、多精度量化、CUDA Graph 等技术的组合，在 Hopper 架构上实现了业界领先的推理性能。虽然在易用性上略逊于 vLLM，且绑定 NVIDIA 硬件，但对于已经部署 A100/H100/H200 基础设施、追求极致吞吐和低延迟的生产环境而言，它是一个经过验证的优选方案。

对于希望落地的团队，建议从官方示例模型入手，逐步掌握引擎构建、量化校准和 Triton 部署的完整链路，再根据业务场景的 batch 分布和延迟要求，调整 TP/PP 的并行策略，以达到成本与性能的最佳平衡。