Transformer推理加速实战:从KV Cache到投机采样的10个优化技巧 前阵子帮某个团队调一个生成式模型的线上服务压测的时候发现吞吐上不去显存倒是快撑爆了。模型结构没问题数据链路也正常问题几乎全出在推理阶段的资源利用上。后来把几个优化手段组合着上单卡吞吐直接翻了一倍多。这里把我的实践记录整理一下聚焦在Transformer架构的推理侧优化。文中涉及的技巧有的是换了推理框架或后端才生效的有的是改模型结构或调整运行时配置但都可以在实际部署中独立或组合使用。先说清楚这里讨论的不是训练加速也不是分布式推理的宏大架构而是单机单卡或单机多卡场景下把Transformer模型的推理吞吐压上去的10个具体切入点。目标很明确在不牺牲太多精度的前提下让单位时间能处理更多请求或者让显存能容纳更大的batch。以下内容偏工程实践涉及原理的地方我会尽量讲透方便你判断自己的业务场景该怎么选。1. 先立一个优化主线KV Cache的命中率与显存账本吞吐上不去的首要瓶颈往往不是算力而是显存和访存带宽。Transformer推理时每个请求都会在解码阶段反复读写历史token的Key和Value矩阵这就是KV Cache。它的大小直接决定你能塞多大的batch也直接影响HBM显存带宽的消耗。1.1 搞清楚KV Cache到底吃了多少显存KV Cache的内存占用可以用很简单的公式估算2K和V两组 × 层数 × 头数 × 每头维度 × token数 × 精度字节数。举个例子13B模型、40层、32头、每头128维用FP16推理每个token大约占2 × 40 × 32 × 128 × 2 655,360字节也就是640KB。处理2048个上下文token时单请求的KV Cache就要占用1.25GB。如果并发32个请求单是KV Cache就是40GB算力再强也跑不动。数据来自公开的模型参数信息公式本身是通用的你可以直接用这个公式给你的模型记一笔账。1.2 把KV Cache的分配策略从预留改成按需多数推理框架包括部分开源的项目实现默认按最大序列长度一次性把KV Cache预留出来。这种静态分配的方式省了动态内存管理的麻烦但代价是显存浪费相当严重。尤其是在多数请求只有短上下文、短输出的场景里预留空间的使用率往往不足一半。我推荐的做法是采用PagedAttention式的按需分配这个思路和操作系统的分页策略一个道理你有一块连续的显存逻辑上用页表按块分配用到了再真正占用物理显存。实测下来在用某开源框架部署Chat类型模型时仅这一项就让显存碎片率降低了接近三成进而让可用batch容量增加了大约40%。如果你的框架不支持按需分配至少把max_seq_len参数往下调别让模型理论支持长度无节制地去占显存。1.3 针对长上下文的KV Cache量化这一步和后面的量化章节有交集但因为KV Cache本身就是显存大头单独拿出来说。当前很多框架支持对KV Cache做INT8或INT4量化直接把这块能耗大户压缩到原体积的1/2或1/4。量化KV Cache需要额外用心因为它直接影响attention计算的精度。实践中我建议先用INT8相对稳妥PPL困惑度损失一般在0.1以内很多模型体感上甚至无感INT4的收益更大但对激活值分布极度敏感个别层可能出现异常。可以按层混合量化让敏感层保留FP16非敏感层用INT4能把平均压缩比做到接近3倍。1.4 用好num_beams1和repetition_penalty的代价权衡这个是使用层面的优化但经常被忽略。贪心解码num_beams1和beam search比如num_beams4对吞吐的影响不是4倍而是更高的倍数因为beam search推理时KV Cache也要按beam数翻倍维护。不少业务场景其实不需要beam search把num_beams从4降到1能直接换取接近4倍的吞吐提升代价是部分长文本质量略降。我认为大多数生成类API服务beam search的收益并不足以抵消开销能关就关。2. 权重换精度架构不动让显存和带宽翻身权重数值的精度几乎是推理优化里性价比最高的一环。Transformer推理时每次前向计算都要读权重精度降一半读取的字节数就少一半带宽压力也直接减半。这种优化不改变任何网络结构却能同时改善显存占用和计算效率。我实测过一组对比同一部署环境下同一个7B模型在A100上分别用FP16、INT8和INT4精度加载当并发请求规模较大时INT8版本的吞吐约为FP16版本的1.6倍INT4版本约为FP16版本的2.2倍。这里参考的是团队内部压测数据不同框架、不同模型会有差异但趋势基本一致。2.1 量化感知训练QAT vs. 训练后量化PTQ谁更适合推理部署当前的开源模型大多已经过了量化适配用PTQ训练后量化就够了。它的核心是拿一小批校准数据通常几百条文本去统计激活值的分布然后确定缩放因子。优点是完全不碰训练流程一个下午就能搞定。QAT量化感知训练会在训练过程中模拟量化噪声让模型对低精度有适应性精度损失更小但需要重新走训练管线。对绝大多数推理侧优化场景来说PTQ已经能满足工程需求。只有在模型对INT8精度损失敏感或者需要压到INT4且精度告急时才值得去折腾QAT。2.2 GPTQ与AWQ两种主流PTQ方案的差异这两个都是针对Transformer权重量化的经典方案核心差异在如何降低量化误差GPTQ基于二阶信息Hessian矩阵逐层优化量化权重把量化误差分摊到后续权重上适合在离线阶段花较长时间把精度尽量压回来。AWQ是一种激活感知的量化方法基于少数关键通道对精度影响更大的观察在量化时按通道做缩放保护校准开销更小速度也更快。实操层面GPTQ常用在7B/13B级别模型的4bit量化上AWQ的INT4方案在部分场景下精度表现更稳。我自己的习惯是先用GPTQ快速出一个候选模型再用校准集做PPL对比如果PPL劣化明显就切AWQ重来。2.3 别盲目追求INT4INT8的兼容性和收益更均衡网上很多帖子一味强调INT4能把7B模型压到5GB以内显得很诱人。但如果你真的要上线服务我会劝你冷静。INT4推理时反量化操作本身有开销而且很多底层算子比如某些矩阵乘的优化实现对INT4的支持并不完整一旦遇到不支持算子的分支框架会把整段计算退回FP16收益直接打折扣。INT8则成熟得多无论CUDA核心还是各种推理框架支持的完整度都高很多。我可以给你一个实际参考部署一个7B模型FP16下吞吐大约每秒2000 token服务总吞吐切到INT8大约能到每秒3500 token而INT4虽然理论更高但因为反量化和部分layer退回FP16实际只比INT8高出10%-15%左右精度损失却明显得多。我认为对于语言模型的在线服务W8A8权重和激活都是INT8是性价比最高的方案兼顾了收益和稳定性。3. 动态batch与投机采样把GPU的空转时间堵上权重精度解决的是显存和带宽这一节解决的是算力和并行度。Transformer解码是逐token串行生成的每生成一个token就要做一次完整的前向计算。如果你每次只给GPU喂一个请求那算力利用率一定是惨不忍睹的。把多个请求攒起来一起推理就是dynamic batching动态批处理的基本盘。3.1 用框架内置的Continuous Batching替代手工攒batch手工累攒batch的老方案是等一批请求凑够数量再统一推理响应延迟高、用户体验差。现在像vLLM、TensorRT-LLM这类框架都已经支持continuous batching连续批处理或叫in-flight batching思路是请求不用等同一批次的其他请求谁先decode完谁先离开新请求可以随时插入空位框架在迭代级别动态调度。这意味着服务一直处于有活干的状态GPU的空转很少。实测中从静态batching切到continuous batching吞吐提升常常是翻倍的并发越低效果越明显因为资源利用率提升幅度最大。3.2 投机采样Speculative Decoding用小模型帮大模型打草稿这是一个让我印象深刻的优化手段应该说效果超过我的预期。核心思路是让一个很小的草稿模型draft model快速生成几个候选token然后大模型并行验证如果验证通过就一并接受如果失败就丢弃重来。这样原本要串行跑4次的大模型前向变成了1次大模型前向几次小模型前向整体延迟和吞吐都有显著改善。这个方案有个前提草稿模型和大模型的tokenizer必须一致否则token映射会乱另外草稿模型质量不能太差否则接受率太低反而白白多耗时间。我在一个场景里用1B草稿模型配合13B主模型接受率大约在0.7左右整体延迟下降了约40%吞吐提升了约70%。数据受场景影响较大但这个优化方向很值得尝试。3.3 解码长度的上限管理别让长输出拖垮整个batch有一个destream的细节很多系统上线前不会考虑就是单个请求的输出上限。如果不限制max_new_tokens一个用户的超长请求会一直占据batch中的位置拖慢所有其他请求离开的速度吞吐直接下降。我见过不少线上事故是某用户疯狂触发长文本生成整个服务的响应延迟就崩了。建议在路由层就给每个请求设置输出上限比如限制512或1024 token超长请求走异步队列或拆分处理。这个优化看着没什么技术含量但对稳定性的帮助非常大。4. 推理框架选型与运行时配置把改好的方案落到实处如果模型结构和量化也调完了batch策略也对了吞吐还是不够往往问题出在框架层的算子性能。同样一个模型在不同推理框架下的吞吐差距可能达到30%-80%这不是玄学而是算子实现、显存管理、调度策略的综合差距。4.1 主流推理框架的取舍逻辑我目前常用的几个框架简单排个序供你参考特性框架A框架B框架C吞吐表现高中高中显存利用率高动态KV较高手动配置多中量化支持INT4/INT8成熟INT8为主支持但算子覆盖一般自定义算子需要改C/CUDAPython算子扩展友好PythonCUDA适合场景高并发线上服务科研/灵活部署原型快速验证这里的框架A/B/C就是实际常见的开源项目你可以按需求直接去对标。我的选择经验是如果只是跑一个demo无所谓如果要做高并发的API服务就直接上框架A这一类自带paged memory和continuous batching的能省掉大量自己造的轮子。4.2 torch.compile与图优化动态图和静态图的差距如果你的服务还在用PyTorch原生动态图跑推理那么不管前面做了什么优化吞吐上限都会被动态图的调度开销卡住。torch.compile在PyTorch 2.x里已经非常成熟可以通过modereduce-overhead把模型中的算子融合成更少、更大的kernel减少kernel launch次数。但要注意torch.compile并不是推理服务的最终形态它生成的图还是可能包含很多小算子。如果你追求极致性能应该考虑用TensorRT或者框架B这类编译型后端做全静态图优化。以我观察从动态图切到静态图优化7B模型的单batch latency能减少30%左右这在实时性要求高的场景里还是值得做的。4.3 CUDA Graph捕获小算子爆炸的最直接手段Transformer模型里的算子数量动辄几百上千每个小算子的launch开销积累起来非常可观。CUDA Graph的思路是把一串kernel的启动过程一次性捕获成一个graph之后每次只需提交一次graph就能按顺序执行所有算子。很多框架已经自动集成了CUDA Graph比如框架A的enforce_eagerFalse模式就是默认去掉eager执行走CUDA Graph但如果你是自己写的推理脚本可以把解码循环内的前向计算包在一个graph里实测对短序列场景的延迟改善很显著。有个坑需要注意CUDA Graph会锁定显存池如果你后续需要动态分配显存比如新请求加入batch会把graph打碎或重建所以它更适合固定batch或配合continuous batching的场景使用。5. 算子级优化不要小看单个矩阵乘的力量batch和框架的优化本质上都是把资源用满但要真正把速度推上去离不开算子本身的优化。Transformer推理时最耗时的部分通常是线性层如QKV投影、FFN的中间层这类算子性能几乎决定一切。5.1 核心算子的融合GELU、LayerNorm与矩阵乘的合并一个典型的操作链是矩阵乘 - bias add - GELU - 下一个线性层。如果你逐个函数去调用cudaMemcpy中间结果和kernel launch的开销会浪费大量时间。正确的做法是写融合算子在同一个kernel里完成矩阵乘、Bias、激活函数甚至跨过了激活后直接再乘下一层权重。这就是FlashAttention系列同样采用的少读少写哲学——减少显存往返比降低计算量更有效。很多框架已经内置了这些融合逻辑但如果你是在自己定制模型结构这里值得重点检查。用Nsight Compute或PyTorch Profiler看一遍时间占比线性层和相关融合机会往往是最先暴露的问题。5.2 FlashAttention与Attention变体FlashAttention的核心不是数学上的改变而是通过分块tiling把注意力计算拆成小块避免在显存中实例化完整的S矩阵即QK^T的结果。这个矩阵在长序列场景下会非常巨大FlashAttention通过在线softmax等技术把中间量控制在片上内存SRAM中速度更快、显存占用更小。如果你的模型仍然使用标准的torch.nn.MultiheadAttention实现那第一件事就是换成支持FlashAttention的内核许多框架已经通过attn_implementationflash_attention_2选项提供。这个改动对长上下文推理的帮助几乎是立竿见影的。5.3 结尾的非关键层selective layer skipping这是一个带偏门性质但实际可以有稳定收益的技巧有些模型尾部层对输出概率分布的影响非常小推理时可以直接跳过。比如某些13B模型跳过最后2个decoder layerPPL只涨不到0.5%但延迟能省2%左右。当模型被量化和batch优化之后这个2%也会显得有价值。这个思路不适合无脑尝试更合理的做法是使用校准集逐层评估模型输出变化找出对最终分布影响最小的若干层然后选择性地跳过它们。这也是架构优化里少数真正动模型结构的做法。6. 小模型替代与蒸馏如果2倍还不够那就换思路前面所有优化都在同一架构内做文章如果2倍吞吐目标在这些手段的组合下还没达成就需要考虑从模型本身下功夫了。这不属于传统意义的推理优化但对线上系统的吞吐影响是决定性的一击。6.1 用更小的模型做同一件事模型的参数量几乎和推理速度成正比。7B和13B模型的单次前向延迟比大约就是7:13在实际算力达到瓶颈的场景更接近。如果业务逻辑允许一定程度的输出去风格化或简化将13B模型替换为7B或更小的模型吞吐提升是倍数级的远高于任何技巧性优化。现在不少团队的实际做法是复杂场景用大模型简单场景用小模型按难度路由分配合适的模型容量。6.2 蒸馏用小模型逼近大模型的效果有句话说得好蒸馏好的小模型效果不一定输给原来的大模型。你完全可以拿一个大模型当老师在大量业务数据上让小模型学习它的输出分布。这个做法在不少中等难度任务上可以做到7B水平的延迟、13B水平的质量这是我在实践中觉得性价比高于任何量化技巧的方案。关键是数据配比和蒸馏温度设置温度不要太高否则会损失真实分布的尖锐度。6.3 分层退出Early Exit让简单样本走捷径又一个偏向长度换时间的思路。给模型加上分类器分支当某层已经足够确信输出时提前结束并返回只有难样本才走完整个模型。这样做的好处是简单样本的平均层数更低服务端的吞吐自然更高。这个方案落地时有个障碍很多模型架构没有天然的Early Exit分支需要训练时附加辅助loss。如果拿现成的开源模型直接改效果会比较勉强。我的建议是先看自己的业务请求分布如果大部分请求都属于简单问题这个思路值得一试如果所有请求都很复杂就跳过它。7. 排查链路吞吐上不去时按什么顺序定位根因把10个技巧过完还是要面对现实不是每个技巧都立即可用。我结合自己的排查经验按优先级列一个定位清单供你参考。7.1 第一步先看显存和batch容量如果你的并发稍微一高就OOM显存溢出那么所有计算优化都得往后放。先按第1章的方式计算出KV Cache的占用再决定是否量化KV Cache或降低序列长度。效能的排序是显存容量 - batch大小 - 吞吐。显存不够后面全是空谈。7.2 第二步看time分布占比利用框架自带的profiler或Nsight Systems抓一段推理过程把时间按算子类别统计。你会发现50%以上的时间往往集中在少数几个kernel上。如果这些kernel中间有大量内存拷贝或小kernel启动那说明融合不够这是优化潜力最大的地方。7.3 第三步排查框架配置的隐形开关很多团队用框架A部署后吞吐上不去其实是因为默认配置没有开启某些优势特性。比如enable_prefix_caching、enable_chunked_prefill这些开关对长上下文或高并发请求有显著影响gpu_memory_utilization设太低时KV Cache容量不足动态batch就会提前打满。把这个值从0.85提到0.95吞吐常常能上升10%-20%。7.4 第四步用小规模实验验证假设无论你从理论推断哪个技巧更有效我都建议先在单请求低并发下跑一遍先看latency再看并发升高时的throughput曲线。这样做的好处是能快速分离出是计算瓶颈还是调度瓶颈。例如在并发1时延迟很高那问题大概率在算子效率并发1时延迟正常但并发16时吞吐上不去那问题大概率在显存或batch策略。8. 综合部署案例参考组合10个技巧的一个可达方案前面分章节讲了技巧这里给一个相对完整的组合部署示例方便你直接参考操作顺序。8.1 目标与初始准备假设目标把一个13B模型在单张A10080GB上部署希望吞吐从初始状态的每秒约1800token提升到3600token以上同时保证PPL困惑度不劣化超过0.5。初始状态FP16精度无量化无动态batch框架为原始PyTorch推理脚本。8.2 分阶段优化过程先把推理后端从原始PyTorch脚本换成框架A支持PagedAttention、continuous batching、CUDA Graph。预期吞吐约3500token/s。这一步就能把基础目标完成大半。开启INT8权重量化W8A8保持显存余量让batch容量变大。预期吞吐约4500token/s。此时已超过2倍目标。开启KV Cache INT8量化配合gpu_memory_utilization0.92。预期吞吐约5000-5500token/s。有余力还能开投机采样进一步提升。对输出层做选择性跳过或蒸馏一小段作为最后的补充。需要说明这个数值是我在一个特定业务请求分布下测出来的换场景会有波动。但顺序的重要性和提升幅度的量级是可以复用的。8.3 实际应用时的注意事项每次改动后都要回到校准集上做PPL或下游任务评测不能只看吞吐数字涨了多少就上线。对API服务来说延迟的尾部P99和吞吐同样重要动态batch和连续批处理在某些极端情况下会拉大延迟抖动需要额外监控。上线时建议做A/B对比同一批请求分别走老配置和新配置对比返回质量再决定是否全量切流。9. 我的实践心路哪些技巧被高估哪些被低估了优化做过一轮之后总会有一些判断修正。这里分享几个比较强烈的个人感受。被高估的INT4量化。并不是说它不好而是实际收益并没有宣传中那么夸张。反量化的开销、算子支持不完整、个别层精度崩塌这些都会在工程落地时突然冒出来。除非你的唯一目标是塞进显存否则INT4带来的麻烦往往比收益多。被低估的连续批处理continuous batching。很多人把它和普通batch混为一谈但实际效果差距很大。静态batch是攒够一桌开饭continuous batching是流水线自助餐GPU的利用率完全是两种量级。如果你只能做一件事我建议优先做这个。被低估的投机采样。初看总觉得多跑一个小模型会更慢实际落地后效果真香。它本质上是把算力空闲和延迟敏感这两个矛盾点用一个小模型给调和了是我近期最推荐的推理加速手段。被低估的CUDA Graph。它不需要改模型结构也不需要额外训练只要框架支持配置两个参数就能见效。对延迟优化来说它比很多花哨的算法都更实在。10. 下一步还能怎么压榨多卡和更深的系统优化如果单卡的2倍吞吐已经满足需求那恭喜你这套组合拳可以收工了。如果还不够还有几条路可以继续走一是把多卡并行从张量并行Tensor Parallelism切到更细粒度的流水线并行让更多计算并行化。二是做prefix caching当多个请求共享同一段前缀比如系统提示词或长文档时缓存住这部分计算结果能省掉大量重复计算。三是在架构层面引入early exit或更激进的分层剪枝对特定领域模型效果突出。但我要提醒一句优化的边际收益是递减的。当单卡吞吐已经到三四倍之后再想翻倍通常意味着重新设计服务架构、多卡调度和数据分片这些工作的复杂度和维护成本会一下子高很多。对多数业务来说把2-4倍提升稳定吃下来已经比很多团队跑得更好了。