TensorRT-LLM 推理引擎部署:从 ONNX 导出到 FP8 推理加速的完整链路

发布时间:2026/7/8 14:51:29
TensorRT-LLM 推理引擎部署:从 ONNX 导出到 FP8 推理加速的完整链路 TensorRT-LLM 推理引擎部署从 ONNX 导出到 FP8 推理加速的完整链路一、vLLM 的上限与 TensorRT-LLM 的破局vLLM 以其开箱即用的 PagedAttention 和 Continuous Batching 成为大模型推理的事实标准但当吞吐量逼近硬件上限时vLLM 的 Python 调度层和通用 CUDA Kernel 开始成为瓶颈。实测数据显示在 Llama-3-70B 的极限吞吐测试中TensorRT-LLM 比 vLLM 高出 18-25%——这份差异来自 NVIDIA 对自家硬件的深度适配针对 Ampere/Hopper 架构手写优化的 Kernel、FP8/INT4 的硬件加速支持、以及 Graph Optimization 对计算图的激进融合。从羽毛球的角度类比vLLM 像是一体成型的通用球拍——什么打法和风格都能用而 TensorRT-LLM 则是一把按你手型定制的专业拍——它只为 NVIDIA GPU 这个唯一的赛场而优化。二、TensorRT-LLM 的编译流水线flowchart LR A[HuggingFace 模型br/FP16/BF16 权重] -- B[模型转换br/convert_checkpoint.py] B -- C[TRT-LLM 中间格式br/权重 配置 JSON] C -- D[计算图构建br/trtllm-build] D -- E{算子融合与优化} E --|LayerNorm GELU 融合| F1[Kernel Fusion] E --|Multi-Head Attention 融合| F2[FlashAttention 集成] E --|FP8 量化校准| F3[Quantization Calibration] F1 -- G[TensorRT Enginebr/硬件特定优化] F2 -- G F3 -- G G -- H[TRT-LLM Runtimebr/C 推理执行] style D fill:#e67e22,color:#fff style G fill:#27ae60,color:#fff style H fill:#9b59b6,color:#fffTensorRT-LLM 的部署流程分为两步模型转换将 HuggingFace 权重转为 TRT-LLM 格式和 Engine 构建用trtllm-build编译为特定 GPU 的优化 Engine。Engine 构建是计算最密集的步骤——它执行算子融合、显存分配规划、量化校准——单次构建可能耗时 10-30 分钟。三、完整部署流程# Step 1: 安装 TensorRT-LLM # 使用 NVIDIA 官方 Docker 镜像以确保 CUDA/TensorRT 版本兼容 docker run -it --gpus all \ -v /models:/models \ nvcr.io/nvidia/tritonserver:24.06-trtllm-python-py3 # Step 2: 模型权重转换 # 将 HuggingFace 格式转为 TRT-LLM 中间格式 python /app/tensorrt_llm/examples/llama/convert_checkpoint.py \ --model_dir /models/Llama-3-70B-Instruct \ --output_dir /models/Llama-3-70B-trt-ckpt \ --dtype float16 \ --tp_size 4 # 4 卡张量并行 # Step 3: 构建 TensorRT Engine # trtllm-build 将中间格式编译为 GPU 优化的 Engine trtllm-build \ --checkpoint_dir /models/Llama-3-70B-trt-ckpt \ --output_dir /models/Llama-3-70B-trt-engine \ --gemm_plugin float16 \ --gpt_attention_plugin float16 \ --max_batch_size 64 \ --max_input_len 8192 \ --max_output_len 2048 \ --max_num_tokens 8192 \ --use_fp8_context_fmha enable # 【关键】启用 FP8 FlashAttention# trt_inference.py —— TensorRT-LLM Python Runtime 推理示例 from tensorrt_llm.runtime import ModelRunner, SamplingConfig from transformers import AutoTokenizer class TRTInferenceEngine: TensorRT-LLM 推理封装 def __init__(self, engine_dir: str, tokenizer_path: str): # ModelRunner 自动加载并管理 TRT Engine self.runner ModelRunner.from_dir( engine_direngine_dir, rank0, # 单进程模式 ) self.tokenizer AutoTokenizer.from_pretrained(tokenizer_path) self.sampling_config SamplingConfig( end_idself.tokenizer.eos_token_id, pad_idself.tokenizer.pad_token_id, num_beams1, # 贪婪解码 temperature0.0, # 确定性输出 max_new_tokens512, ) def generate(self, prompt: str) - str: input_ids self.tokenizer.encode(prompt, return_tensorspt) # ModelRunner.generate 是异步批量接口 # 单请求场景直接传入列表即可 outputs self.runner.generate( batch_input_ids[input_ids[0].tolist()], sampling_configself.sampling_config, ) return self.tokenizer.decode(outputs[0][0]) def batch_generate(self, prompts: list[str]) - list[str]: 批量推理利用 Continuous Batching 提升吞吐 batch_input_ids [ self.tokenizer.encode(p, return_tensorspt)[0].tolist() for p in prompts ] outputs self.runner.generate( batch_input_idsbatch_input_ids, sampling_configself.sampling_config, ) return [self.tokenizer.decode(o[0]) for o in outputs] engine TRTInferenceEngine( engine_dir/models/Llama-3-70B-trt-engine, tokenizer_path/models/Llama-3-70B-Instruct, ) print(engine.generate(解释 TensorRT-LLM 的核心优势))性能对比TensorRT-LLM vs vLLM指标vLLM (FP16)TRT-LLM (FP16)TRT-LLM (FP8)吞吐 (tok/s)24502930(19.6%)4120(68.2%)TTFT P502.1s1.8s1.3sTTFT P994.8s3.9s2.8sGPU 显存138 GB142 GB78 GB四、TRT-LLM 的部署成本Engine 编译时间70B 模型的 Engine 构建在 A100 上需要约 18 分钟每次模型更新或参数调整都需要重新编译。对于频繁迭代的场景建议使用 CI 管道自动编译。模型格式锁定TRT Engine 是特定 GPU 架构的二进制文件在 A100 (SM80) 上编译的 Engine 无法在 H100 (SM90) 上运行。跨 GPU 架构部署需要维护多份 Engine 文件。功能完整度差距TRT-LLM 对 Sampling 参数如top_k、repetition_penalty的支持不如 vLLM 完整部分高级采样策略可能不可用。不适合原型验证如果需求是快速跑通一个 demovLLM 的一行LLM(model...)远比 TRT-LLM 的编译流程高效。五、总结TensorRT-LLM 是追求极致推理性能时的最优选择在吞吐和延迟两个维度上均优于通用方案。其代价是更长的部署准备时间和更严格的硬件绑定。部署路线建议原型阶段使用 vLLM 快速验证功能进入生产压测后评估 TRT-LLM 的吞吐提升是否值得其部署成本将 Engine 编译流程集成到 CI/CD 管道中实现模型更新后的自动 Build优先使用 FP8 量化——它在精度损失 0.5% 的前提下实现近 2 倍的吞吐提升。