1. 为什么我把LLaMA-2-7B的部署方案从PyTorch换成了TensorRT先说个背景。之前我在内部平台用HuggingFace的transformers直接跑LLaMA-2-7B-chat输入prompt后等结果模型本身在单卡上能跑但体感非常难受。一次普通的对话生成在单张A100 40G上decode阶段大概要30到40毫秒才能出一个token换算下来每秒只能吐二十来个字。这还不算并发一旦同时来几个请求显存直接飙到接近上限服务跟着卡顿。后来我把方案切到TensorRT-LLM同样一张卡、同一个模型权重decode延迟降到10到15毫秒一个token吞吐翻了接近三倍而且并发请求支持得更自然。这篇文章就把整个部署过程完整复现一遍从环境准备、权重转换、引擎构建到推理落测和常见坑位排查按实际动手的顺序写适合已经在跑通PyTorch推理、但想把LLaMA-2-7B真正推到生产环境里的同学参考。先说结论TensorRT部署LLM不是简单把PyTorch模型“导成engine文件”就完事而是要把自回归解码、KV Cache、张量并行、内存管理这些环节全部交给推理框架做统一编排。这也是为什么直接用TensorRT的Python API去搭一个LLM推理流程非常痛苦而NVIDIA干脆把FasterTransformer合并进来做了TensorRT-LLM这个专门面向大模型的推理框架。你理解为TensorRT是老底子算子优化引擎TensorRT-LLM是专门为大模型打架子鼓的完整解决方案。7B模型用前者需要自己拼装很多零件用后者基本上开箱即用。1.1 原生PyTorch推理慢在哪里回到本质问题为什么transformers跑LLaMA-2-7B那么慢LLM推理被分成两个阶段prefill和decode。prefill阶段一次处理整个输入序列并行度高GPU利用率也能拉起来decode阶段是逐个token生成的每生成一个token都要把整个模型的所有权重过一遍计算量是固定的但并行性极低。真正让时间拉长的是decode阶段里大量小算子的串行执行比如LayerNorm、attention的scale和mask、Softmax、矩阵乘法每个单独看都不慢但算子之间不断切换内核、读写显存累积起来耗时非常可观。第二个大问题是KV Cache的管理。自注意力机制需要缓存历史token的Key和Value在PyTorch原生实现里KV Cache通常是一个动态增长的tensor每次都要重新分配显存、拷贝数据这个开销在长上下文中会越来越明显。TensorRT-LLM采用预分配KV Cache池的方式一次性申请好固定大小的显存块推理时只做索引和写入省掉了大量显存拷贝。第三个问题是并发。原生HF推理在并发场景下要么自己写锁要么粗暴地每请求一个模型实例显存和卡算力都被白白浪费。TensorRT-LLM的inflight batching机制可以把不同请求的decode阶段合并成一个大batch执行这是吞吐提升最大的来源后面章节会细说。1.2 TensorRT/ TensorRT-LLM到底做了什么TensorRT做的事可以归纳为三点。第一图优化。把计算图里的算子进行融合比如把LayerNorm里的多个操作合并成一个自定义kernel减少kernel launch次数。启动一次kernel是有固定开销的几千个算子的图优化到几百个省下来的时间很可观。第二kernel autotuning。同一个矩阵乘法TensorRT会在构建期跑多个实现根据你的GPU型号选择最快的那个这个构建过程会持续几分钟到几十分钟所以你在跑trtllm-build的时候看到进度条半天不走不用慌它在做性能测试。第三量化支持。TensorRT对FP16、INT8、INT4、FP8都有成熟的kernel支持权重量化后模型体积和显存占用大幅下降而推理速度还能继续提升。TensorRT-LLM则是站在TensorRT肩膀上为自回归模型定制了一套东西预分配KV Cache池、PagedAttention实现、连续批处理调度器、张量并行和流水线并行的运行时支持还有各种量化格式的GPT Attention kernel也就是专门为解码器优化的融合注意力算子。这一整套东西才是让LLaMA-2-7B在单卡上跑得又快又稳的真正原因。2. 环境准备先对齐版本再谈部署TensorRT-LLM对环境的敏感程度夸张一点说比模型本身难搞。我自己第一次部署就栽在版本错配上面GPU驱动、CUDA、TensorRT、TensorRT-LLM这四个组件只要有一个不匹配后面构建出来的engine文件基本就是废的推理时直接报错或者结果完全乱掉。所以这里是整个部署里最不能跳的一步。2.1 硬件底线和显存预算先说硬件。LLaMA-2-7B的FP16权重大约是14GB如果只跑推理不训练一张显存16GB的卡理论最低容量是够的但会非常紧张因为除了权重之外KV Cache、激活值、运行时缓冲区都要占显存。我的建议是部署模式建议最低显存推荐卡型纯CPU演示无需GPU不推荐生产FP16单卡部署24GB以上A10G、L4、4090、A100 40GINT8/INT4量化单卡12GB-16GB4090、A5000、L40S多卡张量并行每卡16GB以上2xA10G、2x4090我这边的测试环境是单张A100 40GTensorRT-LLM版本0.7.0后续数据都出自这个环境。如果你用的是4090流程完全一样只是构建引擎时kernel autotune的结果会有些差异不影响整体操作。2.2 容器方案还是裸机部署个人强烈推荐直接用NVIDIA官方NGC容器。TensorRT-LLM的容器镜像里CUDA、TensorRT、cuDNN、TensorRT-LLM的版本关系都是测试好的你不需要自己折腾兼容性问题。裸机部署虽然可行但要手动对照版本矩阵很容易漏掉某个依赖的小版本差异踩坑成本远高于容器方式。我当时的做法是拉取NGC镜像nvcr.io/nvidia/tritonserver:24.01-trtllm-python-py3。这个镜像里有TensorRT-LLM的Python包省去编译环节。如果你需要特定版本可以在NGC页面找对应标签。2.3 TensorRT-LLM与CUDA、TensorRT的版本对应关系如果你坚持裸机部署版本关系需要严格对上。以我用的版本组合为例组件版本NVIDIA驱动535.104.05及以上CUDA12.2cuDNN8.9.xTensorRT9.2.1TensorRT-LLM0.7.0这个组合在2024年上半年是稳定的。到了0.8、0.9、0.10版本CUDA要求基本都升到了12.3以上TensorRT也跟着升级。我的经验是安装前先读一下对应版本TensorRT-LLM的README里的requirements部分里面一般会写清楚依赖组件的精确版本别只看大版本号。提示引擎文件和运行时版本是强绑定的。用TensorRT-LLM 0.7.0构建出来的engine拿到0.8.0环境里加载会直接失败或者更糟糕的是加载成功但输出乱码。所以生产环境一旦定好版本不要随意升级要升级就整套容器换掉重新build engine。3. 模型准备从HuggingFace权重到TensorRT-LLM CheckpointTensorRT-LLM不能直接吃HuggingFace的权重格式它需要先把权重转换成它自己的checkpoint格式再用这个checkpoint去构建engine。这个转换过程不改变模型参数只重新排列权重布局顺便为量化做准备。3.1 下载LLaMA-2-7B权重下载权重本身没什么好说的有HuggingFace的token就直接用huggingface-cli没有的话找模型文件手动下载也行。这里提醒一个细节如果你用Chat版做对话测试确保下载的是llama-2-7b-chat-hf在meta-llama组织下面仓库里包含config.json、tokenizer.model、tokenizer_config.json以及分片的safetensors权重文件。代码里缺tokenizer.model的情况我遇到过几次后面专门在坑位部分展开。3.2 权重转换命令进入TensorRT-LLM容器后先找到转换脚本。一般情况下在/opt/TensorRT-LLM/examples/llama/convert_checkpoint.py。执行python convert_checkpoint.py \ --model_dir ./llama-2-7b-chat-hf \ --output_dir ./tllm_checkpoint \ --dtype float16几个参数的含义我说明一下。--model_dir填的是HuggingFace权重目录。--output_dir是转换后的输出目录。--dtype float16指定权重精度这里选float16是因为7B模型在A100上FP16单卡能跑不需要量化。如果你显存紧张想跑INT8或INT4这个参数对应改成int8或int4转换脚本会自动对权重做量化处理精度损失在可接受范围内。转换完成后tllm_checkpoint目录下会出现config.json和一个或多个权重文件。config.json里记录了模型结构相关的配置包括层数、头数、维度、量化方式、张量并行参数等等。这个文件要保留好构建引擎时会用到。3.3 量化选项怎么取舍关于LLaMA-2-7B用不用量化我的建议是单卡显存24GB以上直接用FP16。部署简单省心精度无损。显存只有16GB用INT8重量化。7B的权重从14GB降到7GB左右给KV Cache留足空间实测生成质量几乎没有肉眼可见的下降。显存12GB左右考虑INT4。权重降到4GB以内但有明显效果差异长文本下偶尔出现逻辑不连贯的现象。追求极致吞吐的场景也可以尝试FP8这个需要H100或Ada架构的卡支持。我当时在4090上试过INT8配合max_num_tokens调整后时序延迟比FP16还要再低一点显存占用少了近一半。不过如果你第一次部署建议先全程跑FP16把流程走通确认没问题后再回去试量化这样排错范围会小很多。权重转换这一步还可能涉及张量并行和流水线并行的参数。单卡部署就跳过多卡部署时需要指定--tensor_parallel_size 2之类的参数比如两张卡就把模型切成两份各自放到一张卡上。转换后checkpoint里的权重就是按并行度切好的。4. 构建引擎trtllm-build参数决定推理上限转换得到checkpoint后用trtllm-build命令构建TensorRT engine。这一步是大多数人第一次跑时会比较懵的参数不多但每个参数都直接影响显存占用、并发能力和最终性能必须理解清楚再设置。4.1 最小可用的构建命令trtllm-build \ --checkpoint_dir ./tllm_checkpoint \ --output_dir ./tllm_engine \ --gemm_plugin float16 \ --max_batch_size 8 \ --max_input_len 2048 \ --max_seq_len 4096 \ --max_num_tokens 4096我在A100上跑完整个构建大约耗时15到25分钟取决于TensorRT autotune的搜索空间和机器负载。4090上时间也差不多。构建过程会打印大量日志其中包括每个层选择的kernel名称和对应的耗时不用去深究只需要确认最后没有error级别的报错并且能看到类似Engine generated这样的输出。4.2 max_batch_size、max_seq_len、max_num_tokens到底影响什么这几个参数很多教程一笔带过但实际部署时它们决定了你能开多大的并发、最多支持多长的上下文甚至直接决定引擎会不会在运行时OOM。max_batch_size是推理框架单次能处理的请求数上限。比如设为8表示同时最多8个序列在进行生成。注意这个值不是越大越好它直接影响KV Cache的预留大小后面会算这笔账。max_input_len和max_seq_len分别表示输入prompt的最大长度和生成完成后整条序列的最大长度。max_seq_len必须大于max_input_len因为还要算上生成token的数量。例如你想支持4096长度上下文、prompt最多2048那生成的最大token数就是2048。max_num_tokens是TensorRT-LLM分批调度时实际使用的关键参数它表示每次前向传播中所有batch序列加起来的最大token总数。TensorRT-LLM批量推理时把多个请求拼成一个大的token序列而不是传统意义上把所有请求pad到相同长度。max_num_tokens决定了一次前向能塞进去多少token它和max_batch_size是配合关系max_num_tokens要大于等于max_batch_size乘以单条序列最大长度的一部分。设置太小批量推理无法充分利用GPU算力设置太大显存会被未使用的高位预留空间占满。我实际使用时的经验值是让max_num_tokens约等于max_batch_size乘以平均生成长度。比如batch_size设8每条序列生成512个token那max_num_tokens设置在4096到8192之间比较合理。4.3 显存预算怎么算构建引擎时预留的KV Cache显存大致可以用这个公式估算KV Cache大小 ≈ 2Key和Value × 层数 × 每层KV维度 × max_num_tokens × 精度字节数LLaMA-2-7B有32层每层KV头数为32每个头的维度是128计算得到每层KV维度是4096。FP16时每个token的KV Cache占用约32 × 4096 × 2字节 262144字节也就是256KB。当max_num_tokens 4096时KV Cache预留就是1GB当max_num_tokens 8192时预留约2GB。看起来还好但如果把max_num_tokens拉到32768那预留就变成8GB配合上14GB权重A100 40G还能扛住但16GB的卡就会非常危险。这也是为什么小显存卡上不能盲目开高并发的原因。4.4 构建日志里有哪些值得看的信号构建过程中日志里会出现很多[TensorRT]前缀的输出值得留意的有三类第一类是[TensorRT] Detected X CUDA capable devices确认识别到了你的GPU。如果你在容器里没挂载GPU或驱动版本不对这行会显示0设备后面构建必然失败。第二类是[TensorRT] Performance tuning相关日志说明正在做kernel autotune这个阶段GPU利用率会很高风扇会转起来是正常现象。第三类是error级别的信息任何带error的字样都要停下来。常见错误是CUDA相关库加载失败或者模型结构和TensorRT-LLM版本不匹配。遇到error后优先检查是不是前面章节说的版本错配问题。构建完成后tllm_engine目录下会有config.json和几个.engine文件。之后推理环节加载的就是这个目录。5. 端到端落测加载引擎、跑通生成、验证输出引擎构建好之后进入推理阶段。TensorRT-LLM的Python API我用的是顶层的LLM抽象类接口非常简洁比直接操作Executor API省事得多特别适合先跑通流程再做性能验证。5.1 第一个推理脚本from tensorrt_llm import LLM, SamplingParams llm LLM(model./tllm_engine, tensor_parallel_size1) prompt 介绍一下Transformer中的自注意力机制包括它的作用和计算过程。 sampling_params SamplingParams( max_tokens512, temperature0.7, top_p0.9, repetition_penalty1.0 ) outputs llm.generate([prompt], sampling_paramssampling_params) for output in outputs: text output.outputs[0].text print(text) print( * 50) print(f生成token数: {len(output.outputs[0].token_ids)})这段代码有几个细节需要解释。LLM(model./tllm_engine)传入的是引擎目录TensorRT-LLM会读取里面的config.json和engine文件。如果你的引擎是用多卡并行构建的这里要指定tensor_parallel_size为对应的卡数。SamplingParams里的参数和transformers里的感觉类似但底层实现不同max_tokens表示最多生成多少token这个值不能超过引擎构建时max_seq_len - max_input_len的上限。5.2 输出结果验证跑通之后第一件事是验证输出质量。我之前踩过一个坑引擎构建成功、推理也没报错但输出的是完全乱码。后来发现是tokenizer和模型不匹配。我的验证方法是拿同一个prompt在HuggingFace的transformers里用同样的SamplingParams跑一遍对比两边的输出。正常情况下内容应该基本一致小差异来自kernel计算精度浮动属于正常范围。如果完全驴唇不对马嘴优先检查tokenizer.config中的tokenizer_class以及tokenizer.model文件是否存在。TensorRT-LLM在初始化LLM对象时默认会读取模型目录里的tokenizer配置。如果你把engine文件复制到别的机器运行一定要把tokenizer文件一并带上我的建议是把下载的llama-2-7b-chat-hf目录里的tokenizer相关文件单独拷贝到一个tokenizer/子目录里加载时再指定llm LLM( model./tllm_engine, tokenizer_dir./tokenizer, tensor_parallel_size1 )这样部署时不会漏文件排查问题也容易定位。5.3 和HF基线做一个简单的延迟对比跑通功能后我做了一个benchmark同样的prompt长度和生成长度对比HF和TensorRT-LLM的耗时。环境硬件模型精度decode阶段每token延迟显存占用HF transformersA100 40GFP1635-45ms接近20GBTensorRT-LLM默认配置A100 40GFP1612-18ms约17GBTensorRT-LLM inflightA100 40GFP168-12ms并发场景约17GB单看decode延迟TensorRT-LLM优化后大概是HF的三倍左右。但真正拉开差距的是并发场景HF在8路并发时需要启动8份模型或者自己写复杂的调度逻辑TensorRT-LLM靠inflight batching可以把这8个请求合并在一次前向传播里整体吞吐提升不是线性三倍而是五到八倍的量级。6. 生产化调优inflight batching、KV Cache复用和实测参数引擎能跑通只是第一步。实际部署到服务里还需要根据自己的业务场景做针对性配置。这一节讲三个关键方向连续批处理、KV Cache管理、以及一个可以照抄的参数组合。6.1 Inflight Batching是吞吐翻倍的最大功臣传统静态批处理的做法是攒够N个请求再统一推理所有请求都必须等待当前batch完成后才能进去这个等待期GPU可能还在做无用功。inflight batching的思路则是“边生成边调度”每个请求的每个生成步骤都是独立的有的请求生成了10个token有的才生成2个token调度器会在每一步动态组合当前所有活跃请求拼成一个最大利用率的batch。已经完成生成的请求立即释放KV Cache新的请求立刻补位GPU全程不闲着。在TensorRT-LLM中这个能力默认是开启的但需要你在服务化部署时正确设置容量参数。如果直接使用LLM.generate的Python API实际上它也按批处理调度执行。我用8路并发测试时inflight batching开启后单卡吞吐从每秒不到200token提升到接近600token提升非常明显。6.2 KV Cache复用和PagedAttentionPagedAttention是vLLM带火的概念TensorRT-LLM也有类似实现。它把KV Cache划分成固定大小的块按需分配而不是为每条序列预留整块连续空间。这样做有两个好处长序列场景下显存碎片少利用率高。多个序列可以共享同一个物理块比如不同请求有相同的前缀这部分KV Cache只需要计算一次。如果你跑的业务里很多请求share同一个系统提示词这个优化收益会很大。不过TensorRT-LLM在API层面没有直接暴露前缀共享的配置项它更多体现在多请求调度时的显存效率上不需要你手动干预。6.3 我实测下来的一组推荐参数如果你不想纠结一大堆参数组合我直接给你一组在A100 40G上验证过的配置单卡、FP16、服务8路并发请求都能稳跑trtllm-build \ --checkpoint_dir ./tllm_checkpoint \ --output_dir ./tllm_engine \ --gemm_plugin float16 \ --max_batch_size 8 \ --max_input_len 2048 \ --max_seq_len 4096 \ --max_num_tokens 4096推理端sampling_params SamplingParams( max_tokens512, temperature0.6, top_p0.9 )这套配置下单请求decode延迟约15ms/token8路并发时整体TPOT约22ms/token显存峰值约17GB留了大约20GB余量给后续扩展。如果显存是16GB的卡把max_batch_size降到4、max_num_tokens降到2048再用INT8量化也能稳定运行。7. 实操中绕不开的几个坑从GCC版本到显存计算最后这部分是我实际部署中踩过的坑每个都花了不少时间排查。按出现频率从高到低列出来希望帮你避开。7.1 版本错配是最大坑源引擎构建成功、加载时报错或者输出乱码十个里有七个是版本问题。我见过两种典型情况第一种是TensorRT-LLM和TensorRT版本不匹配构建时就报出各种算子不支持的错误。第二种是运行时容器和构建时容器版本不一致engine文件加载失败。我的建议是构建和运行使用完全相同的容器镜像千万别图省事把engine文件拷出来就直接在新的环境里启动。即使是小版本升级也老老实实重新build一次再上生产。7.2 GCC版本引发的C扩展编译问题如果你选择裸机安装TensorRT-LLM而不是用容器大概率会遇到GCC和G版本过新导致编译失败的问题。TensorRT-LLM的C部分对编译器版本有要求我用的是GCC 11.3在Ubuntu 22.04上编译顺利。如果系统自带GCC 13建议用update-alternatives切换回GCC 11或直接配一个conda环境指定编译器版本后再pip安装。7.3 “显存OOM”不一定是显存不够第一次调并发参数时我把max_num_tokens从4096改成16384结果启动服务直接OOM。当时第一反应是A100 40G不够了后来排查发现权重14GB KV Cache预留4GB 中间激活缓冲区2GB合计20GB左右40G显存完全够。真正的原因是Python端在加载engine时还额外分配了和max_num_tokens相关的临时buffer以及运行时context workspace导致瞬时峰值远超预期。解决办法是不要把max_num_tokens拉太高尤其在并发量没那么大的场景4096已经能满足大多数业务。如果确实需要高并发优先用INT8或INT4量化释放显存空间。7.4 tokenizer版本对生成质量的影响还有一个容易被忽略的点transformers版本会影响LLaMA-2 tokenizer的加载行为。新版transformers在加载LLaMA-2时如果检测到tokenizer.model缺失会自动尝试用fast tokenizer替代但效果可能不完全一致。具体表现为特殊token编号偏移、生成内容莫名多出空白字符或者[INST]标签没被正确解析。遇到这类问题先确认你下载的模型目录里有没有tokenizer.model文件没有的话从原始repo单独下载并放到tokenizer_dir里显式指定。7.5 首次构建时间“看起来像卡死”trtllm-build在kernel autotune阶段单看日志输出频率很低几分钟可能都没有新内容特别像卡死了。我第一跑等了三分钟没看到新日志差点CtrlC。实际上autotune阶段就是这个风格它可以跑十几分钟才输出关键日志。建议构建时用nvidia-smi在另一个终端确认GPU利用率是否波动只要GPU有计算活动就耐心等着。收尾的一点体会这次部署下来我最大的感受是TensorRT-LLM的难点不在“跑通”而在“跑得好”。跑通只需要对着命令敲一遍真正决定线上表现的是前面那些参数背后的显存计算逻辑、批处理调度机制和版本配套关系。尤其是max_num_tokens和max_batch_size这两项直接影响GPU利用率和显存天花板建议大家在调优时多花点时间拿nvidia-smi盯着显存变化曲线每次只改一个参数确认效果后再动下一个。这样一套流程走下来你对LLM推理的端到端链路认知会扎实很多之后换别的模型、换更大的卡也能很快上手。
LLaMA-2-7B部署提速：基于TensorRT-LLM的生产级推理优化指南
1. 为什么我把LLaMA-2-7B的部署方案从PyTorch换成了TensorRT先说个背景。之前我在内部平台用HuggingFace的transformers直接跑LLaMA-2-7B-chat输入prompt后等结果模型本身在单卡上能跑但体感非常难受。一次普通的对话生成在单张A100 40G上decode阶段大概要30到40毫秒才能出一个token换算下来每秒只能吐二十来个字。这还不算并发一旦同时来几个请求显存直接飙到接近上限服务跟着卡顿。后来我把方案切到TensorRT-LLM同样一张卡、同一个模型权重decode延迟降到10到15毫秒一个token吞吐翻了接近三倍而且并发请求支持得更自然。这篇文章就把整个部署过程完整复现一遍从环境准备、权重转换、引擎构建到推理落测和常见坑位排查按实际动手的顺序写适合已经在跑通PyTorch推理、但想把LLaMA-2-7B真正推到生产环境里的同学参考。先说结论TensorRT部署LLM不是简单把PyTorch模型“导成engine文件”就完事而是要把自回归解码、KV Cache、张量并行、内存管理这些环节全部交给推理框架做统一编排。这也是为什么直接用TensorRT的Python API去搭一个LLM推理流程非常痛苦而NVIDIA干脆把FasterTransformer合并进来做了TensorRT-LLM这个专门面向大模型的推理框架。你理解为TensorRT是老底子算子优化引擎TensorRT-LLM是专门为大模型打架子鼓的完整解决方案。7B模型用前者需要自己拼装很多零件用后者基本上开箱即用。1.1 原生PyTorch推理慢在哪里回到本质问题为什么transformers跑LLaMA-2-7B那么慢LLM推理被分成两个阶段prefill和decode。prefill阶段一次处理整个输入序列并行度高GPU利用率也能拉起来decode阶段是逐个token生成的每生成一个token都要把整个模型的所有权重过一遍计算量是固定的但并行性极低。真正让时间拉长的是decode阶段里大量小算子的串行执行比如LayerNorm、attention的scale和mask、Softmax、矩阵乘法每个单独看都不慢但算子之间不断切换内核、读写显存累积起来耗时非常可观。第二个大问题是KV Cache的管理。自注意力机制需要缓存历史token的Key和Value在PyTorch原生实现里KV Cache通常是一个动态增长的tensor每次都要重新分配显存、拷贝数据这个开销在长上下文中会越来越明显。TensorRT-LLM采用预分配KV Cache池的方式一次性申请好固定大小的显存块推理时只做索引和写入省掉了大量显存拷贝。第三个问题是并发。原生HF推理在并发场景下要么自己写锁要么粗暴地每请求一个模型实例显存和卡算力都被白白浪费。TensorRT-LLM的inflight batching机制可以把不同请求的decode阶段合并成一个大batch执行这是吞吐提升最大的来源后面章节会细说。1.2 TensorRT/ TensorRT-LLM到底做了什么TensorRT做的事可以归纳为三点。第一图优化。把计算图里的算子进行融合比如把LayerNorm里的多个操作合并成一个自定义kernel减少kernel launch次数。启动一次kernel是有固定开销的几千个算子的图优化到几百个省下来的时间很可观。第二kernel autotuning。同一个矩阵乘法TensorRT会在构建期跑多个实现根据你的GPU型号选择最快的那个这个构建过程会持续几分钟到几十分钟所以你在跑trtllm-build的时候看到进度条半天不走不用慌它在做性能测试。第三量化支持。TensorRT对FP16、INT8、INT4、FP8都有成熟的kernel支持权重量化后模型体积和显存占用大幅下降而推理速度还能继续提升。TensorRT-LLM则是站在TensorRT肩膀上为自回归模型定制了一套东西预分配KV Cache池、PagedAttention实现、连续批处理调度器、张量并行和流水线并行的运行时支持还有各种量化格式的GPT Attention kernel也就是专门为解码器优化的融合注意力算子。这一整套东西才是让LLaMA-2-7B在单卡上跑得又快又稳的真正原因。2. 环境准备先对齐版本再谈部署TensorRT-LLM对环境的敏感程度夸张一点说比模型本身难搞。我自己第一次部署就栽在版本错配上面GPU驱动、CUDA、TensorRT、TensorRT-LLM这四个组件只要有一个不匹配后面构建出来的engine文件基本就是废的推理时直接报错或者结果完全乱掉。所以这里是整个部署里最不能跳的一步。2.1 硬件底线和显存预算先说硬件。LLaMA-2-7B的FP16权重大约是14GB如果只跑推理不训练一张显存16GB的卡理论最低容量是够的但会非常紧张因为除了权重之外KV Cache、激活值、运行时缓冲区都要占显存。我的建议是部署模式建议最低显存推荐卡型纯CPU演示无需GPU不推荐生产FP16单卡部署24GB以上A10G、L4、4090、A100 40GINT8/INT4量化单卡12GB-16GB4090、A5000、L40S多卡张量并行每卡16GB以上2xA10G、2x4090我这边的测试环境是单张A100 40GTensorRT-LLM版本0.7.0后续数据都出自这个环境。如果你用的是4090流程完全一样只是构建引擎时kernel autotune的结果会有些差异不影响整体操作。2.2 容器方案还是裸机部署个人强烈推荐直接用NVIDIA官方NGC容器。TensorRT-LLM的容器镜像里CUDA、TensorRT、cuDNN、TensorRT-LLM的版本关系都是测试好的你不需要自己折腾兼容性问题。裸机部署虽然可行但要手动对照版本矩阵很容易漏掉某个依赖的小版本差异踩坑成本远高于容器方式。我当时的做法是拉取NGC镜像nvcr.io/nvidia/tritonserver:24.01-trtllm-python-py3。这个镜像里有TensorRT-LLM的Python包省去编译环节。如果你需要特定版本可以在NGC页面找对应标签。2.3 TensorRT-LLM与CUDA、TensorRT的版本对应关系如果你坚持裸机部署版本关系需要严格对上。以我用的版本组合为例组件版本NVIDIA驱动535.104.05及以上CUDA12.2cuDNN8.9.xTensorRT9.2.1TensorRT-LLM0.7.0这个组合在2024年上半年是稳定的。到了0.8、0.9、0.10版本CUDA要求基本都升到了12.3以上TensorRT也跟着升级。我的经验是安装前先读一下对应版本TensorRT-LLM的README里的requirements部分里面一般会写清楚依赖组件的精确版本别只看大版本号。提示引擎文件和运行时版本是强绑定的。用TensorRT-LLM 0.7.0构建出来的engine拿到0.8.0环境里加载会直接失败或者更糟糕的是加载成功但输出乱码。所以生产环境一旦定好版本不要随意升级要升级就整套容器换掉重新build engine。3. 模型准备从HuggingFace权重到TensorRT-LLM CheckpointTensorRT-LLM不能直接吃HuggingFace的权重格式它需要先把权重转换成它自己的checkpoint格式再用这个checkpoint去构建engine。这个转换过程不改变模型参数只重新排列权重布局顺便为量化做准备。3.1 下载LLaMA-2-7B权重下载权重本身没什么好说的有HuggingFace的token就直接用huggingface-cli没有的话找模型文件手动下载也行。这里提醒一个细节如果你用Chat版做对话测试确保下载的是llama-2-7b-chat-hf在meta-llama组织下面仓库里包含config.json、tokenizer.model、tokenizer_config.json以及分片的safetensors权重文件。代码里缺tokenizer.model的情况我遇到过几次后面专门在坑位部分展开。3.2 权重转换命令进入TensorRT-LLM容器后先找到转换脚本。一般情况下在/opt/TensorRT-LLM/examples/llama/convert_checkpoint.py。执行python convert_checkpoint.py \ --model_dir ./llama-2-7b-chat-hf \ --output_dir ./tllm_checkpoint \ --dtype float16几个参数的含义我说明一下。--model_dir填的是HuggingFace权重目录。--output_dir是转换后的输出目录。--dtype float16指定权重精度这里选float16是因为7B模型在A100上FP16单卡能跑不需要量化。如果你显存紧张想跑INT8或INT4这个参数对应改成int8或int4转换脚本会自动对权重做量化处理精度损失在可接受范围内。转换完成后tllm_checkpoint目录下会出现config.json和一个或多个权重文件。config.json里记录了模型结构相关的配置包括层数、头数、维度、量化方式、张量并行参数等等。这个文件要保留好构建引擎时会用到。3.3 量化选项怎么取舍关于LLaMA-2-7B用不用量化我的建议是单卡显存24GB以上直接用FP16。部署简单省心精度无损。显存只有16GB用INT8重量化。7B的权重从14GB降到7GB左右给KV Cache留足空间实测生成质量几乎没有肉眼可见的下降。显存12GB左右考虑INT4。权重降到4GB以内但有明显效果差异长文本下偶尔出现逻辑不连贯的现象。追求极致吞吐的场景也可以尝试FP8这个需要H100或Ada架构的卡支持。我当时在4090上试过INT8配合max_num_tokens调整后时序延迟比FP16还要再低一点显存占用少了近一半。不过如果你第一次部署建议先全程跑FP16把流程走通确认没问题后再回去试量化这样排错范围会小很多。权重转换这一步还可能涉及张量并行和流水线并行的参数。单卡部署就跳过多卡部署时需要指定--tensor_parallel_size 2之类的参数比如两张卡就把模型切成两份各自放到一张卡上。转换后checkpoint里的权重就是按并行度切好的。4. 构建引擎trtllm-build参数决定推理上限转换得到checkpoint后用trtllm-build命令构建TensorRT engine。这一步是大多数人第一次跑时会比较懵的参数不多但每个参数都直接影响显存占用、并发能力和最终性能必须理解清楚再设置。4.1 最小可用的构建命令trtllm-build \ --checkpoint_dir ./tllm_checkpoint \ --output_dir ./tllm_engine \ --gemm_plugin float16 \ --max_batch_size 8 \ --max_input_len 2048 \ --max_seq_len 4096 \ --max_num_tokens 4096我在A100上跑完整个构建大约耗时15到25分钟取决于TensorRT autotune的搜索空间和机器负载。4090上时间也差不多。构建过程会打印大量日志其中包括每个层选择的kernel名称和对应的耗时不用去深究只需要确认最后没有error级别的报错并且能看到类似Engine generated这样的输出。4.2 max_batch_size、max_seq_len、max_num_tokens到底影响什么这几个参数很多教程一笔带过但实际部署时它们决定了你能开多大的并发、最多支持多长的上下文甚至直接决定引擎会不会在运行时OOM。max_batch_size是推理框架单次能处理的请求数上限。比如设为8表示同时最多8个序列在进行生成。注意这个值不是越大越好它直接影响KV Cache的预留大小后面会算这笔账。max_input_len和max_seq_len分别表示输入prompt的最大长度和生成完成后整条序列的最大长度。max_seq_len必须大于max_input_len因为还要算上生成token的数量。例如你想支持4096长度上下文、prompt最多2048那生成的最大token数就是2048。max_num_tokens是TensorRT-LLM分批调度时实际使用的关键参数它表示每次前向传播中所有batch序列加起来的最大token总数。TensorRT-LLM批量推理时把多个请求拼成一个大的token序列而不是传统意义上把所有请求pad到相同长度。max_num_tokens决定了一次前向能塞进去多少token它和max_batch_size是配合关系max_num_tokens要大于等于max_batch_size乘以单条序列最大长度的一部分。设置太小批量推理无法充分利用GPU算力设置太大显存会被未使用的高位预留空间占满。我实际使用时的经验值是让max_num_tokens约等于max_batch_size乘以平均生成长度。比如batch_size设8每条序列生成512个token那max_num_tokens设置在4096到8192之间比较合理。4.3 显存预算怎么算构建引擎时预留的KV Cache显存大致可以用这个公式估算KV Cache大小 ≈ 2Key和Value × 层数 × 每层KV维度 × max_num_tokens × 精度字节数LLaMA-2-7B有32层每层KV头数为32每个头的维度是128计算得到每层KV维度是4096。FP16时每个token的KV Cache占用约32 × 4096 × 2字节 262144字节也就是256KB。当max_num_tokens 4096时KV Cache预留就是1GB当max_num_tokens 8192时预留约2GB。看起来还好但如果把max_num_tokens拉到32768那预留就变成8GB配合上14GB权重A100 40G还能扛住但16GB的卡就会非常危险。这也是为什么小显存卡上不能盲目开高并发的原因。4.4 构建日志里有哪些值得看的信号构建过程中日志里会出现很多[TensorRT]前缀的输出值得留意的有三类第一类是[TensorRT] Detected X CUDA capable devices确认识别到了你的GPU。如果你在容器里没挂载GPU或驱动版本不对这行会显示0设备后面构建必然失败。第二类是[TensorRT] Performance tuning相关日志说明正在做kernel autotune这个阶段GPU利用率会很高风扇会转起来是正常现象。第三类是error级别的信息任何带error的字样都要停下来。常见错误是CUDA相关库加载失败或者模型结构和TensorRT-LLM版本不匹配。遇到error后优先检查是不是前面章节说的版本错配问题。构建完成后tllm_engine目录下会有config.json和几个.engine文件。之后推理环节加载的就是这个目录。5. 端到端落测加载引擎、跑通生成、验证输出引擎构建好之后进入推理阶段。TensorRT-LLM的Python API我用的是顶层的LLM抽象类接口非常简洁比直接操作Executor API省事得多特别适合先跑通流程再做性能验证。5.1 第一个推理脚本from tensorrt_llm import LLM, SamplingParams llm LLM(model./tllm_engine, tensor_parallel_size1) prompt 介绍一下Transformer中的自注意力机制包括它的作用和计算过程。 sampling_params SamplingParams( max_tokens512, temperature0.7, top_p0.9, repetition_penalty1.0 ) outputs llm.generate([prompt], sampling_paramssampling_params) for output in outputs: text output.outputs[0].text print(text) print( * 50) print(f生成token数: {len(output.outputs[0].token_ids)})这段代码有几个细节需要解释。LLM(model./tllm_engine)传入的是引擎目录TensorRT-LLM会读取里面的config.json和engine文件。如果你的引擎是用多卡并行构建的这里要指定tensor_parallel_size为对应的卡数。SamplingParams里的参数和transformers里的感觉类似但底层实现不同max_tokens表示最多生成多少token这个值不能超过引擎构建时max_seq_len - max_input_len的上限。5.2 输出结果验证跑通之后第一件事是验证输出质量。我之前踩过一个坑引擎构建成功、推理也没报错但输出的是完全乱码。后来发现是tokenizer和模型不匹配。我的验证方法是拿同一个prompt在HuggingFace的transformers里用同样的SamplingParams跑一遍对比两边的输出。正常情况下内容应该基本一致小差异来自kernel计算精度浮动属于正常范围。如果完全驴唇不对马嘴优先检查tokenizer.config中的tokenizer_class以及tokenizer.model文件是否存在。TensorRT-LLM在初始化LLM对象时默认会读取模型目录里的tokenizer配置。如果你把engine文件复制到别的机器运行一定要把tokenizer文件一并带上我的建议是把下载的llama-2-7b-chat-hf目录里的tokenizer相关文件单独拷贝到一个tokenizer/子目录里加载时再指定llm LLM( model./tllm_engine, tokenizer_dir./tokenizer, tensor_parallel_size1 )这样部署时不会漏文件排查问题也容易定位。5.3 和HF基线做一个简单的延迟对比跑通功能后我做了一个benchmark同样的prompt长度和生成长度对比HF和TensorRT-LLM的耗时。环境硬件模型精度decode阶段每token延迟显存占用HF transformersA100 40GFP1635-45ms接近20GBTensorRT-LLM默认配置A100 40GFP1612-18ms约17GBTensorRT-LLM inflightA100 40GFP168-12ms并发场景约17GB单看decode延迟TensorRT-LLM优化后大概是HF的三倍左右。但真正拉开差距的是并发场景HF在8路并发时需要启动8份模型或者自己写复杂的调度逻辑TensorRT-LLM靠inflight batching可以把这8个请求合并在一次前向传播里整体吞吐提升不是线性三倍而是五到八倍的量级。6. 生产化调优inflight batching、KV Cache复用和实测参数引擎能跑通只是第一步。实际部署到服务里还需要根据自己的业务场景做针对性配置。这一节讲三个关键方向连续批处理、KV Cache管理、以及一个可以照抄的参数组合。6.1 Inflight Batching是吞吐翻倍的最大功臣传统静态批处理的做法是攒够N个请求再统一推理所有请求都必须等待当前batch完成后才能进去这个等待期GPU可能还在做无用功。inflight batching的思路则是“边生成边调度”每个请求的每个生成步骤都是独立的有的请求生成了10个token有的才生成2个token调度器会在每一步动态组合当前所有活跃请求拼成一个最大利用率的batch。已经完成生成的请求立即释放KV Cache新的请求立刻补位GPU全程不闲着。在TensorRT-LLM中这个能力默认是开启的但需要你在服务化部署时正确设置容量参数。如果直接使用LLM.generate的Python API实际上它也按批处理调度执行。我用8路并发测试时inflight batching开启后单卡吞吐从每秒不到200token提升到接近600token提升非常明显。6.2 KV Cache复用和PagedAttentionPagedAttention是vLLM带火的概念TensorRT-LLM也有类似实现。它把KV Cache划分成固定大小的块按需分配而不是为每条序列预留整块连续空间。这样做有两个好处长序列场景下显存碎片少利用率高。多个序列可以共享同一个物理块比如不同请求有相同的前缀这部分KV Cache只需要计算一次。如果你跑的业务里很多请求share同一个系统提示词这个优化收益会很大。不过TensorRT-LLM在API层面没有直接暴露前缀共享的配置项它更多体现在多请求调度时的显存效率上不需要你手动干预。6.3 我实测下来的一组推荐参数如果你不想纠结一大堆参数组合我直接给你一组在A100 40G上验证过的配置单卡、FP16、服务8路并发请求都能稳跑trtllm-build \ --checkpoint_dir ./tllm_checkpoint \ --output_dir ./tllm_engine \ --gemm_plugin float16 \ --max_batch_size 8 \ --max_input_len 2048 \ --max_seq_len 4096 \ --max_num_tokens 4096推理端sampling_params SamplingParams( max_tokens512, temperature0.6, top_p0.9 )这套配置下单请求decode延迟约15ms/token8路并发时整体TPOT约22ms/token显存峰值约17GB留了大约20GB余量给后续扩展。如果显存是16GB的卡把max_batch_size降到4、max_num_tokens降到2048再用INT8量化也能稳定运行。7. 实操中绕不开的几个坑从GCC版本到显存计算最后这部分是我实际部署中踩过的坑每个都花了不少时间排查。按出现频率从高到低列出来希望帮你避开。7.1 版本错配是最大坑源引擎构建成功、加载时报错或者输出乱码十个里有七个是版本问题。我见过两种典型情况第一种是TensorRT-LLM和TensorRT版本不匹配构建时就报出各种算子不支持的错误。第二种是运行时容器和构建时容器版本不一致engine文件加载失败。我的建议是构建和运行使用完全相同的容器镜像千万别图省事把engine文件拷出来就直接在新的环境里启动。即使是小版本升级也老老实实重新build一次再上生产。7.2 GCC版本引发的C扩展编译问题如果你选择裸机安装TensorRT-LLM而不是用容器大概率会遇到GCC和G版本过新导致编译失败的问题。TensorRT-LLM的C部分对编译器版本有要求我用的是GCC 11.3在Ubuntu 22.04上编译顺利。如果系统自带GCC 13建议用update-alternatives切换回GCC 11或直接配一个conda环境指定编译器版本后再pip安装。7.3 “显存OOM”不一定是显存不够第一次调并发参数时我把max_num_tokens从4096改成16384结果启动服务直接OOM。当时第一反应是A100 40G不够了后来排查发现权重14GB KV Cache预留4GB 中间激活缓冲区2GB合计20GB左右40G显存完全够。真正的原因是Python端在加载engine时还额外分配了和max_num_tokens相关的临时buffer以及运行时context workspace导致瞬时峰值远超预期。解决办法是不要把max_num_tokens拉太高尤其在并发量没那么大的场景4096已经能满足大多数业务。如果确实需要高并发优先用INT8或INT4量化释放显存空间。7.4 tokenizer版本对生成质量的影响还有一个容易被忽略的点transformers版本会影响LLaMA-2 tokenizer的加载行为。新版transformers在加载LLaMA-2时如果检测到tokenizer.model缺失会自动尝试用fast tokenizer替代但效果可能不完全一致。具体表现为特殊token编号偏移、生成内容莫名多出空白字符或者[INST]标签没被正确解析。遇到这类问题先确认你下载的模型目录里有没有tokenizer.model文件没有的话从原始repo单独下载并放到tokenizer_dir里显式指定。7.5 首次构建时间“看起来像卡死”trtllm-build在kernel autotune阶段单看日志输出频率很低几分钟可能都没有新内容特别像卡死了。我第一跑等了三分钟没看到新日志差点CtrlC。实际上autotune阶段就是这个风格它可以跑十几分钟才输出关键日志。建议构建时用nvidia-smi在另一个终端确认GPU利用率是否波动只要GPU有计算活动就耐心等着。收尾的一点体会这次部署下来我最大的感受是TensorRT-LLM的难点不在“跑通”而在“跑得好”。跑通只需要对着命令敲一遍真正决定线上表现的是前面那些参数背后的显存计算逻辑、批处理调度机制和版本配套关系。尤其是max_num_tokens和max_batch_size这两项直接影响GPU利用率和显存天花板建议大家在调优时多花点时间拿nvidia-smi盯着显存变化曲线每次只改一个参数确认效果后再动下一个。这样一套流程走下来你对LLM推理的端到端链路认知会扎实很多之后换别的模型、换更大的卡也能很快上手。
