TensorRT-LLM思考：从入门到生产级GPU推理优化

共计 4786 个字符，预计需要花费 12 分钟才能阅读完成。

---

一、引言：为什么选择 TensorRT-LLM？

随着大语言模型（如 Llama、Falcon、ChatGLM、Qwen 等）在企业中的广泛应用，推理延迟 、 吞吐量 和显存占用 成为关键瓶颈。NVIDIA 推出的 TensorRT-LLM 是专为 LLM 推理优化的开源库，具备以下优势：

• 极致性能：利用 TensorRT 引擎，结合 CUDA、cuBLAS、cuDNN 等底层库，实现低延迟、高吞吐推理。
• 支持主流模型：Llama、Llama2/3、Mistral、Falcon、Gemma、Qwen、Baichuan、ChatGLM 等。
• 动态批处理（Dynamic Batching） 和 连续批处理（Continuous Batching）：提升 GPU 利用率。
• 多 GPU / 多节点推理：支持张量并行（Tensor Parallelism）和流水线并行（Pipeline Parallelism）。
• 量化支持：INT4/INT8 量化，大幅降低显存占用。
• Python/C++ API：便于集成到生产系统。

---

二、环境准备

1. 硬件要求

• NVIDIA GPU（建议 A100/H100/L40S 或 RTX 4090，支持 Tensor Core）
• 驱动版本 ≥ 535
• CUDA ≥ 12.2

2. 软件依赖

# 安装 Docker（推荐方式）sudo apt-get install docker.io nvidia-docker2

# 拉取官方镜像（以 TensorRT-LLM 0.10 为例）docker pull nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3

或从源码构建（适用于定制需求）：

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
./docker/build.sh  # 构建开发镜像

3. 安装依赖库（容器内）

pip install transformers accelerate sentencepiece

---

三、模型转换：从 Hugging Face 到 TensorRT-LLM

TensorRT-LLM 不直接加载 .bin 或 .safetensors 文件，而是通过 构建 TensorRT 引擎 实现优化。

步骤 1：准备 Hugging Face 模型

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

步骤 2：使用 TensorRT-LLM 脚本转换

TensorRT-LLM 提供了 examples/ 目录下的转换脚本。以 Llama 为例：

cd examples/llama
python3 convert_checkpoint.py \
  --model_dir /path/to/llama-2-7b-hf \
  --output_dir ./trt_llm_checkpoint \
  --dtype float16

该脚本会生成一个包含权重和配置的中间格式（非 TensorRT 引擎）。

步骤 3：构建 TensorRT 引擎

trtllm-build \
  --checkpoint_dir ./trt_llm_checkpoint \
  --output_dir ./trt_llm_engine \
  --gpt_attention_plugin float16 \
  --gemm_plugin float16 \
  --max_batch_size 8 \
  --max_input_len 512 \
  --max_output_len 256 \
  --max_beam_width 1

关键参数说明：

• --gpt_attention_plugin：启用优化的注意力内核（必须启用以获得高性能）
• --max_batch_size：最大批处理大小（影响显存）
• --max_input_len / --max_output_len：输入 / 输出序列长度上限

构建完成后，./trt_llm_engine 目录包含 .engine 文件，即最终可部署的推理引擎。

---

四、推理执行：Python API 示例

import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
runner = ModelRunner.from_dir(
    engine_dir="./trt_llm_engine",
    lora_dir=None,
    rank=tensorrt_llm.mpi_rank())

input_text = "Explain the theory of relativity in simple terms."
input_ids = tokenizer.encode(input_text, return_tensors="pt").to("cuda")

output_ids = runner.generate(
    input_ids,
    max_new_tokens=128,
    end_id=tokenizer.eos_token_id,
    pad_id=tokenizer.pad_token_id
)

output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

---

五、高级优化技巧

1. 启用连续批处理（Continuous Batching）

TensorRT-LLM 支持异步请求队列，自动合并不同长度的请求：

from tensorrt_llm.bindings import GptSession

session = GptSession(
    engine_path="./trt_llm_engine",
    max_batch_size=32,
    max_beam_width=1
)
# 支持异步 enqueue_request()

