做LLM推理优化的朋友应该都有这种体会模型越来越大输入越来越长可自回归生成却一次只能吐一个tokenGPU算力被卡得死死。我这次在项目里尝试用FlashSpec实现推测解码speculative decoding目标很直接在不改变目标模型结构、不牺牲输出质量的前提下把解码速度提上去。整个过程从理论到落地花了差不多两周把草稿模型选型、采样一致性、KV缓存复用、拒绝采样这些环节通通过了一遍也踩了不少坑。这篇文章就把这套实践完整拆开重点讲讲FlashSpec的核心机制、具体实现步骤以及那些不在生产环境里撞一次根本不会注意到的细节给正在做推理加速的朋友一个可参照的路线。1. 先搞清楚推测解码到底在解决什么问题1.1 自回归生成为什么这么慢大模型生成文本的时候用的是自回归方式每生成一个token都要把当前所有输入重新过一遍网络然后只取最后一个位置的概率分布做采样。问题就出在这里——解码阶段每次只计算一个token但GPU要加载的参数量是完整的模型权重和KV cache。对一个7B模型来说哪怕只生成一个token显存里也要搬几十GB的数据计算量反而不大。这种场景叫“memory-bound”卡在内存带宽上不是卡在算力上。我拿A100 40G实测过一个7B模型prefill阶段处理512个token的输入吞吐可以做到几千token每秒但进入decode阶段之后单个请求的生成速度可能只有20到40 token每秒差距非常夸张。原因就是decode阶段的矩阵形状太“瘦”了GPU的大规模并行能力根本施展不开大部分算力都在空转。这个空转的窗口就是推测解码想要利用的机会。1.2 投机采样的两条核心思路推测解码的基本思路可以概括成两句话先用一个更小更快的草稿模型去猜接下来几个token再用目标模型一次性校验这些猜测。如果猜对了批量接受如果猜错了在错的位置截断并做修正。整个过程不会改变目标模型的输出分布理论上和直接用目标模型一个一个生成是等价的。这个思路的关键在于两个环节。第一草稿模型要足够快生成的候选序列不能拖后腿。第二目标模型的校验要“一次前向验证多个token”而不是对每个候选token分别做一次前向。也就是说预填充和验证阶段可以充分利用GPU并行度让原来每次只能前进一个step的串行过程变成“一个step内前进k个token”的过程。需要注意的是因为目标模型的概率分布和草稿模型不一样所以不能简单“猜对了就收下、猜错了就全丢”。正确的做法是基于两个模型在每个位置的输出概率做一次拒绝采样用概率比值来决定接受还是拒绝。这样即使草稿模型在某些位置给出很低概率的token目标模型只要没反对也能被接受从而保持整体输出分布和真实模型一致。1.3 FlashSpec的方案定位FlashSpec是我们在对比了几种开源实现之后选择的方向。它本质上是一套轻量级的推测解码框架不改目标模型通过外挂草稿模型和一套并行验证调度逻辑来加速。和完整跑一遍“两阶段精调蒸馏模型”这类方案相比FlashSpec的实现成本低主要靠推理时的投机采样来提效非常适合已经有现成大模型、想快速压低单请求延迟的落地场景。它对环境的要求也比较常规只要目标模型能用Transformers或vLLM加载草稿模型选一个参数小、生成速度快的就能跑起来。后面我会详细讲我用的模型搭配和具体参数这套组合在我实测环境中拿到了约1.6到2.3倍的解码加速不算夸张但已经很实用。2. FlashSpec 实现的关键设计拆解2.1 草稿模型的选型与参数量平衡草稿模型不是越小越好也不是越大越好。太小了猜不准接受率低校验时拒绝位置太靠前加速收益完全被浪费太大了生成候选序列的时间本身就很可观反而拖慢整个循环。这里要引入一个概念接受率。它表示一个候选序列中被目标模型接受的平均token数。比如草稿模型一次性生成8个候选token目标模型校验后平均接受了3个那么有效加速比大概就是“每轮实际前进3个token / 每轮总耗时”。如果草稿模型生成8个候选耗时是目标模型解码一步的0.3倍校验耗时是0.8倍那么理想加速比大概是3 / (0.3 0.8) ≈ 2.7但这是建立在接受率为3的基础上的。我实测下来一个7B级别的目标模型搭配以下草稿模型的差别很大草稿模型参数量平均接受token数相对目标模型速度实际加速比GPT-2 124M0.1B1.8约30倍以上1.5xGPT-2 1.5B1.5B3.5约8倍2.0x自定义小Transformer 3B3B4.2约3倍1.6x原因很直观草稿模型太小猜测的质量不够经常第一个token就被拒绝等于白白多跑了一轮草稿生成草稿模型太大生成候选序列本身变慢而且占用显存影响目标模型的batch调度。我最后选定的方案是1.5B左右的草稿模型在显存占用和接受率之间取了一个相对舒适的平衡点。2.2 并行校验与KV缓存复用推测解码真正加速的核心在“并行校验”这一步。目标模型可以直接把草稿序列作为输入一次前向计算出所有位置的logits而不是像正常解码那样一个一个来。在校验过程中必须复用KV cache否则每校验一个候选位置都要重新计算之前的key和value那就完全没有加速效果了。具体实现上我采用的是两步式推进草稿模型先独立生成k个候选token同时保留它在每个位置的KV cache这样下一轮预测候选序列时不需要重新计算前面内容。目标模型把原始输入和k个候选token拼接成一段序列一次性前向得到每个位置的目标分布然后从第一个候选token开始做拒绝采样。这里有个容易被忽略的点目标模型在验证候选序列时注意力掩码必须要覆盖到完整的“原始输入 候选token”而不是只覆盖候选token。一个错误做法是把候选token当成一个独立batch只算它们之间的注意力这样校验出来的logits完全不对因为候选token没有看到原始输入的上下文。2.3 采样一致性与拒绝采样拒绝采样是保证分布等价的核心算法也是很多人容易写错的地方。标准流程是这样的草稿模型对当前位置i给出分布p_i并从中采样出候选token t_i。目标模型校验时给出分布q_i也是在同一位置的条件下。如果t_i在两个分布下都满足条件则按min(1, q_i(t_i) / p_i(t_i))的概率接受这个token。如果被拒绝则从修正分布 (q_i - p_i)_ / sum(...) 中采样一个token替代然后这一轮的投机采样结束。也就是说即使草稿模型在每个位置给的token概率都很高只要目标模型不同意就可能被拒绝反过来如果草稿模型概率低但目标模型概率高接受概率会被限制到1不会超采。这套机制保证最终输出和单独跑目标模型的分布完全一致只是走了个“捷径”。实际操作中我遇到的最大坑是“采样参数没有对齐”。如果草稿模型用temperature0.8采样而目标模型校验时用的是temperature0或者两端top_p不一致那么p_i和q_i根本不在同一个温度空间下拒绝采样公式里的概率比值就失去了意义。表现就是接受率骤降加速比直接变成负数。3. 实操从零搭一个 FlashSpec 推理流水线3.1 环境、硬件与依赖清单我这次跑通用的是一张A100 40G目标模型是LLaMA-2-7B-ChatFP16草稿模型是GPT-2 1.5B。软件环境如下CUDA 11.8PyTorch 2.1.0Transformers 4.36.2FlashAttention 2有一点要提前说FlashAttention 2不是必须的但建议装。推理时KV cache和长序列的注意力计算很耗显存带宽用FlashAttention之后目标模型校验前向的速度有肉眼可见的提升。如果环境里没有编译FA2也可以用标准attention跑只是加速比会打折扣。另外草稿模型我用的是FP16精度显存占用大概在3GB左右。如果目标模型占用了大部分显存需要确认剩余空间足够容纳草稿模型和KV cache否则就得考虑CPU offload方案后面踩坑部分会细说。3.2 核心实现代码草稿、校验、采样三件套下面的代码是精简后的核心逻辑完整工程里还要加上并发调度、日志和异常保护但主体就这三步import torch import torch.nn.functional as F def draft_sampling(draft_model, input_ids, past_kv, max_tokens, temperature, top_p): candidate_ids [] current_ids input_ids current_kv past_kv for _ in range(max_tokens): with torch.no_grad(): outputs draft_model( input_idscurrent_ids, past_key_valuescurrent_kv, use_cacheTrue, ) logits outputs.logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) # 简单的top-p截断 sorted_probs, sorted_idx probs.sort(descendingTrue) cumsum sorted_probs.cumsum(dim-1) mask cumsum - sorted_probs top_p sorted_probs[mask] 0.0 probs sorted_probs / sorted_probs.sum(dim-1, keepdimTrue) next_id torch.multinomial(probs, num_samples1) candidate_ids.append(next_id) current_ids next_id current_kv outputs.past_key_values if next_id.item() draft_model.config.eos_token_id: break return torch.cat(candidate_ids, dim1), current_kv def verify_with_target(target_model, input_ids, candidate_ids, past_kv, temperature): # 把原始输入和候选token拼接在一起一次前向 full_ids torch.cat([input_ids, candidate_ids], dim1) with torch.no_grad(): outputs target_model( input_idsfull_ids, past_key_valuespast_kv, use_cacheTrue, ) # 取每个候选token位置对应的logits verify_logits outputs.logits[:, input_ids.size(1) - 1 : -1, :] / temperature return verify_logits def speculative_sampling(draft_model, target_model, input_ids, past_kv, gamma6, temperature1.0, top_p0.9): # 1. 草稿模型生成gamma个候选 candidate_ids, new_draft_kv draft_sampling( draft_model, input_ids, past_kv, gamma, temperature, top_p ) # 2. 目标模型并行校验 target_logits verify_with_target( target_model, input_ids, candidate_ids, past_kv, temperature ) target_probs F.softmax(target_logits, dim-1).squeeze(0) # 3. 逐位置拒绝采样 accepted [] last_rejected_sample None draft_cpu candidate_ids.squeeze(0).cpu() for i in range(candidate_ids.size(1)): token draft_cpu[i].item() p_prob get_draft_prob(draft_model, draft_cpu[: i 1], token) q_prob target_probs[i, token].item() if q_prob p_prob: accepted.append(token) else: if torch.rand(1).item() q_prob / p_prob: accepted.append(token) else: last_rejected_sample renormalize_and_sample( target_probs[i], draft_probs, token ) break n_accepted len(accepted) new_input_ids torch.cat( [input_ids, torch.tensor(accepted, deviceinput_ids.device).unsqueeze(0)], dim1, ) return new_input_ids, n_accepted, new_draft_kv这段代码里的get_draft_prob和renormalize_and_sample是辅助函数。前者需要从草稿模型的缓存里重新取对应位置的logits后者负责从修正分布中采样替代token。真实工程里不要忘了处理“候选全部被接受”的情况此时还需要额外从目标模型最后一个位置多采样一个token把这个token也当作本轮新增token避免输出永远不会超过候选序列长度。3.3 数据流与加速比计算整个推理循环的结构是prefill一遍然后进入投机采样循环。每次循环会产生一个新的候选序列目标模型校验一次然后决定接受多少个token接受的部分被拼接到输出序列同时KV cache更新到被接受的最后一个位置然后草稿模型从这个新位置继续生成下一轮候选。加速比可以用一个简单的公式估算有效加速比 平均每轮接受token数 / (草稿生成耗时 目标校验耗时 采样开销)注意这个公式里的三个时间项都是相对于“目标模型正常解码一步”的时间来算的。我实际测下来的典型值是平均接受token数3.4草稿生成6个token耗时约0.4个目标decode步长目标校验6个候选耗时约0.9个目标decode步长采样修正开销约0.1个目标decode步长那么加速比 3.4 / (0.4 0.9 0.1) 2.43。但因为接受率在不同输入上有波动实际平均下来在2.0左右。别小看这个数字如果线上服务几千路并发每个请求都快一倍对应的GPU成本直接降一半。4. 实测效果到底能快多少4.1 测试模型与基准我用三组数据做了对比测试分别是CNN/DailyMail摘要、GSM8K数学题、以及一组代码生成任务。目标模型统一是LLaMA-2-7B-Chat输入长度控制在512到1024个token之间输出长度限制在256个token以内。对比基准是关闭推测解码的原始解码逻辑采样参数全部保持一致。测量指标主要有两个单请求平均生成耗时以及端到端吞吐token/s。我专门关掉了动态batching因为推测解码主要优化的是单请求延迟如果混入服务端batching结论容易混淆。4.2 性能对比结果下面是CNN/DailyMail摘要任务上的结果方案平均生成耗时(s)吞吐(tokens/s)相对加速原始解码12.820.11.0xFlashSpec GPT-2 124M8.430.51.5xFlashSpec GPT-2 1.5B6.539.41.9x在GSM8K上因为生成内容包含更多推理步骤草稿模型更难猜中目标模型的下一个token平均接受率降到了2.1左右加速比大概1.5倍。而在代码生成这种局部模式很强的任务上接受率能飙到4.5加速比到了2.3倍。这说明推测解码的实际收益和任务类型强相关不要指望所有场景都有统一效果。4.3 什么场景不适合FlashSpec有一点我在一开始就踩过坑投机采样对单请求、低并发场景非常友好但对高并发服务端场景不见得划算。原因很简单动态batching本身已经把多路请求的decode step合并成了更大的batchGPU利用率已经拉满此时再叠加投机采样优势会被严重稀释甚至因为多了一层调度开销而变慢。我实测batch size到8之后FlashSpec相比原始解码的加速比掉到1.1到1.2左右再往上基本没有收益。因此我的结论是需要单流低延迟的场景比如交互式对话、代码补全、实时翻译优先考虑推测解码追求整体吞吐、跑离线批量任务、或已经大量并发打满GPU的场景收益有限优先优化batching和调度策略。5. 踩坑记录与排查思路5.1 采样参数没对齐校验全崩这是我在完成第一版集成后遇到的最诡异的问题接受率掉到不足1.0生成速度比原始解码还慢30%。排查了很久才发现草稿模型使用了默认的greedy采样temperature0而目标模型校验用的是temperature0.8的采样结果。两边拿到的概率分布根本不是同一个采样策略下的分布拒绝采样的接受逻辑完全失效。修复方案是确保草稿生成、目标校验、拒绝采样三处使用完全一致的temperature和top_p并且采样过程要固定随机种子尽量保证可复现。这个坑的隐蔽之处在于代码不会报错只会表现为“性能异常”很容易让人误以为是并行校验写错了。5.2 KV cache 索引错位结果乱成一团第二次大坑出现在KV cache的复用上。一开始我为了节省显存在目标模型校验时传了past_key_values但没有注意position_ids要从input_ids.size(1) - 1的位置开始。结果前几轮生成还能用到第三轮开始输出完全错乱甚至出现重复片段。问题在于目标模型的KV cache里已经缓存了原始输入和之前接受的token拼接候选序列时新token的position id应该是“已处理长度”而不是从0开始。我在调查时对比了Transformers不同版本的use_cache行为发现有些版本会自动推断有些版本严格依赖传入的position_ids一旦没传就默认从0算起。最后我统一改成显式构造position_ids并且在拼接候选token时用attention_mask把原始输入长度标出来问题才彻底解决。5.3 草稿模型放GPU还是CPU显存与延迟的权衡目标模型加载完毕后40G显存里大概剩了不到5G。如果草稿模型用GPT-2 1.5B FP16还剩一点空间但如果同时开多个服务实例显存就非常紧张。我尝试过把草稿模型offload到CPU结果每轮生成候选序列都要从CPU搬运权重到计算设备通信开销差点把推测解码的收益全吃掉。最终我采用了两路并行的折中方案草稿模型放在GPU但用FP16并把最大候选长度从8降到6同时使用CUDA Graph把草稿模型的前向调用固定下来减少kernel launch开销。实测下来显存占用稳定控制在目标模型草稿KV cache整体35GB左右留够了余量。如果你连2GB都挤不出来我建议把草稿模型换到0.1B级别接受率低一点但至少能保证整个推理流程不崩。5.4 混合精度导致概率分布漂移有一次优化时我把目标模型校验阶段的logits也改成了FP16计算结果发现单测时输出分布和原始模型出现肉眼可见的差异。原因是目标模型在FP16下计算logits时尾部概率的精度不够导致拒绝采样公式里q_prob / p_prob出现较大误差尤其是当草稿概率很低、目标概率作为分子时FP16舍入误差会被放大。解决办法很简单校验阶段的logits和softmax计算放到FP32下执行只在KV cache和线性层计算保留FP16加速。这个改动只影响极小一部分显存和计算量但输出分布的一致性能明显提升。5.5 问题速查表为了方便大家排查我把这轮的几个典型问题和解决方法整理成一张表现象根因处理方法接受率低于1.0速度下降草稿模型和目标模型采样参数不一致统一temperature、top_p固定随机种子输出开始重复/乱码position_ids从0开始导致KV cache错位显式传入position_ids处理attention_mask显存不足或频繁OOM草稿模型过大或候选序列过长换小草稿模型、降低gamma、使用FP16、考虑CPU offload输出分布和原始模型不一致校验阶段FP16精度不够logits和softmax切到FP32高并发下加速比消失动态batching已经拉满GPU利用率集中优化batching或仅在低并发场景启用推测解码6. 一点个人体会前前后后调了几周我的最大感受是推测解码不是银弹它适合的场景非常明确——单请求延迟敏感、GPU算力没有完全打满、任务有比较强的上下文惯性比如代码和摘要。在这些前提下FlashSpec这套“小草稿模型目标模型并行校验”的方案能带来接近2倍的端到端提升已经很值了。如果让我给后来人一个建议先别急着调代码先把你自己的场景数据测一遍算清楚“理论加速比”到底有没有空间。草稿模型选型的试错成本比调参低得多多测几个规模找到接受率和生成耗时的最佳交叉点比盲目堆候选长度有用得多。这个方向后面我还在继续探索比如在目标模型本身支持前缀缓存的情况下配合投机采样能不能进一步压低首token延迟等有结论了再来分享。
