TensorRT-LLM 优化推理引擎
TensorRT-LLM 是 NVIDIA 推出的大语言模型推理优化引擎，基于 TensorRT 编译器将模型转换为高度优化的 GPU 执行图，支持多种先进优化技术。
核心特性
TensorRT-LLM 的主要优化技术包括：
- Kernel Fusion：将多个算子融合为单个 GPU Kernel，减少显存访问
- Continuous Batching：原生支持连续批处理调度
- In-Flight Batching：NVIDIA 专有的动态批处理方案
- FP8/INT8 量化：支持多种精度推理
- Tensor Parallel：原生张量并行支持
python
# TensorRT-LLM 模型编译
from tensorrt_llm import LLM, SamplingParams
# 编译并加载模型
llm = LLM(
model="meta-llama/Llama-2-7b-hf",
tensor_parallel_size=2,
dtype="float16",
max_batch_size=32,
max_input_len=2048,
max_output_len=512,
)
# 推理
params = SamplingParams(max_tokens=256, temperature=0.7)
outputs = llm.generate(["解释 Transformer 架构"], params)编译优化流程
TensorRT-LLM 的编译流程：
- 模型解析：将 HuggingFace 模型转换为 TRT-LLM 中间表示
- 图优化：算子融合、常量折叠、死代码消除
- Kernel 选择：为每个算子选择最优 GPU Kernel
- 引擎构建：生成序列化的执行引擎
- 运行时加载：反序列化引擎并执行推理
编译时间
首次编译可能需要 10-30 分钟，但生成的引擎可以序列化保存，后续加载仅需数秒。建议在 CI/CD 流程中预编译引擎。
性能对比
| 模型 | 精度 | A100 吞吐 (tok/s) | 延迟 P99 (ms) |
|---|---|---|---|
| Llama-2-7B | FP16 | 2800 | 45 |
| Llama-2-7B | INT8 | 4200 | 32 |
| Llama-2-70B | FP16 (4xTP) | 850 | 120 |
| Llama-2-70B | INT8 (4xTP) | 1300 | 85 |
硬件绑定
TensorRT-LLM 编译的引擎与 GPU 架构强绑定。在 A100 上编译的引擎无法在 H100 上运行，需要针对目标硬件重新编译。