适用于高并发在线服务（如 API 服务器）。

2. 量化（INT4/INT8）

大幅降低显存，牺牲少量精度：

# 转换时启用 AWQ 或 SmoothQuant
python3 convert_checkpoint.py \
  --model_dir /path/to/model \
  --output_dir ./quant_checkpoint \
  --dtype float16 \
  --qformat int4_awq \
  --awq_path /path/to/llama2-7b-awq.pt

注意：需提前使用 autoawq 或 smoothquant 工具生成量化参数。

3. 多 GPU 张量并行

mpirun -n 4 --allow-run-as-root \
  trtllm-build \
    --checkpoint_dir ./trt_llm_checkpoint \
    --output_dir ./trt_llm_engine_tp4 \
    --gpt_attention_plugin float16 \
    --max_batch_size 16 \
    --tp_size 4  # Tensor Parallelism = 4

推理时自动通过 MPI 分发。

4. KV Cache 优化

启用 paged_kv_cache 减少显存碎片：

trtllm-build ... --paged_kv_cache --tokens_per_block 64

---

六、生产部署：集成 Triton Inference Server

TensorRT-LLM 官方推荐与 NVIDIA Triton 集成。

1. 准备模型仓库

model_repo/
└── llama-2-7b-trtllm/
    ├── config.pbtxt
    └── 1/
        └── model.engine

config.pbtxt 示例：

name: "llama-2-7b-trtllm"
platform: "tensorrt_llm"
max_batch_size: 8
input [
  {
    name: "input_ids"
    data_type: TYPE_INT32
    dims: [-1]
  }
]
output [
  {
    name: "output_ids"
    data_type: TYPE_INT32
    dims: [-1]
  }
]
parameters: {
  key: "gpt_model_type"
  value: {string_value: "LLaMA"}
}

2. 启动 Triton Server

docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v $(pwd)/model_repo:/models \
  nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3 \
  tritonserver --model-repository=/models

3. 客户端调用（HTTP/gRPC）

import tritonclient.http as httpclient

client = httpclient.InferenceServerClient(url="localhost:8000")
input_ids = np.array([[1, 2, 3, ...]], dtype=np.int32)
inputs = [httpclient.InferInput("input_ids", input_ids.shape, "INT32")]
inputs[0].set_data_from_numpy(input_ids)
result = client.infer(model_name="llama-2-7b-trtllm", inputs=inputs)
output = result.as_numpy("output_ids")

---

七、性能基准与调优建议

优化项	效果
启用 gpt_attention_plugin	2–3x 加速
使用 FP16 + Tensor Core	显存减半，速度翻倍
INT4 量化	显存减少 75%，吞吐提升 2x
连续批处理	GPU 利用率提升 3–5x
Paged KV Cache	支持更长上下文（32K+）

调优 checklist：

• 使用最新版 TensorRT-LLM（每月更新）
• 启用所有插件（attention, gemm, paged kv）
• 根据实际负载设置 max_batch_size 和 max_input_len
• 监控 GPU 利用率（nvidia-smi dmon）
• 使用 trtexec 分析引擎性能

---

八、常见问题（FAQ）

Q1：是否支持自定义模型？
A：支持！需实现 tensorrt_llm/models/ 下的模型类，并注册到构建流程。

Q2：如何支持 LoRA 微调？
A：TensorRT-LLM 原生支持 LoRA，转换时指定 --lora_dir，推理时动态加载。

Q3：能否与 vLLM 对比？
A：TensorRT-LLM 在 NVIDIA 硬件上通常更快（尤其 H100），但 vLLM 更易用、社区活跃。建议根据团队技术栈选择。

---

九、结语

TensorRT-LLM 是当前 NVIDIA GPU 上 LLM 推理的黄金标准。虽然学习曲线较陡，但一旦掌握，可将推理性能推向极致。建议从官方 examples/ 入手，逐步尝试量化、多卡、Triton 集成等高级功能。

📚 延伸阅读：

• TensorRT-LLM GitHub
• NVIDIA 官方文档
• 《Optimizing LLM Inference on GPUs》NVIDIA 技术白皮书