投机推理无损加速大模型解码:接受-拒绝机制与工程实践
1. 投机推理到底在解决什么问题1.1 自回归生成的串行瓶颈从哪来大模型推理有一个绕不开的物理事实自回归生成是串行的。每生成一个token都要把当前全部上下文送进模型做一次前向计算得到下一个token的概率分布采样出结果再把这个结果拼回输入重复下一轮。这个循环里第t个token的计算必须等第t-1个token算完才能开始前后存在硬依赖。这个特性带来的直接后果是GPU的并行算力在解码阶段被严重浪费。训练阶段可以一次性并行处理整条序列的所有位置矩阵乘法的规模很大算力利用率高而解码阶段每个step只处理一个位置的新token虽然注意力部分要跟历史KV做交互但整体计算量相对小访存占比高GPU大量时间花在等显存带宽上而不是在算。我实测过一个典型场景7B级别的模型batch size为1时单步解码的GPU利用率经常在20%以下瓶颈几乎全在显存读取权重和KV cache上。这就是所谓的memory-bound。你堆再多卡单条请求的延迟也降不下来因为串行链条本身没被打破。1.2 投机推理的核心思路用便宜模型猜用贵模型验投机推理Speculative Decoding的切入点很巧妙既然串行链条难破那就让一次前向计算尽量多产出几个token。具体做法是引入一个草稿模型draft model它比目标大模型小得多、快得多先自回归地连续猜出γ个token然后把这γ个猜测token一次性送进目标大模型让大模型并行地对这γ个位置做一次前向得到每个位置的真实概率分布最后用一个接受-拒绝采样规则决定这γ个猜测里哪些能被接受。关键在于目标大模型这一次前向是并行的γ个位置同时算计算量和单步解码相比只是线性增加了一点点因为注意力要覆盖更长的上下文但产出可能不止一个token。如果草稿模型猜得准一次前向就能确认多个token等效于把串行步数压缩了。标题里说的数学严密保分布指的就是这个接受-拒绝机制保证了最终输出的分布和目标模型单独解码时的分布完全一致。这不是近似加速是无损加速。这一点非常重要也是它区别于量化、蒸馏、剪枝等改变模型行为的加速手段的根本所在。1.3 为什么它能加速2到3倍加速比的理论上限取决于草稿模型的接受率。假设草稿模型每个token的接受概率是α那么一次投机循环期望接受的token数是(1-α^(γ1))/(1-α)。当α接近1时加速比可以接近γ1当α很低时加速比趋近于1甚至因为额外开销而变慢。实测中用同系列的小模型做草稿比如用1B级模型给7B或13B模型做草稿α通常在0.6到0.8之间γ取4到8端到端加速2到3倍是很常见的。如果草稿模型和目标模型分布差异大α会掉到0.3以下这时候加速就很不明显了。所以这个技术的收益高度依赖草稿模型的质量。选草稿模型不是越小越好而是要在快和准之间找平衡点。这一点后面会展开讲。2. 接受-拒绝机制是怎么保证分布不变的2.1 一次投机循环的完整流程先把流程拆清楚。设目标模型为M_p草稿模型为M_q当前已生成的上下文为prefix。草稿模型以prefix为条件自回归生成γ个候选tokenx_1, x_2, ..., x_γ同时记录每个位置的草稿概率q(x_i | prefix, x_1...x_{i-1})。把prefix加上这γ个候选token组成长度为len(prefix)γ的序列一次性送进目标模型做前向得到每个位置的目标概率p(x_i | prefix, x_1...x_{i-1})。从i1开始逐个判定对每个位置采样一个均匀随机数r~U(0,1)如果r min(1, p(x_i)/q(x_i))则接受该token否则拒绝并从修正分布norm(max(0, p-q))中重新采样一个token作为该位置的输出然后停止后续判定。如果γ个token全部被接受还可以从目标模型在最后一个位置的分布里额外采样一个token这叫bonus token然后进入下一轮。如果某个位置被拒绝该位置用修正分布采出的token替代之后的候选全部丢弃进入下一轮。这个流程里接受概率取min(1, p/q)是精心设计的。它保证了对任意token x最终被输出的概率恰好等于p(x)。推导不复杂接受路径贡献q(x)·min(1,p(x)/q(x)) min(q(x), p(x))拒绝路径贡献(1-α)·norm(max(0,p-q))(x)两者相加正好是p(x)。这就是数学严密保分布的含义。2.2 为什么不能用简单的概率比阈值判定有人会想直接判断p(x) q(x)就接受否则拒绝不行吗不行。这样做会系统性地偏向高概率token输出分布会偏移生成质量会变。投机推理的接受-拒绝规则必须满足细致平衡式的条件才能保证无损。我见过一些简化实现为了省事用固定阈值或者直接用argmax比对结果就是生成结果的多样性和原始模型不一致长文本生成时会出现重复或退化。这类实现跑benchmark可能看着还行但实际业务里会出问题。所以如果你要自己实现接受-拒绝这一步千万别偷懒。2.3 修正分布的采样细节修正分布norm(max(0, p-q))在实现时有个坑当p和q非常接近时max(0, p-q)可能全为零归一化会除零。工程上一般加一个极小值eps或者在这种情况下直接从p采样。另外这个修正分布只在拒绝位置用一次之后的位置全部丢弃所以计算量很小不影响整体性能。还有一个细节草稿模型和目标模型的词表必须一致否则p和q没法逐token比较。如果两个模型词表不同需要做词表对齐映射这会引入额外误差破坏无损性。所以实践中一般选同系列、同tokenizer的模型做草稿。3. 草稿模型的选型与调优3.1 同系列小模型是首选最省事的方案是用目标模型同系列的小模型做草稿。比如目标用13B草稿用1B或2B目标用70B草稿用7B。同系列意味着tokenizer一致、训练数据分布接近、输出风格相似接受率天然就高。我做过一组对比目标模型固定草稿模型分别用同系列1B、同系列3B、以及一个不同系列的2B。结果同系列1B的接受率约0.68同系列3B约0.79不同系列2B只有0.41。不同系列那个虽然参数量居中但因为分布差异大接受率反而最低端到端速度还不如同系列1B。所以选草稿模型的第一原则是先看tokenizer和训练分布再看参数量。参数量决定草稿模型自身的速度分布匹配度决定接受率两者要一起看。3.2 草稿模型多大合适草稿模型不能太小太小了接受率低也不能太大太大了单次草稿生成的开销就上来了。经验上草稿模型参数量是目标模型的1/10到1/20比较合适。7B目标配0.5B到1B草稿13B目标配1B到2B草稿70B目标配7B草稿。但这个比例不是死的。如果目标模型是经过指令微调的草稿模型最好也做过类似的微调否则在对话场景下接受率会明显下降。我试过用base模型做草稿给chat模型做投机接受率比同是chat的草稿低了将近15个百分点。3.3 γ值怎么定γ是每轮投机的候选token数。γ越大理论上限越高但草稿模型连续猜多个token时越往后越容易猜错接受率会衰减。而且γ大意味着目标模型一次前向要处理更长的序列计算量增加。实测下来γ取4到8是甜点区。γ4时加速比稳定γ8时如果接受率高能更快但接受率一掉就亏。我一般先用γ5跑一轮看接受率如果接受率高于0.7就调到8低于0.5就降到3。还有一个动态γ的策略根据上一轮的接受情况调整下一轮的γ。上一轮全接受就增大γ上一轮早早被拒就减小γ。这个策略在长文本生成里能再挤出10%到15%的加速。4. 工程实现里的关键细节4.1 KV cache的管理投机推理对KV cache的管理比普通解码复杂。草稿模型生成γ个token时它自己的KV cache要往前推γ步目标模型验证时要一次性把这γ个位置的目标KV算出来。如果某些token被拒绝目标模型的KV cache需要回滚到拒绝位置之前。回滚这一步很容易出错。我踩过的坑是拒绝后没有正确截断目标模型的KV cache导致下一轮验证时上下文错位生成结果出现乱码。正确做法是维护一个逻辑长度指针拒绝时把指针回退到接受位置物理显存可以延迟回收。另外草稿模型和目标模型的KV cache是分开的显存占用是两者之和。如果显存紧张可以考虑草稿模型用更小的cache精度或者草稿模型不保留完整历史只保留最近若干token。4.2 批处理下的投机单请求投机好理解但线上服务都是批处理的。批处理下每个请求的接受情况不同有的请求一轮接受5个token有的只接受1个。如果按最慢的请求对齐快的请求就被拖累了。工程上有两种处理方式一种是同步批所有请求等这一轮全部验证完再进入下一轮实现简单但效率有损失另一种是异步批每个请求独立推进接受多的请求先进入下一轮。异步批实现复杂但吞吐提升明显。我实测异步批在混合负载下比同步批吞吐高30%以上。不过异步批对调度器要求高要处理好请求的加入和退出避免长尾请求一直占着资源。这块建议先用成熟推理框架的现成实现自己从头写容易出并发bug。4.3 与连续批处理的配合连续批处理continuous batching是现在推理服务的标配它允许新请求在任意step加入。投机推理和连续批处理结合时要注意草稿阶段的batch和验证阶段的batch可能不一致。草稿阶段每个请求独立生成γ个token验证阶段要把所有请求的候选拼在一起做一次大前向。这里有个优化点可以把多个请求的验证合并成一个大batch提高目标模型的算力利用率。因为验证阶段是并行的合并后矩阵规模更大反而更划算。我在一个并发16的测试里合并验证比逐个验证吞吐高了将近40%。5. 常见问题与排查实录5.1 加速比不达预期甚至变慢这是最常见的问题。排查顺序如下现象可能原因排查方法接受率低于0.4草稿模型与目标模型分布差异大换同系列草稿检查tokenizer是否一致接受率正常但速度慢草稿模型太大或γ太大降草稿参数量降γ测单轮耗时首token延迟高草稿模型加载和预热开销预加载草稿warmup若干轮长文本越跑越慢KV cache膨胀回滚逻辑有bug检查cache截断限制最大长度我遇到过一次加速比只有1.1的情况排查半天发现是草稿模型用了float32而目标模型用float16草稿模型单步耗时比预期高了一倍多。统一精度后加速比回到2.4。这种精度不一致的坑很隐蔽建议一开始就把两个模型的dtype对齐。5.2 生成质量下降理论上投机推理是无损的如果质量下降一定是实现有问题。重点检查三处接受-拒绝的随机数是否正确采样、修正分布是否用了max(0,p-q)而不是其他形式、bonus token是否从正确的分布采样。还有一个容易忽略的点温度参数。如果目标模型解码时用了temperature草稿模型也必须用相同的temperature否则p和q不在同一尺度上接受率会异常。top-k和top-p同理两个模型的采样参数要一致。5.3 显存溢出草稿模型虽然小但加上目标模型的KV cache显存占用可能超出预期。尤其是γ较大时验证阶段要一次性算γ个位置的KV峰值显存会上升。缓解办法限制最大batch size、降低γ、草稿模型用更激进的量化、及时释放被拒绝位置的KV。我一般会留20%的显存余量给峰值不然容易在长序列上OOM。5.4 不同任务的接受率差异代码生成任务的接受率通常比开放域对话高因为代码的局部确定性更强草稿模型容易猜对。数学推理任务接受率偏低因为每一步的推理路径依赖强草稿模型容易跑偏。翻译任务介于两者之间。这意味着同一个投机配置在不同业务上的加速比可能差一倍。上线前一定要用真实业务数据测接受率别拿通用benchmark的数据直接套。6. 实操配置参考6.1 一套可复现的起步配置以目标模型13B、草稿模型1B为例给一套我验证过的起步参数草稿模型同系列1Bfloat16γ初始5根据接受率动态调整范围3到8温度草稿和目标都用0.7top-p0.9两个模型一致batch size单请求先测通再上连续批处理KV cache目标模型用paged attention草稿模型用连续cache先跑100条真实请求统计接受率和端到端延迟再决定是否调γ和换草稿。6.2 接受率监控上线后要持续监控接受率。接受率突然下降通常意味着输入分布变了比如业务从通用对话切到了专业领域问答。这时候要么换更匹配的草稿模型要么对草稿模型做领域微调。我一般会在服务里埋一个接受率的滑动窗口统计按小时聚合。如果连续几个窗口接受率低于0.5就触发告警人工介入看是不是需要调整。6.3 什么时候不该用投机推理投机推理不是万能的。如果目标模型本身很小比如1B以下草稿模型很难再小加速空间有限。如果业务对延迟不敏感但对吞吐敏感直接上大batch普通解码可能更简单。如果显存极度紧张加一个草稿模型的额外开销可能不划算。我的判断标准是目标模型参数量在7B以上、单请求延迟是瓶颈、显存有至少30%余量这三个条件满足时投机推理值得上。否则先优化其他环节。7. 我对这项技术的实际体会投机推理最吸引我的地方是它的无损特性。很多加速手段都是在精度和速度之间做trade-off而投机推理在数学上保证了输出分布不变这让它在对生成质量敏感的场景里特别有吸引力。你不需要重新训练模型不需要校准只要挂一个草稿模型就能拿到2到3倍加速工程改造成本相对可控。但它也不是没有代价。草稿模型的维护、接受率的监控、KV cache回滚的正确性这些都是额外的工程负担。我见过团队兴冲冲上线结果因为回滚bug导致线上生成乱码回滚修复花了两周。所以我的建议是先在离线环境把接受-拒绝逻辑和cache管理测透用单元测试覆盖各种拒绝位置再上线上。另外草稿模型的选择是个持续调优的过程。业务变了草稿模型可能就要换。把它当成一个需要长期维护的组件而不是一次性的优化心态上会更稳。后续如果草稿模型也能做量化或者蒸馏进一步压低它的开销整个方案的性价比还能再上一个台阶。