Cohere推出硬件感知的动态推测解码:推理速度翻倍


当你的显卡闲着没事干的时候,它不是在休息,它是在等我把下一批数据喂给它。

大语言模型生成文本是一个字一个字往外蹦的,这种蜗牛速度让程序员们头疼了很久。
后来有人想了个办法,弄个小模型先一口气写好几个字,再让大模型去检查这些字对不对。只要小模型猜得准,大模型一次就能确认好几个字,速度一下子就上去了。这个办法叫推测解码。

听起来完美无缺对吧?但问题在于,小模型猜字的数量如果固定不变,在某些情况下反而会让速度变得更慢。生产环境里的请求数量忽多忽少,显卡的工作状态也跟着变来变去,固定的猜字数量就像一双不合脚的鞋,有时候跑得快,有时候磨脚。

Cohere公司的工程师们搞出了一套动态推测解码技术,让猜字数量跟着硬件状态实时调整,这才算把这个问题真正解决了。

显卡干活也有两种完全不同的状态

你去观察显卡干活的样子,会发现它有两种截然不同的状态。一种状态是闲得发慌,显卡里面的计算单元大部分时间都在空转,因为数据从显存搬到缓存的速度太慢了,计算单元只能干等着。这时候显卡的主要瓶颈在于搬运数据,搬运工忙不过来,流水线上的工人全部停工。另一种状态是忙得要死,数据搬运的速度足够快,计算单元一刻不停地做矩阵乘法,显卡的全部火力都打开了。

这两种状态之间的切换取决于一个东西,叫做批次大小。想象你开了一家奶茶店,如果客人不多,你大部分时间在等收银员下单,做奶茶的机器经常空着,这就叫内存带宽瓶颈。如果客人排长队,收银员下单飞快,做奶茶的机器一刻不停,这就叫计算瓶颈。显卡干活也是一样的道理,批次小的时候显卡在等数据,批次大的时候显卡算到冒烟。

推测解码能提速的本质,就是利用了显卡在内存带宽瓶颈时期闲置的计算能力。你让大模型一次验证好几个小模型写的字,等于让本来闲着没事干的计算单元忙起来。但你得注意一个致命问题,大模型每验证一个批次要处理的字数是批次大小乘以猜字数量加一。如果猜字数量固定为三个,大模型一次就要处理四倍的量。在批次小的时候这没问题,反正计算单元闲着也是闲着。但在批次大的时候,大模型本来就已经算不过来了,你还让它多算四倍,这不是帮忙,这是在帮倒忙。

固定猜字数量在生产环境中经常翻车

你把推测解码部署到真正的生产系统里试试看,马上就会撞上一堵南墙。生产环境里的批次大小很少固定不变,它会随着用户请求的数量忽高忽低。白天人多批次大,晚上人少批次小,搞促销活动的时候突然涌入大量请求,批次直接爆炸。固定猜字数量在这种场景下面根本没法用,因为你要么照顾小批次让大批次变慢,要么照顾大批次让小幅度的提速完全浪费掉。

强化学习训练的场景更是雪上加霜。强化学习分训练阶段和生成阶段,生成阶段要产出大量数据用来训练模型。这个生成阶段能吃掉整个训练流程百分之八十五的计算资源,是个不折不扣的资源黑洞。更气人的是,推理模型在生成的时候会出现长尾分布,意思是一个批次里面可能有一个请求特别话痨,一直在那里疯狂输出,把整个批次的其他请求都堵在那里干等。

推测解码本来特别适合处理这种长尾生成,因为长尾请求意味着批次有效变小了,正好是推测解码发挥威力的时候。但问题在于,你没法提前知道一个请求会不会变成话痨,所以系统只能按照最坏情况来配置。如果你为了照顾少数长尾请求把猜字数量调高,正常批次下就会翻车。如果你为了正常批次把猜字数量调低,长尾请求的提速又拿不到。这个两难困境让固定猜字数量的推测解码在强化学习场景里特别鸡肋。

动态猜字数量像给显卡装了个自动变速箱

