AI 推理优化实践-TensorRT-LLM 部署：A100 上 Llama-3-70B 吞吐量提升 8 倍

一、为什么 HF 原生推理跑不快先看 HuggingFacetransformers默认推理路径的瓶颈# HF 原生 generate逐 token、Python 循环、无批处理优化 for _ in range(max_new_tokens): outputs model(input_ids) # 全量前向 next_token sample(outputs) input_ids append(next_token)问题清单TensorRT-LLM 就是为消灭这些瓶颈而生的。二、TensorRT-LLM 四大核心优化2.1 算子融合Kernel Fusion把多个小算子融合成一个 CUDA kernel减少 kernel launch 次数和显存读写融合前LayerNorm → MatMul → BiasAdd → GeLU (4 次 kernel launch) 融合后FusedAttention(单一 kernel) (1 次 kernel launch)Transformer 层里大量 GEMM 激活 归一化被融合launch 开销下降一个量级。2.2 In-Flight Batching飞行中批处理这是 TensorRT-LLM 的杀手锏等价于 vLLM 的 Continuous Batching传统静态批处理 req1 ━━━━━━━━━━━等长都得等 req2 跑完 req2 ━━━━━提前结束但被占用直到 req1 完 ​ In-Flight Batching step1: [req1, req2, req3] ─┐ step2: [req1, req2, req4] ─┤ req3 结束立即插入 req4 step3: [req1, req5, req4] ─┘ req2 结束插入 req5新请求在迭代间隙动态插入GPU 利用率从 30% 拉到 90%。2.3 PagedAttention与 vLLM 同源思想KV Cache 分页管理消除显存碎片支持更大并发。2.4 量化FP8 / INT8 / AWQ# 权重 FP8H100 原生A100 用 INT8 近似 # KV Cache 也可量化FP8 KV Cache 省一半显存FP8 权重 FP8 KV Cache 在 H100 上可把吞吐再翻一倍。A100 不支持 FP8 硬件用 INT8 权重 SmoothQuant 近似。三、核心架构走读Session管理引擎执行和调度是 In-Flight Batching 的调度核心Tokenizer/SamplingPython 层负责前后处理TRT Enginetrtllm-build编译出的优化引擎包含所有融合 kernel四、部署实战从权重到服务4.1 环境准备# 官方镜像最省事包含所有依赖 docker run --gpus all -it --rm \ -v $PWD:/workspace \ nvcr.io/nvidia/tensorrt-llm:latest \ /bin/bash ​ pip install tensorrt_llm -U --pre4.2 第一步权重转换把 HF 格式转成 TensorRT-LLM 格式# 克隆转化脚本以 Llama-3-70B 为例 python convert_checkpoint.py \ --model_dir /workspace/Llama-3-70B \ --output_dir /workspace/llama3-70b-trt \ --dtype float16 \ --use_weight_only \ --weight_only_precision int8 \ # INT8 权重量化 --tp_size 4 # 4 卡张量并行--tp_size 4表示用 4 张 GPU 做张量并行切分 70B 模型每卡约 35GB FP16或 INT8 约 17.5GB。4.3 第二步编译引擎trtllm-build \ --checkpoint_dir /workspace/llama3-70b-trt \ --output_dir /workspace/llama3-70b-engine \ --gemm_plugin float16 \ --max_batch_size 64 \ --max_input_len 2048 \ --max_seq_len 4096 \ --paged_kv_cache enable \ --use_fused_mlp enable--paged_kv_cache enable开启分页 KV Cache--use_fused_mlp开启 MLP 融合。4.4 第三步推理测试# 命令行快速验证 python run.py \ --engine_dir /workspace/llama3-70b-engine \ --tokenizer_dir /workspace/Llama-3-70B \ --max_output_len 256 \ --input_text 用一句话解释什么是张量并行。 ​ # 或用 C 高性能服务 trtllm-serve \ --engine_dir /workspace/llama3-70b-engine \ --tokenizer_dir /workspace/Llama-3-70B \ --host 0.0.0.0 --port 80004.5 Python API 推理import tensorrt_llm from tensorrt_llm.runtime import ModelRunner ​ runner ModelRunner.from_dir( engine_dir/workspace/llama3-70b-engine, rank0, ) outputs runner.generate( batch_input_ids[[1, 15043, 29892, 1234]], # tokenized prompt max_new_tokens256, end_id2, pad_id2, )五、生产部署Triton Inference Server高性能生产环境用 Triton 托管 TensorRT-LLM# 启动 trtllm_backend docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \ -v $PWD/trt_llm_backend:/workspace \ nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3 config.pbtxt 关键配置 name: tensorrt_llm backend: tensorrtllm max_batch_size: 64 instance_group [{ kind: KIND_GPU, count: 4 }] # 4 GPU 实例 dynamic_batching { preferred_batch_size: [ 32, 64 ] max_queue_delay_microseconds: 500 }Triton 提供 gRPC/HTTP 接口自带动态批处理配合 In-Flight Batching 效果最佳。六、性能实测A100 三方案对比6.1 环境GPU4×A100-80G张量并行 tp4模型Llama-3-70B-Instruct输入 512 token输出 256 token测试工具官方bench.py6.2 结果* FP8 需 H100A100 上用 INT8 权重 SmoothQuant 近似。关键结论相比 HF 原生单并发快约 3.4 倍18→62 等效高并发快约 10.7 倍110→1180相比 vLLMTensorRT-LLM INT8 在 A100 上再快约 1.7 倍690→1180代价是部署复杂度更高显存占用从 140GB 降到 52GB4 卡轻松容纳还能上更大并发6.3 什么时候选谁场景推荐快速验证、灵活切换模型、社区生态vLLM极致吞吐、NVIDIA 硬件、生产固化TensorRT-LLM Triton国产/非 NVIDIA 卡昇腾、昆仑对应厂商推理引擎七、调优参数参数作用建议max_batch_size最大并发按显存设到 64/128max_num_tokens总 token 预算影响 In-Flight 调度paged_kv_cache分页 KV必须 enable量化精度性能/精度权衡H100 FP8A100 INT8tp_size张量并行度卡数或能整除层数的组合常见坑tp_size 不整除70B 用 tp4 可以每层 70 层/417.5 不整除实际按 hidden 维切没问题但 tp3 会报错max_seq_len 设太小长上下文直接截断INT8 精度掉点用 SmoothQuant 校准集否则激活值溢出Triton 队列延迟max_queue_delay_microseconds调大可提升批大小但增延迟八、总结TensorRT-LLM 是 NVIDIA 生态下的大模型推理性能天花板算子融合 In-Flight Batching PagedAttention 量化四管齐下A100 上 Llama-3-70B 高并发吞吐比 HF 原生快约 10 倍部署链路convert_checkpoint → trtllm-build → trtllm-serve/Triton代价是部署复杂度和硬件绑定需要权衡下一篇预告量化部署是推理优化的重头戏——AWQ 与 GPTQ 在 TensorRT-LLM / vLLM 上的落地差异以及端侧 INT4 推理下期继续。往期回顾AvLLM PagedAttention 解析显存利用率从 40% 提升到 90% 的秘密AI 推理优化实践vLLM Continuous Batching5800 t/s 吞吐量的核心引擎AI 推理优化:LoRA/QLoRA 微调原理单卡 24G 显存如何微调 70B 大模型关注获取更新本文属于「AI 推理优化实战」专栏点击关注获取推理加速系列后续内容。