TensorRT-LLM 实战：从模型导出到 FP8 推理加速

摘要: vLLM 适合"开箱即用、动态批处理"，TensorRT-LLM（TRT-LLM）适合"同一模型极致压榨"。本文走完整链路：HuggingFace 模型 → TRT-LLM 格式转换 → FP8 量化 → 引擎构建 → 推理服务化 → 基准对比，所有命令可复现。

0. 实验目标

• 把 HF 格式模型转成 TRT-LLM 引擎，跑通推理
• 用 FP8 量化把 70B 模型塞进单卡 80G（或显著提吞吐）
• 理解 trtllm-build 核心参数的语义（max-batch-size / max-seq-len / dtype）
• 与 vLLM 同卡同模型做基准对比，给出选型建议

1. 环境与格式转换

# 官方镜像（推荐，避免环境坑）
docker run --gpus all -it --shm-size=16g \
  nvcr.io/nvidia/tritonserver:24.07-trtllm-python-py3

# HF 模型 → TRT-LLM 检查点（拿 LLaMA-3-8B 举例）
python3 convert_checkpoint.py \
  --model_dir meta-llama/Llama-3-8B \
  --output_dir ./trt-ckpt/llama-8b \
  --dtype float16 \
  --tp_size 1

逐段拆解：

• convert_checkpoint.py：把 HF 权重重排成 TRT-LLM 需要的布局（权重、KV 缓存、MoE 路由等按 TRT 图结构预先分好）。转换一次、引擎可多次构建。
• --tp_size 1：张量并行数。8B 单卡够用设 1；70B 需要 --tp_size 2/4/8（多卡时同时决定 KV cache 与 attention 的切分）。
• 先确认 GPU 驱动支持（nvidia-smi），FP8 量化（下文）需要 H100 或 Ada 系架构（H100 最佳）。

2. FP8 量化（关键一步）

# 方案 A：HF 量化权重转换（如果模型已发布 FP8 版本）
python3 convert_checkpoint.py \
  --model_dir meta-llama/Llama-3-8B-Instruct-FP8 \
  --output_dir ./trt-ckpt/llama-8b-fp8 \
  --dtype float16 \
  --qformat fp8

# 方案 B：PTQ 后量化（无 FP8 权重时，用校准集量化）
python3 quantize.py \
  --model_dir meta-llama/Llama-3-8B \
  --dtype float16 \
  --qformat fp8 \
  --calib_dataset ./calib-data.jsonl \
  --output_dir ./trt-ckpt/llama-8b-fp8-ptq \
  --calib_size 512

原理讲解：

• FP8（E4M3 权重 + E4M3 激活 / E5M2 部分场景）：权重和 KV Cache 各减半显存，70B 从 ~140GB（FP16）降到 ~70GB，正好落进单卡 80G。
• PTQ 校准：用几百条代表性样本跑一遍，统计激活分布，给 FP8 找到最优 scale。校准集要贴近真实负载（别用与训练集重复的语料）。
• Per-tensor vs Per-channel：TRT-LLM FP8 默认逐层/逐通道量化，精度损失通常在 1% 以内（对生成任务几乎无感）；如果校准集偏差大，生成质量会明显退化——校准集质量 = 量化质量。

3. 构建引擎

trtllm-build \
  --checkpoint_dir ./trt-ckpt/llama-8b-fp8 \
  --output_dir ./engines/llama-8b-fp8 \
  --max_batch_size 64 \
  --max_seq_len 8192 \
  --max_num_tokens 8192 \
  --gemm_plugin auto \
  --kv_cache_dtype fp8 \
  --use_fused_mlp

逐段拆解：

• --max_batch_size 64 / --max_num_tokens 8192：引擎是按最大尺寸静态规划的。想改并发上限？必须重新构建引擎——这是 TRT-LLM 与 vLLM 的最大差异（vLLM 动态，TRT 静态）。
• --kv_cache_dtype fp8：KV Cache 也用 FP8，长上下文显存省一半。
• --gemm_plugin auto / --use_fused_mlp：把矩阵乘和 MLP 用 CUDA 内核融合。这些 plugin 开关是 TRT-LLM 提速的核心，默认开。

⚠️ 构建耗时：8B 引擎构建约 3-10 分钟，70B 可达 30-60 分钟。构建参数（尤其 max_seq_len）定了就别随便改，否则全量重来。

4. 启动推理服务

# 方案 A：trtllm-serve（单进程快速验证）
trtllm-serve --engine_dir ./engines/llama-8b-fp8 \
  --host 0.0.0.0 --port 8001 \
  --max_beam_width 1

# 方案 B：Triton + TRT-LLM 后端（生产推荐，带监控/批处理/动态并发）
docker run --gpus all --shm-size=16g --rm \
  -p 8001:8001 -v $(pwd)/engines:/engines \
  nvcr.io/nvidia/tritonserver:24.07-trtllm-python-py3 \
  tritonserver --model-repository=/engines/triton-repo

工程启示：

• trtllm-serve 暴露 OpenAI 兼容 API（/v1/chat/completions），快速验证用。
• 生产上 Triton + TRT-LLM 后端才是正路：动态批处理（dynamic batching）、Prometheus 指标、多模型管理、Grace 排队都在这里。
• 和 vLLM 一样：服务永不直接暴露公网，网关统一收口。

5. 基准对比：TRT-LLM vs vLLM（同卡同模型）

# 同样用 /metrics 或自写压测脚本，32 并发、128 token
# 结果记录格式：吞吐(token/s) | P50 首 token(ms) | P95 首 token(ms)

实验结果（FP8 · 单卡 80G · 8B 模型 · 示例量级）：

引擎	吞吐（token/s）	P50 首 token	备注
vLLM（FP16）	~3.1k	中	动态批处理、易用
TRT-LLM（FP16）	~3.6k	中	静态引擎、融合内核
TRT-LLM（FP8）	~4.4k	低	KV Cache 也 FP8，并发更高

数值为相对量级参考。结论不是"TRT 一定赢"：vLLM 在动态负载、快速迭代场景更合适；TRT-LLM 在"固定模型、固定规格、极致性能"场景有 10-40% 优势，代价是构建耗时与规格锁定。

6. 选型建议（可直接抄进方案）

场景	推荐	理由
快速上线 / 多模型切换 / 实验迭代	vLLM	动态批处理、无需重新构建
单一主力模型、规格固定、极致吞吐	TRT-LLM	融合内核 + FP8，性能上限更高
70B+ 单卡推理（显存吃紧）	TRT-LLM FP8	FP8 权重+KV 显存减半
大模型服务网关 / 多模型统一入口	Triton	多后端、动态批处理、监控完备

7. 工程启示汇总

1. TRT-LLM 是"静态规格"引擎：max_seq_len/batch 定了就不能动态改，变更要重新构建。
2. FP8 三步走：驱动支持（H100/Ada）→ 权重转换或 PTQ → 构建引擎开 kv_cache_dtype fp8。
3. 校准集质量 = 量化质量：用贴近真实负载的样本校准，别用训练语料。
4. plugin 开关是提速核心：gemm_plugin/use_fused_mlp 默认开，别关。
5. 生产走 Triton 后端：动态批处理、监控、多模型管理一站搞定。
6. 与 vLLM 互补而非互斥：快速迭代用 vLLM，主力模型压榨用 TRT-LLM。
7. 构建参数版本化管理：引擎构建参数写进 CI 配置，换参即换引擎版本，可回溯。

8. 延伸阅读

• TensorRT-LLM 官方文档
• Triton Inference Server 文档
• 本知识库：12. 推理引擎优化 章节 · vLLM 生产化部署