动态推测解码解决这个问题的思路特别简单粗暴,每时每刻根据硬件状态重新算一遍最优猜字数量。硬件状态说白了就是当前批次大小和显卡的内存带宽利用率,这两样东西就能决定你是处在搬运瓶颈还是计算瓶颈。处在搬运瓶颈就多猜几个字把空闲计算资源用起来,处在计算瓶颈就少猜几个字别给大模型添乱。

对于稠密模型来说,猜字数量和批次大小之间的关系是一条单调递减的曲线。批次越小猜越多,批次越大猜越少,这个规律非常直观。但对于混合专家模型来说,这个关系就变得好玩了,它不是单调的,而是一条山谷形状的曲线。批次很小的时候最优猜字数量反而低,因为验证阶段会额外激活一些专家模块,这些专家的权重要从显存搬进来,额外增加了搬运开销。批次稍微大一点的时候最优猜字数量反而变高了,因为大多数专家已经被激活了,验证额外猜的字不再增加多少搬运开销。批次再大的时候又变回和稠密模型一样,计算成为瓶颈所以猜字数量又降下来。

这个现象背后藏着一个反直觉的道理,在混合专家模型里,验证第一个猜的字反而是代价最高的,因为要加载额外的专家权重。后面再猜的字反而变得便宜了,因为专家已经在显存里了。这和稠密模型的逻辑完全相反,稠密模型里每个猜的字代价都差不多。这说明动态推测解码在不同模型架构上的表现差异大得惊人,你不能用一个套路去打天下。

用一个简单的公式找到最优猜字数量

那么问题来了,怎么找到当前硬件状态下的最优猜字数量呢?你去翻学术论文会发现一大堆复杂到让人头晕的公式和算法。但Cohere的工程师们把这个问题简化成了一个小学算术级别的指标,叫好效益。好效益等于接受长度除以令牌间延迟。接受长度是小模型猜的字被大模型接受的数量,令牌间延迟是生成一个完整批次的总耗时。

这个公式的美妙之处在于它直接反映了用户体验。接受长度越长用户每次刷新看到的新字越多,令牌间延迟越短用户等得越着急。两个东西相除得到的结果越大说明速度越快。你要做的就是不断尝试不同的猜字数量,分别算出它们的好效益,挑那个数值最大的用就行了。

工程师们做了一件特别聪明的事情,他们提前在离线环境里跑一遍各种批次大小和猜字数量的组合,把最优值做成一张查表。运行时直接查表就行,不需要在线计算,彻底解决了系统启动时不知道怎么选猜字数量的问题。而且这张表可以随时更新,如果系统运行时发现实际表现和离线测试有偏差,动态调整表里的数值就行。

稠密模型实测结果证明动态方案全面碾压

工程师们用MT-Bench数据集做了实测,对比了三套方案。第一套是不做推测解码的原始方案,第二套是猜字数量固定为三的传统推测解码,第三套是动态猜字数量的新方案。测试对象是Cohere自家的Command A稠密模型和Command A+混合专家模型。

在稠密模型上,动态方案的表现堪称完美碾压。批次小的时候动态方案和传统推测解码速度一样快,都拿到了推测解码应有的提速。批次中等的时候动态方案开始展现优势,比传统方案快一截。批次大的时候传统方案已经开始翻车了,速度比原始方案还慢,但动态方案依然保持优势。到了批次特别大的时候传统方案彻底崩盘,动态方案回落到和原始方案差不多的速度,至少没变得更慢。

具体数据更有说服力,在批次大小为一百二十八和二百五十六的时候,动态方案比传统方案快大约百分之二十三。这还不是最厉害的,在批次一百二十八的时候动态方案比原始方案快百分之七点五,在批次二百五十六的时候快百分之一点八。传统方案在这两个批次下都比原始方案慢。动态方案的精髓就在这里,能提速的时候使劲提,没提速空间的时候至少不拖后腿,这是一个永远不会让你亏钱的策略。

混合专家模型的结果揭示了一个隐藏问题

