不换硬件,10个架构优化让Transformer推理吞吐翻倍
一股脑堆GPU是把Transformer推理做得又慢又贵的常见路子。真上了线才发现显存翻倍、带宽翻倍吞吐却只涨了几成账户余额反而先见底。不少团队在优化推理时总把目光放在加机器、换显卡上忽略了模型架构本身和运行时调度里的大量可挖空间。这篇东西想聊的就是怎么不换硬件靠10个架构优化技巧把吞吐量往上顶目标2倍实测能到多少、瓶颈卡在哪我都会写清楚。适合准备做模型部署、推理服务性能调优的朋友尤其是大模型应用团队的算法同学和推理引擎开发者。1. 先认清瓶颈算力、访存还是调度1.1 解码阶段为什么这么慢Transformer的生成阶段是自回归式的每生成一个token都要把历史token的键值缓存也就是KV Cache和模型权重重新读一遍。这里的关键变量不是浮点算力而是显存带宽。假设某个70亿参数的模型用FP16存储光权重就是大约14GB而一张主流数据中心卡显存带宽大概在2TB/s量级理论上一轮前向就要7毫秒。再加上KV Cache的读取实际延迟只会更高。算力反而常常闲着GPU的利用率很低。所以推理性能优化的第一性原则挺明确减少每一步读取的数据量或者让读取出来的数据被更多token复用。这也是后面这些技巧的共同出发点。1.2 优化前的基线测量没有基线就谈优化等于闭眼开车。我会先跑一个固定输入长度的压测统计三个指标首token延迟、每个token的增量延迟、整机吞吐量tokens/s。顺便记录GPU利用率和显存占用。操作上要注意压测时不能只看平均延迟很容易被长尾拖累。要分位数p50、p99都得记。实测中发现很多模型在并发上来之后自回归延迟并没有变慢但吞吐就是上不去这时候多半是批处理调度、padding策略或者显存交换出了问题跟内核计算关系不大。做基线时最好把batch从1、4、8、16、32逐档测一遍画出吞吐随batch变化的曲线后续所有优化效果都能在同一张图上对比。2. 计算内核优化把GPU喂饱2.1 技巧一算子融合GPU执行一次内核要从显存里搬数据搬一次数据比算一次贵得多。原先一个算子接一个算子中间结果要写回显存下一个算子再读出来白白浪费带宽。算子融合就是把这几个读写过程合并成一个内核让中间结果留在寄存器或共享内存里。举一个常见例子把LayerNorm和后面的线性层融合把残差相加和下一个归一化融合。融合之后同样的计算任务内核启动次数减少显存读写总量也明显下降。不少开源引擎和编译工具已经把这套做成了自动优化但如果你是在自研推理框架里干活手动融合几个高频算子往往能先吃下一波收益。2.2 技巧二分块注意力FlashAttention类实现标准注意力里有一个很重的中间结果Q乘K转置之后得到N乘N的注意力矩阵这个矩阵会完整写入显存再做softmax然后乘V。序列一长显存直接爆掉带宽也被大量消耗。分块注意力把N乘N的矩阵拆成小块在线计算softmax的统计量全程不落回显存。空间复杂度从O(N^2)降到O(N)访存量大幅下降。长序列场景收益尤其明显但短序列未必赚因为分块和重算本身也有开销。我的做法是序列长度超过512甚至1024时才打开这个优化短序列维持普通注意力反而更快。2.3 技巧三CUDA Graph与编译级优化GPU上跑模型不只是算还得不断由CPU下发内核启动命令。小模型、短序列场景下CPU启动开销甚至能占到单次迭代时间的一半。CUDA Graph把一串内核启动序列捕获成一张图之后一次性重放大幅降低CPU的调度负担。使用时要留意捕获阶段要求模型shape固定动态shape场景会踩坑。实操时可以把batch维度或序列长度维度固定在某个最大值内部再用mask或padding处理。你要是用某些支持运行时编译的框架可能在底层已经自动应用了类似机制不需要手工开。但自己写推理脚本时千万别忘了查一下内核启动占比。3. 模型结构改造让模型更轻、更省缓存3.1 技巧四压缩KV Cache自回归解码必须保存历史token的K和V。KV Cache的显存占用约等于batch size乘以序列长度乘以层数乘以每个头的维度再乘以数值精度占用的字节数。算一笔账batch32、上下文长度2048、层数32、隐层维度4096、FP16理论就需要32×2048×32×4096×2字节大约32GB。这数字比很多模型的权重还大。压缩方案通常分几类一是分页式管理只在真正需要时分配显存不按最大长度整体预留二是把K和V量化到INT8精度损失通常可控三是用滑动窗口或全局token混合策略只保留最近一段历史极少场景会影响质量。这里不一定要全用先算算自己场景里KV Cache占了多大再决定从哪个维度下手。3.2 技巧五量化压缩FP16/INT8/INT4模型在推理时主要瓶颈是访存降低权重精度等于直接减少每次要搬的数据量。FP16权重一个数占2字节INT8是1字节INT4可以再到0.5字节。以70亿参数模型为例FP16权重约14GBINT8约7GBINT4约3.5GB。省下来的显存可以喂给更大的批量数据搬移量也同步下降吞吐自然往上走。常见路线是先做PTQ训练后量化用校准集统计激活范围选择per-tensor或per-channel的量化粒度如果精度掉得厉害再考虑QAT或对某些敏感层做混合精度。实操上我会先把容易量化的线性层和注意力层全部转INT8观察输出质量质量没问题再试探INT4。切记量化之前先保存一份原模型方便随时对比。3.3 技巧六剪枝与稀疏化剪枝不是阉割能力而是去掉冗余参数。非结构化剪枝把单个权重置零但GPU目前对稀疏矩阵的计算加速并不普及实际收益很有限还容易带来内存碎片。真正对推理有用的是结构化剪枝判断出哪些注意力头或者FFN中间神经元贡献很小整组删掉再微调恢复精度。以多头注意力为例可以给每个头算一个重要性分数比如看它在验证集上对输出梯度的贡献或者激活值幅度排名靠后的头直接移除。常见模型剪掉10%~20%的头质量下降幅度往往小于显存和算力的下降幅度。注意剪枝后通常需要短期微调否则输出质量会明显波动。如果你没有微调条件这个方案慎用。3.4 技巧七MQA/GQA注意力变体标准多头注意力里每个头都各自存一份K和V。假设32个头KV Cache体量就是单份的32倍。要是让所有头共享同一组K、V就变成MQAKV Cache大幅缩小分成若干组共享则是GQA是卡在质量和速度之间的折中。对已经训好的模型把注意力结构改成GQA并不容易基本要在预训练阶段就定下来。但如果你是在微调一个底座本来就带GQA的模型那恭喜你KV Cache和访存量天然就低一截优化负担小很多。选型时建议优先挑GQA结构的开源底座别自己硬改。4. 系统级优化调度、并行与投机4.1 技巧八动态批处理文本推理和图片分类不同每个请求的输入输出长度差得很大。如果用固定batch等所有请求都结束再一起归还慢请求会一直占着算力吞吐被拖垮。动态批处理连续批处理的思路是把调度粒度从“整批请求”降到“单个迭代”每轮迭代从等待队列里挑最该执行的请求完成一轮就把位置让给新请求。GPU利用率因此被填满尾部延迟也会好很多。很多成熟推理引擎已经把动态批处理做成默认能力但如果你在自研服务就得自己维护一个迭代级调度器。常见实现是按当前迭代所有序列的最长长度做padding尽量把长度相近的请求聚到同一批减少无效计算。4.2 技巧九张量并行与批量并行单卡放不下模型或者单卡带宽不够时可以把权重按层切到多张卡上。张量并行把注意力和MLP的权重矩阵切成多份各卡各算一部分算完再通过通信规约汇总。它的特点是显存压力分散但通信开销会上来所以对计算量大的模型收益更明显。如果你的卡也不少可以反过来做数据并行每张卡放完整模型各跑不同请求吞吐线性扩展。实测中小模型用数据并行往往比张量并行更省心因为通信占比太大。大模型到单卡装不下的程度才值得上张量并行。还有一种组合方式先张量并行降低单卡显存再数据并行提升整机吞吐这也是很多推理服务的主流架构。4.3 技巧十投机解码自回归是串行的一次只能生成一个token这是Transformer推理低效的根源之一。投机解码想了个办法先让一个小草稿模型快速生成一串候选token再让大模型用一次前向同时验证这些token。如果都能通过就相当于一次迭代推进了好几步串行迭代数大幅下降。实际加速取决于草稿模型和主模型的“默契程度”太小的草稿模型预测不准接受率低太大会拖慢草稿阶段本身。按我的经验草稿模型参数取主模型的几十分之一到十分之一之间接受率能达到0.6甚至更高的时候才划算。这个技巧属于工程上稍微复杂但收益很直接的一类值得排进最终优化清单。5. 组合优化实战如何把2倍吞吐落地5.1 一条推荐的优化路线10个技巧全上不一定最优硬件和模型结构不同收益排序也不同。我按通用场景排过一条路线先开FlashAttention和量化把访存压力打下来接着上动态批处理和KV Cache优化把显存利用率和调度效率提上去最后叠加CUDA Graph和投机解码榨干CPU启动和迭代串行的水分。举个实际案例某个约70亿参数的模型原始FP16、固定batch8吞吐在800 tokens/s左右量化到INT8之后batch能开到16配合动态批处理和FlashAttention稳定跑到1400到1600再开CUDA Graph和投机解码最终能摸到1700到2000。整个过程没有换过显卡显存占用反而降了。硬件配置不变的情况下2倍并不夸张。5.2 实测中的几个“坑”第一个坑是盲目加大batch导致显存OOM。加batch之前先用公式算一遍权重加KV Cache的显存占用量留出至少20%余量。第二个坑是FlashAttention在短序列下反而更慢别迷信新算法要用profiler看实际耗时。第三个坑是量化后生成质量突然崩校准集太窄或者某些层特别敏感这时需要扩展校准数据或者改用per-channel、混合精度。第四个坑是CUDA Graph遇上动态shape直接报错可以先固定最大长度或做长度分桶。这些坑我基本都踩过写下来就是希望大家别再交学费。5.3 常见问题与排查清单这里直接列一个速查表遇到问题先对着查。问题现象可能原因排查思路吞吐上不去但GPU利用率低动态批处理未生效或padding浪费看调度器是否迭代级切批统计padding占比量化后输出质量下跌明显校准集不具代表性/敏感层全量量化扩大校准集敏感层回退FP16首token延迟高增量延迟低batch过大或输入padding过多检查批内序列长度分布按长度分桶长序列时显存暴涨KV Cache没有分页管理开启分页式KV管理必要时量化KV多卡扩展后吞吐不增反降张量并行通信占比过高减小切分粒度或改用数据并行CUDA Graph捕获失败模型中有动态shape或循环依赖CPU逻辑固定最大长度用mask规避动态分支排查问题要有顺序先看显存和带宽再看调度最后才怀疑内核。很多“优化无效”的结论其实都是运行方式变了但测量方法没跟上。6. 最后想说点实操体会这几轮优化做下来我最大的体会是不管网上给出的技巧多炫都得回到自己的模型和部署环境里做对照实验。比如同一个量化策略对某个模型收益极好换一个模型可能精度掉到没法忍同一个投机解码某一类文本接受率很高换成长篇技术文档可能收益就缩水。所以建议把上面的优化技巧当成一个工具箱按“量化 → 访存 → 调度 → 编译 → 解码策略”的顺序逐个尝试每加一项就重新跑一遍基线用数据说话。另外KV Cache量化值得单独试试因为它几乎不影响延迟但对显存压力影响很大。你要是服务已经上线还可以把10个技巧按收益和风险排个优先级先做低风险高收益的再慢慢挑战需要验证的项目。优化做完之后别忘了把最优配置固化到部署脚本里这些经验回头能给团队省不少时间。