ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Hopper架构TensorRT推理优化:TMA、WGMMA与FP8实战指南

Hopper架构TensorRT推理优化:TMA、WGMMA与FP8实战指南 1. 为什么 Hopper 架构值得单独聊 TensorRT 优化手里有 H100 或者 H200 的兄弟应该都有感觉这两张卡跑常规的 TensorRT 推理和 A100 比起来提升确实明显但如果你只是把 A100 上的 engine 直接搬过来重新编译一遍那基本只发挥了 Hopper 六成左右的功力。我从去年开始陆续在 H100 集群上做推理部署踩了不少坑之后才慢慢摸清楚Hopper 这一代在硬件层面加了很多专门为推理服务的东西TensorRT 也针对性地做了大量适配但这些特性默认不一定全部打开需要你在构建 engine 的时候有意识地配置。这篇文章主要面向已经在用 TensorRT 做推理部署、手里有 H100 或 H200 的工程师或者正在评估要不要上 Hopper 做推理服务的团队。我会把 Hopper 上 TensorRT 的几个关键新特性拆开讲清楚包括它们背后的硬件原理、在什么场景下能吃到红利、怎么在构建 engine 的时候把这些能力释放出来。同时也会聊到硬件级 kernel 这个话题很多人对 TensorRT 生成的 kernel 到底跑在什么单元上、为什么 Hopper 上某些 kernel 能快这么多其实是一知半解的我尽量用大白话把这块讲透。先给一个整体判断Hopper 上的 TensorRT 优化核心抓手就三个——Transformer Engine 的 FP8 支持、TMATensor Memory Accelerator带来的访存效率提升、以及分布式推理场景下的多卡通信优化。这三个方向各自对应不同的硬件单元和 kernel 实现下面逐个展开。2. Hopper 硬件级 kernel 到底新在哪里2.1 从 SM 到 TMA访存路径的重新设计要理解 Hopper 上 TensorRT 为什么能快得先知道 Hopper 的 SMStreaming Multiprocessor内部发生了什么变化。A100 那一代数据从 global memory 搬到 shared memory 再送到 register 做计算这条路径上的搬运工作主要由线程自己发指令完成也就是你写 CUDA kernel 时常用的cp.async那一套。到了 HopperNVIDIA 直接塞了一个专门的硬件单元叫 TMA全称 Tensor Memory Accelerator它可以理解成一个专职做多维张量搬运的 DMA 引擎。TMA 的价值在于它把地址计算、边界处理、数据搬运这些脏活从计算线程手里接过去了。以前一个 warp 里得专门派几个线程去算地址、发 load 指令现在你只需要在 host 端或者 device 端描述一个 tensor map告诉 TMA 数据在哪、什么形状、要搬到哪块 shared memory剩下的它自己搞定。计算线程被彻底解放出来可以全程泡在 tensor core 或者 FP32 单元里做计算。TensorRT 在 Hopper 上生成的很多 kernel尤其是那些卷积和 GEMM 相关的底层已经切到了 TMA 路径。你在 profiler 里看 kernel 名字带tma或者bulk字样的基本就是走了这条新路径。实测下来在 batch size 比较大、feature map 尺寸不规则的情况下TMA 带来的收益最明显因为不规则形状的地址计算开销在 A100 上是很可观的。2.2 WGMMAwarpgroup 级别的矩阵乘Hopper 另一个硬件级的大改动是 WGMMA全称 Warpgroup Matrix Multiply-Accumulate。A100 上的 MMA 指令是 warp 级别的一个 warp 32 个线程协作完成一个矩阵乘加。Hopper 把粒度提升到了 warpgroup也就是 4 个 warp 共 128 个线程一起做一个更大的矩阵运算。这个改动对 TensorRT 意味着什么简单说同样一个 GEMMHopper 上可以用更少的指令完成指令发射的开销降低了tensor core 的利用率上去了。但代价是你的数据布局得配合 WGMMA 的要求来排shared memory 里的数据摆放方式、swizzle 模式都有讲究。TensorRT 在构建 engine 的时候会自动做这些 layout 转换但如果你自己写 plugin 或者自定义 kernel这块就得自己处理。我遇到过一个问题自己写的一个 plugin 在 A100 上跑得好好的搬到 H100 上性能反而掉了排查半天发现是 shared memory 的 swizzle 模式和 WGMMA 要求的对不上导致 bank conflict 严重。后来改成 TensorRT 推荐的 layout 就正常了。这个坑后面在常见问题部分会详细说。2.3 FP8 与 Transformer Engine 的配合Hopper 的 tensor core 原生支持 FP8 格式这是这一代在精度和速度平衡上最大的卖点。TensorRT 从 8.6 版本开始对 FP8 的支持就比较完整了到了 9.x 版本基本是开箱即用。但 FP8 不是说你设个 flag 就能直接用的它涉及到 scale 的选取、校准数据的准备、以及哪些层适合降到 FP8 哪些层必须保留 FP16。Transformer Engine 是 NVIDIA 专门为 Transformer 类模型做的 FP8 训练和推理库TensorRT 在 Hopper 上会调用它的一些底层能力来做 FP8 的量化。实际部署中我一般建议对 attention 部分的 QKV 投影和 FFN 的两层线性做 FP8layer norm 和 softmax 保持 FP16这样精度损失基本在可接受范围内速度提升能有 1.5 到 1.8 倍。3. 在 H100/H200 上构建 TensorRT engine 的完整实操3.1 环境准备与版本选择先把环境搞对这一步看着简单但坑最多。TensorRT 的版本和 CUDA 版本、驱动版本之间有严格的对应关系Hopper 支持是从 CUDA 11.8 开始正式完整的但要用上 TMA 和 WGMMA 这些特性建议 CUDA 12.2 以上TensorRT 用 9.0 以上的版本。我目前在生产环境用的是这套组合驱动 535 以上CUDA 12.3TensorRT 9.2cuDNN 8.9。这个组合在 H100 和 H200 上都验证过稳定性没问题。安装方式我推荐用 tar 包手动装不要用 pip 装 TensorRT因为 pip 版本有时候会缺一些 Hopper 相关的 kernel 库而且版本管理比较混乱。# 下载 TensorRT tar 包后解压 tar -xzvf TensorRT-9.2.0.5.Linux.x86_64-gnu.cuda-12.3.tar.gz export TRT_PATH/path/to/TensorRT-9.2.0.5 export LD_LIBRARY_PATH$TRT_PATH/lib:$LD_LIBRARY_PATH export PATH$TRT_PATH/bin:$PATH # 验证安装 trtexec --version装完之后用trtexec跑一个简单的 benchmark 确认 Hopper 的 kernel 能正常调用。如果trtexec报找不到某些 so 库大概率是 LD_LIBRARY_PATH 没设对或者 CUDA 版本和 TensorRT 编译时的版本不匹配。3.2 ONNX 导出与精度配置从 PyTorch 导出 ONNX 这一步opset 版本建议用 17 或更高低版本 opset 有些算子导出后 TensorRT 解析会出问题。导出的时候 dynamic axes 要设对尤其是 batch 维度和序列长度维度这直接决定了后面 engine 能不能支持动态 shape。import torch import torch.onnx model.eval() dummy_input torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} } )导出完成后强烈建议用 onnxsim 做一遍简化把那些冗余的 reshape、transpose 合并掉。TensorRT 的 ONNX parser 虽然能处理这些但简化之后构建出来的 engine 结构更干净性能也会好一些。精度配置这块Hopper 上你有几个选择FP32、FP16、FP8、INT8。我的经验是视觉类模型 FP16 基本无损可以直接上Transformer 类模型如果对精度敏感先跑 FP16确认精度达标后再尝试 FP8INT8 需要校准集Hopper 上 INT8 的吞吐虽然高但精度损失需要仔细评估。3.3 构建 engine 时的关键参数用trtexec或者 Python API 构建 engine 的时候有几个参数直接决定了你能不能吃到 Hopper 的硬件红利。trtexec \ --onnxmodel.onnx \ --saveEnginemodel_h100.engine \ --fp16 \ --memPoolSizeworkspace:8192 \ --builderOptimizationLevel5 \ --hardwareCompatibilityLevelampere \ --useCudaGraph \ --verbose这里重点说几个参数。--builderOptimizationLevel5是最高优化级别TensorRT 会花更多时间搜索最优的 kernel 组合在 Hopper 上这个搜索空间比 A100 大很多因为多了 TMA 和 WGMMA 这些选项。--useCudaGraph在 Hopper 上收益很明显因为 Hopper 的 kernel launch 开销相对更高用 CUDA Graph 把多个 kernel 打包成一个图执行能省不少。还有一个容易被忽略的是--memPoolSizeHopper 的 shared memory 比 A100 大H100 每个 SM 有 228KB 的 shared memoryA100 只有 164KB。给 workspace 分配足够的内存TensorRT 才能用上更大的 tile size充分发挥 TMA 的搬运能力。我一般设 8GB 到 16GB具体看模型大小。3.4 验证 engine 是否用上了 Hopper 特性engine 构建完之后怎么确认它真的用上了 Hopper 的硬件特性最直接的办法是用trtexec --loadEnginemodel_h100.engine --dumpProfile看每一层的耗时然后对比 A100 上的 profile。如果某些层在 Hopper 上快得离谱那大概率是走了 TMA 或者 WGMMA 路径。另一个办法是用 Nsight Systems 抓 timeline看 kernel 名字里有没有tma、wgmma、fp8这些关键字。如果全是implicit_convolve或者sm80_xmma这种老 kernel 名字说明你的 engine 没有针对 Hopper 做优化可能需要检查 TensorRT 版本或者构建参数。4. 多卡部署与 kernel 层面的通信优化4.1 H100 千卡部署的通信瓶颈单卡优化做完scale 到多卡的时候问题就来了。H100 支持 NVLink 4.0单卡带宽 900GB/sH200 也是类似。但如果你做的是 tensor parallel 或者 pipeline parallel 的推理卡间通信的开销会吃掉相当一部分计算收益。我实测过一个 70B 的模型用 8 卡 tensor parallel通信占了总时间的 30% 左右。TensorRT 本身对多卡推理的支持是通过 TensorRT-LLM 这个上层框架来做的底层通信走的是 NCCL。NCCL 在 Hopper 上有一些针对性的优化比如利用 NVLink 的 multicast 能力做 all-gather但需要你在构建 engine 的时候把 parallel config 配对。4.2 通信 kernel 与计算 kernel 的 overlap多卡推理优化的核心思路是让通信和计算 overlap 起来。Hopper 的异步执行能力比 A100 强SM 可以在等待通信结果的同时继续执行其他计算任务。TensorRT-LLM 里有一个叫communication overlap的选项打开之后它会把 all-reduce 拆成多个 chunk每个 chunk 通信完就触发对应的计算而不是等整个 all-reduce 做完再算。这个优化在 H100 上效果特别明显因为 H100 的 NVLink 带宽足够高chunk 拆小之后通信时间短overlap 的窗口就大。我在一个 13B 模型上试过打开 overlap 之后端到端延迟降了 22%。4.3 硬件级 kernel 的调试方法当你怀疑某个 kernel 没有跑在预期的硬件单元上时有几个调试手段。第一是用cuobjdump把 engine 里的 cubin 反汇编出来看 SASS 代码里有没有HMMA、QMMA这些 Hopper 特有的指令。第二是用 Nsight Compute 抓单个 kernel 的详细指标重点看 tensor core 的利用率和 shared memory 的 bank conflict 情况。# 反汇编 engine 中的 kernel cuobjdump -sass model_h100.engine sass_dump.txt # 搜索 Hopper 特有指令 grep -E HMMA|QMMA|TMA|BULK sass_dump.txt如果 SASS 里全是FFMA这种标量浮点指令说明 tensor core 根本没被用上那性能肯定上不去。这种情况通常是精度配置或者 layer 融合出了问题需要回头检查构建参数。5. 常见问题与排查速查5.1 构建 engine 时报错 unsupported operator这是最常见的问题ONNX 里有些算子 TensorRT 的 parser 不认识。解决办法有几个一是用 TensorRT 的 plugin 机制自己实现一个二是用 ONNX GraphSurgeon 把不支持的算子替换成支持的组合三是升级 TensorRT 版本看新版本有没有支持。我遇到比较多的是LayerNormalization和Gelu这两个算子在老版本 TensorRT 上需要手动替换。TensorRT 9.x 之后基本都原生支持了所以升级版本往往是最省事的办法。5.2 FP8 精度掉得厉害FP8 的精度问题通常出在 scale 上。TensorRT 默认用的是 per-tensor scale对于激活值分布不均匀的层per-tensor scale 会导致量化误差很大。解决办法是改用 per-channel scale或者对敏感层保留 FP16。具体操作是在构建 engine 的时候提供一个 calibration cache里面记录了每一层的 scale。如果你没有校准数据可以用 TensorRT 的IInt8EntropyCalibrator2接口自己写一个校准器拿几百张代表性图片跑一遍就行。5.3 多卡推理时 NCCL 超时NCCL 超时一般有两个原因一是网络配置问题NVLink 或者 InfiniBand 的拓扑没配对二是某个 rank 的计算时间过长导致其他 rank 等太久。前者需要检查nvidia-smi topo -m的输出确认卡间连接是 NVLink 而不是 PCIe。后者需要看每个 rank 的 profile找出那个拖后腿的。我遇到过一次8 卡里有 1 张卡的温度偏高触发了降频导致那个 rank 的计算时间比其他卡长了 15%NCCL 就超时了。后来加了散热就好了。这种硬件层面的问题很容易被忽略但排查起来其实不难看nvidia-smi -q -d TEMPERATURE就行。5.4 常见问题速查表问题现象可能原因排查方法解决方案engine 构建失败ONNX 算子不支持看 trtexec 报错信息升级 TensorRT 或用 plugin精度掉点严重FP8 scale 不合理对比 FP16 和 FP8 输出改 per-channel scale多卡通信慢拓扑不是 NVLinknvidia-smi topo -m检查硬件连接kernel 性能差没用上 TMA/WGMMANsight Compute 看指令调整构建参数显存不够workspace 太小看构建日志加大 memPoolSize5.5 几个实操心得第一个心得是关于 builder optimization level 的。这个参数设成 5 的时候构建 engine 的时间会很长有时候要几十分钟。我的做法是先用 level 1 快速构建一个 engine 验证精度和功能确认没问题之后再花时间用 level 5 构建最终版本。第二个心得是关于 CUDA Graph 的。CUDA Graph 虽然能省 launch 开销但它要求输入输出的地址固定。如果你的推理服务需要动态 batch 或者动态 shapeCUDA Graph 就用不了。折中方案是用cudaGraphExecKernelNodeSetParams在运行时更新 kernel 参数但这需要你对 CUDA Graph 的 API 比较熟。第三个心得是关于 H200 的。H200 相比 H100 主要是显存大了、带宽高了计算单元其实是一样的。所以 H100 上的 TensorRT 优化经验在 H200 上基本通用不需要重新调。但 H200 的显存大了之后你可以把 batch size 开得更大这时候 TMA 的收益会更明显因为大 batch 下数据搬运的量更大TMA 的硬件加速效果更突出。6. 一些关于 kernel 选择的个人体会最后聊点偏经验的东西。TensorRT 在 Hopper 上生成的 kernel其实是一个自动搜索的结果它会根据你的模型结构、输入 shape、精度配置从一堆候选实现里挑一个它认为最快的。但这个“认为最快”是基于它的启发式规则不一定对你的具体场景最优。我一般会做一件事用trtexec的--timingCacheFile参数把每次构建的 timing cache 存下来然后对比不同构建参数下的 kernel 选择。有时候你会发现某个参数微调一下TensorRT 就选了一个完全不同的 kernel 实现性能差出 20% 都很正常。另外不要迷信 FP8。FP8 在 Hopper 上确实快但不是所有模型都适合。我试过一个检测模型FP8 之后 mAP 掉了 3 个点怎么调 scale 都救不回来最后只能退回 FP16。所以精度和速度的平衡还是得拿实际数据说话不能一概而论。Hopper 这一代的硬件复杂度比 A100 高了不少TensorRT 帮你屏蔽了很多底层细节但要想真正榨干性能还是得对底层的 kernel 和硬件单元有一定了解。希望这篇东西能给正在折腾 H100/H200 推理的兄弟一些参考少走点弯路。
返回列表