混合专家模型Command A+的测试结果就有意思了。动态方案选出来的最优猜字数量在大部分批次下都是三,刚好和传统方案的固定值一样,所以两者速度差不多。只有在批次十六到三十二之间动态方案选了五,但这个更高的猜字数量没有转化成更快的速度。

工程师们分析后发现了一个藏在表象下面的问题。他们的猜字模型是EAGLE架构,这个架构在训练的时候只针对猜一步进行优化,但在实际使用中要连续猜好多步。当猜字数量从三增加到五的时候,后面多猜的那几个字准确率不够高,大模型验证的时候经常把它们拒掉,所以整体速度没提上去。这说明动态猜字数量虽然找到了理论上最优的数值,但这个数值受限于猜字模型本身的能力没法兑现。

工程师们很坦诚地指出了这个问题,并预测更新的猜字模型比如EAGLE-3或者DFlash在这种场景下表现会更好。这种实事求是的态度在技术博客里相当罕见,大部分公司只会展示自己牛逼的地方,不会主动暴露软肋。

在vLLM框架里落地时碰到两个硬骨头

把动态推测解码塞进vLLM这个高度优化的推理框架里,才是整个项目最难啃的骨头。vLLM做了大量底层优化,动态猜字数量几乎跟每一处优化都有冲突,需要逐个解决兼容性问题。

第一个硬骨头是异步调度。vLLM为了不让CPU调度拖慢GPU计算,搞了一套异步机制。CPU在给当前批次做调度的时候,GPU同时在跑上一个批次的计算。调度器跑在CPU上,计算跑在GPU上,两边同时干活互不等待。但调度器需要提前知道下一个批次要猜多少个字,才能提前准备好内存空间。如果猜字数量每个批次都在变,调度器和计算引擎之间的信息同步就变得特别复杂。

工程师们重构了调度器和计算引擎之间的通信协议。调度器提前告诉计算引擎本轮要猜几个字,计算引擎跑完本轮后告诉调度器实际验证了几个字。两边之间传递的都是轻量级的数字,不传递具体的内容,这样异步调度依然能正常工作,不会被动态猜字数量打断。

第二个硬骨头是全CUDA图。CUDA图把一大堆内核启动合并成一次启动,省掉了大量的CPU开销。vLLM对解码请求用了全CUDA图,每个批次大小和猜字数量的组合都要提前录制好一个图。传统推测解码只需要录制每个批次大小配固定猜字数量就行,但动态方案要为每个批次大小录制多个不同猜字数量的图。

工程师们修改了CUDA图管理器,让它能处理批次大小和猜字数量的二维组合。一个批次里有八个请求的时候,要同时录制猜字数量为一和猜字数量为三的图。为什么跳过猜字数量为二呢,因为猜字数量为二的时候八个请求没法平均分配到验证阶段,走不到全CUDA图这个路径。这种细节层面的工程考虑,体现了框架级优化的繁琐程度远超算法本身。

速度不是唯一的胜利条件不拖后腿才是真本事

动态推测解码给我们的最大启示不是它有多快,而是它知道什么时候该快什么时候该慢。传统推测解码像一辆只有油门没有刹车的车,能冲的时候确实快,但遇到弯道就翻。动态方案像一辆装了自动变速箱和刹车系统的车,该快的时候快,该慢的时候慢,全程保持可控。

在生产环境里部署任何加速技术,稳定性往往比峰值速度更重要。一个技术在百分之九十的情况下提速百分之五十,但在百分之十的情况下拖慢百分之两百,整体算下来可能还不如不加速。动态方案的思路就是消除那百分之十的拖慢情况,哪怕在某些场景下提速不那么激进,整体收益反而更高。

从强化学习到大模型服务,从稠密模型到混合专家模型,动态推测解码证明了一个道理。硬件状态在不断变化,算法策略也必须跟着变,固定不变的优化策略最终会被变化的 workload 淘汰。那些看起来最优雅的算法往往不是固定套路的,而是能够感知环境并自我调整的。

这个技术已经被合并进了vLLM主分支,任何使用vLLM的用户都能直接享受到动态推测解码带来的收益。不需要改代码,不需要重新训练模型,升级框架版本就能自动根据硬件状态选择最优猜字数量。