作为一个成天和大模型推理耗时间的人我一直在关注怎么把生成速度再往上顶一顶。之前用vLLM做常规优化TGI也试过到瓶颈之后开始折腾推测解码Speculative Decoding。一开始用的Medusa头和EAGLE效果有但总觉得训练成本高、实现也绕。直到看到FlashSpec这个思路实测之后发现它在EAGLE系列里属于比较“聪明”的一类而且实现起来没有想象中复杂。这篇博文直接聊聊我实践FlashSpec的过程重点说说它到底是什么、我在接入时怎么改造的、以及中间踩过的那几个坑。不堆公式尽量说人话给想动手复现的同学一条相对顺滑的路径。1. 推测解码到底是什么FlashSpec的切入点在哪里想用好FlashSpec得先弄清楚推测解码的问题在哪。它不是什么魔法思路很直白小模型先生成一批候选token大模型一次性验证一批而不是一个token一个token地等。理想情况下你生成4个token大模型一次前向就把这4个全接受了理论上吞吐能翻几倍。但实际做的时候有两个坎绕不开。第一个是草稿模型的“智商”问题小模型生成的候选如果经常被大模型拒绝那还不如不推。第二个是验证环节的并行度问题动态形状、KV Cache管理、树形验证这些工程细节特别容易拖后腿导致理论上限很高实际落地打对折。FlashSpec的切入点就很明确它走的是EAGLE这条路利用LLM的深层隐藏状态来做草稿模型的输入。说得再直白一点EAGLE系列的核心洞察是如果你知道第i层的隐藏状态那么预测第i1层的内容比单纯用token embedding去猜要准得多。FlashSpec在这个基础上进一步做了浅层权重共享和分块并行验证把草稿模型的训练成本压到很低同时验证阶段不再一条路走到黑而是分块并行去验。所以FlashSpec适合谁适合那些已经跑通EAGLE、但觉得训练草稿模型太贵、或者验证阶段加速比上不去的团队。也适合在长序列生成上有刚需的人比如代码生成、长文档摘要因为这些场景下循环次数多推测收益才大。短文本对话不是它的主战场收益有限这是我实测后的直观感受。1.1 它和EAGLE、Medusa这类方案的本质区别先说说Medusa。Medusa的思路是多头并行预测一个位置同时预测多个未来token然后去做树状验证。它的优点是理论上限高但缺点是训练不稳定而且推理时显存占用是真的吓人。我试过一次Medusa多头配置直接把我那张A100干到OOM边缘。EAGLE就不一样。EAGLE是把深层隐藏状态和token embedding拼接起来作为草稿模型的输入来预测下一个token。这相当于草稿模型是“借大模型的脑子”来做预测所以接受率天然比独立小模型要高。EAGLE-2还加了动态树机制根据置信度调整树结构。FlashSpec在EAGLE基础上做了两个关键改动。第一草稿模型不再从零训练而是复用LLM的前几层浅层权重只训练中间一小部分参数这就把训练成本从“训练一个小模型”降到了“微调几个层”。第二验证阶段不再一次性等草稿模型生成完整序列后再验证而是分块生成、分块验证让草稿模型和验证模型在时间上重叠起来。这两点合在一起就是FlashSpec能同时兼顾训练成本和推理速度的根本原因。1.2 为什么长序列场景收益更明显实测下来FlashSpec在生成128到512个token的段落时收益最明显。原因是草稿模型在一次prefill之后后续每走一步只需要增量更新KV Cache这个模式下分块并行验证的优势才能完全释放。而如果你生成短句比如“你好”这种草稿模型还没进入状态就结束了反而白白浪费一次prefill开销。而且长序列场景下FlashSpec的分块策略有两个实实在在的好处。一是草稿模型的KV Cache可以提前算好验证阶段直接复用不用反复重算。二是总生成步数变多之后接受率的微小优势会被放大比如接受率从0.6提到0.7看似只提升了10个百分点但总加速比可能从1.8跳到2.5。这中间的差额全部来自“等待大模型验证”的时间被压缩了。另外说一下FlashSpec论文里报告的是在LLaMA系列和自回归模型上的结果但我自己试过把它迁移到Qwen架构上也没有太大障碍。关键点是理解它设计上对“分层”的依赖这个后面在改造时会详细展开。2. 动手前的准备环境搭建和模型选型这一节先把基础设施说清楚方便直接抄作业。硬件环境我跑通整套流程用的是一张A100 80G。FlashSpec整体显存开销比EAGLE大一点点因为草稿模型需要额外存一份权重但比Medusa还是小不少。如果你是V100或者3090建议用小模型比如7B级别草稿模型层数减半。A100的话13B到14B的模型都问题不大。软件栈Python 3.10CUDA 12.1PyTorch 2.1.2transformers 4.37.2accelerate 0.27.2基于FlashAttention 2.3的flashinfer做分块验证时这个库很关键这里多说一句transformers版本最好锁定在4.37左右。太新版我遇到过兼容问题尤其是cache_position参数在推测解码时容易报错。太旧版又不支持static cache导致KV Cache管理非常痛苦。基座模型选择如果你的目标是复现论文效果选LLaMA系列最稳妥因为FlashSpec官方实验是基于LLaMA的。但如果你只想在自己的业务上体验一下提速效果我建议直接从Qwen2-7B或者Mistral-7B开始原因很简单——推理框架对这两个架构的支持更成熟遇到奇怪的bug时社区也更容易给出答案。我最终选的是Llama-2-7B-hf作为基座模型权重可以从HuggingFace直接拉不需要做额外处理FlashSpec的草稿模型训练脚本会自动把浅层权重拷贝到草稿模型上。2.1 浅层权重共享的原理为什么草稿模型瘦身这么多FlashSpec里草稿模型的构造比较特殊——不是一个小型独立Transformer而是把基座模型的前N层直接搬过来复用然后在后面接一个轻量的预测头。这样做最大的好处是草稿模型不需要重新学习底层的语法和词法知识这些信息在浅层就已经编码好了。打个比方基座模型就像一个有经验的老师傅前几层相当于他的“基本功”。草稿模型直接继承这个“基本功”只学怎么快速做出预判而不是从头学说话。这样草稿模型的参数量可以压缩到基座模型的1/4甚至更少训练数据需求也大幅下降。实际操作中我取的是基座模型的前8层作为共享层。这个数字不是拍脑袋定的而是根据模型的隐藏层数量动态算的。Llama-2-7B隐藏层总共32层取前1/4刚好是8层。太少了草稿模型学习能力不足太多了显存和训练时间会显著增加收益却不明显。2.2 环境配置的细节坑环境搭建这一关我栽了一个不小的跟头。刚开始用的flash-attn版本是2.5.0结果在加载草稿模型时频繁报thread同步错误。后来定位发现是flashinfer和flash-attn之间版本冲突。我的建议是flash-attn锁到2.3.0不要盲目追新flashinfer用0.1.2版本CUDA最好用12.1而不是11.8因为torch.compile在12.1环境下对分块验证的支持更友好如果你想用Docker我提供一个可用的镜像组合pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime作为基础镜像然后手动安装flashinfer 0.1.2。不建议直接用flashinfer官方镜像因为它里面捆绑的transformers版本太新跑FlashSpec时容易踩到generation_config不兼容的坑。3. FlashSpec核心改造流程四步走整个改造的核心本质上是三步构造草稿模型训练草稿模型推理阶段分块验证。我把它拆成四个严格有序的步骤按这个顺序做能少走很多弯路。3.1 构造草稿模型浅层共享预测头这个步骤的目的是生成一个参与推理的草稿模型对象。它和基座模型共享前8层权重但后续层是随机初始化的。代码逻辑可以这样理解import torch import torch.nn as nn from transformers import LlamaForCausalLM, LlamaConfig base_model LlamaForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf) base_config base_model.config # 草稿模型的配置只有8层隐藏维度、注意力头数保持一致 draft_config LlamaConfig( vocab_sizebase_config.vocab_size, hidden_sizebase_config.hidden_size, intermediate_sizebase_config.intermediate_size, num_hidden_layers8, # 关键只保留8层 num_attention_headsbase_config.num_attention_heads, num_key_value_headsbase_config.num_key_value_heads, max_position_embeddingsbase_config.max_position_embeddings, ) # 创建草稿模型并复用基座模型的前8层权重 draft_model LlamaForCausalLM(configdraft_config) # 权重拷贝浅层参数共享 draft_state_dict draft_model.state_dict() base_state_dict base_model.state_dict() for name, param in draft_state_dict.items(): if model.layers. in name: layer_idx int(name.split(.)[2]) if layer_idx 8: # 对应基座模型里同名的层 source_name name draft_state_dict[name] base_state_dict[source_name] # 注意这里要冻结权重训练时不更新 param.requires_grad False draft_model.load_state_dict(draft_state_dict)这个构造过程有几个隐藏边界条件词嵌入层embed_tokens和输出层lm_head一定要从基座模型复制这是草稿模型能正常工作的基础前8层的requires_grad必须设为False否则一训练就把宝贵的共享权重破坏了第8层之后新增的层我这里加了2层可训练的适配层才是真正需要训练的部分通常用nn.Linear加SiLU激活函数做轻量预测头预测头我是这样设计的class DraftPredictionHead(nn.Module): def __init__(self, hidden_size, vocab_size): super().__init__() self.fc1 nn.Linear(hidden_size, hidden_size // 2) self.act nn.SiLU() self.fc2 nn.Linear(hidden_size // 2, vocab_size) def forward(self, hidden_states): return self.fc2(self.act(self.fc1(hidden_states)))这个设计参考了EAGLE的做法意图是让前8层共享的特征经过一个轻量非线性变换后再映射到词表空间比直接接一个线性的lm_head收敛更快最终接受率也更高。3.2 训练草稿模型小数据、快收敛草稿模型训练是整个实践里最“像普通微调”的一步但有几个关键差异点。首先训练数据不需要和基座模型的预训练数据一个量级几十万条高质量文本就够。其次训练目标很简单让草稿模型学会预测“基座模型会接受的token”而不是单纯预测下一个token。这里有一个细节需要注意训练时的输入不应该只用真实的上一轮token而是要用基座模型的输出去做教师强迫teacher forcing。也就是说我们把输入喂给基座模型拿到基座模型的隐藏状态让草稿模型学习从这个隐藏状态去预测下一个token。这背后的逻辑是草稿模型服务的是“验证阶段”这时候大模型已经产出了真实历史token草稿模型要做的是基于这些历史token和基座模型的隐藏状态来预测候选token。如果训练时只给它看真实token不告诉它基座模型是怎么理解的推理时就会水土不服。训练脚本的关键部分如下from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling from datasets import load_dataset # 构造训练数据用基座模型生成隐藏状态作为草稿模型的输入 def preprocess_function(examples, tokenizer, base_model): inputs tokenizer(examples[text], return_tensorspt, max_length512, truncationTrue) with torch.no_grad(): outputs base_model(input_idsinputs[input_ids].to(cuda), output_hidden_statesTrue) hidden_states outputs.hidden_states[8] # 取第8层的隐藏状态 return {input_ids: inputs[input_ids], hidden_states: hidden_states} training_args TrainingArguments( output_dir./draft_model, per_device_train_batch_size4, learning_rate5e-5, num_train_epochs3, gradient_accumulation_steps8, fp16True, save_steps500, )这里取outputs.hidden_states[8]是关键中的关键它表示第8层的隐藏状态也是草稿模型前8层共享输出的最后一层。训练完成后草稿模型就能根据这个隐藏状态去映射到下一个token的分布。训练时间方面我的实测数据是50万条数据8张A100大概4小时能收敛到比较理想的状态。如果只有单卡A100建议把数据缩减到15万条训练3轮也能达到可接受的效果只不过确实差一些。3.3 推理阶段的分块验证理解核心机制训练好草稿模型之后真正的重头戏是推理阶段的分块验证。这也是FlashSpec和EAGLE最大的区别所在。EAGLE的验证流程是典型的串行草稿模型先预测K个token然后大模型一次性验证这K个token接收部分候选再继续预测下一轮。这有一个明显的等待时间浪费草稿模型生成第K个token时大模型完全闲着。FlashSpec的分块验证思路是把K个候选token切成B块草稿模型每生成一块大模型就验证一块而不是等全部生成完。这样草稿模型的生成和大模型的验证在时间上重叠起来从整体延迟角度看等于草稿模型的生成时间被“藏”进了大模型的验证时间里。看一下具体实现逻辑def speculative_generate( draft_model, target_model, input_ids, max_new_tokens256, num_speculative_tokens5, block_size2, # 分块大小 ): # 缓存初始KV状态 draft_kv init_kv_cache(draft_model, input_ids) target_kv init_kv_cache(target_model, input_ids) for _ in range(max_new_tokens // num_speculative_tokens): # 分块生成 candidate_tokens [] draft_hidden_states get_hidden_states(draft_model, input_ids, draft_kv) for block_start in range(0, num_speculative_tokens, block_size): block_end min(block_start block_size, num_speculative_tokens) needed_tokens block_end - block_start # 草稿模型生成一个block的候选token block_candidates draft_model.generate( hidden_statesdraft_hidden_states, kneeded_tokens, kv_cachedraft_kv, ) candidate_tokens.extend(block_candidates) # 立即交给目标模型验证这一块 accepted_tokens, target_kv verify_and_update( target_model, input_ids, block_candidates, target_kv ) if not all_accepted(accepted_tokens): break input_ids torch.cat([input_ids, torch.tensor([candidate_tokens])], dim-1) return input_ids这个伪代码的意图很直白草稿模型每生成一小块立刻喂给大模型验证。如果这一块全被接受了就继续生成下一块如果中途被拒绝就从这个断点重新开始。这样做的结果就是草稿模型生成和验证模型的推理是流水线式并行的整个生成过程的总延迟大幅下降。实测下来分块大小取2到3之间最稳。太大会退化成EAGLE的串行模式失去并行优势太小会导致频繁切换上下文增加调度开销。3.4 验证器的具体实现如何判断“接受”还是“拒绝”分块验证的核心逻辑在于验证器怎么判断草稿模型的哪些token可以被接受。这里用的是经典的拒绝采样rejection sampling策略但FlashSpec做了优化。最简单的理解方式是这样的草稿模型给出一个候选token同时也给出了它对这个token的概率估计记为p_draft。大模型在验证时会对这个候选位置计算自己的概率分布记为p_target。接受条件就是如果p_target p_draft直接接受因为这个token确实是大模型自己也会选择的如果p_target p_draft则以p_target / p_draft的概率随机接受否则从修正后的分布里重新采样一个token这个机制的意图很清楚它保证了最终输出分布严格等于大模型自身的分布不会因为草稿模型的加入而产生偏置。这也是推测解码“无损”的底气所在。代码层面验证器实现如下def verify_and_update(target_model, input_ids, candidate_tokens, target_kv): with torch.no_grad(): logits target_model( input_idsinput_ids, past_key_valuestarget_kv, use_cacheTrue, ).logits[:, -1, :] # 目标模型对当前位置的预测分布 target_probs torch.softmax(logits, dim-1) accepted_tokens [] for cand in candidate_tokens: p_draft draft_prob(cand) # 草稿模型的概率 p_target target_probs[0, cand] # 目标模型对这个token的概率 acceptance_prob min(1.0, p_target / p_draft) if torch.rand(1).item() acceptance_prob: accepted_tokens.append(cand) target_probs target_probs.clone() # 注意这里要继续更新target的KV cache这里简化为循环内逐步更新 else: # 从修正分布里重新采样 adjusted_probs torch.relu(target_probs - draft_probs) adjusted_probs adjusted_probs / adjusted_probs.sum() accepted_tokens.append(torch.multinomial(adjusted_probs, 1).item()) break return accepted_tokens, target_kv这个验证器有几个需要注意的坑logits取的位置要正确。在past_key_values存在的情况下logits[:, -1, :]表示当前最后位置的下一个token预测直接对应草稿模型预测的下一个位置不需要额外偏移被拒绝后重新采样的分布是max(0, p_target - p_draft)这个修正分布保证了最终输出和大模型单独生成时的分布一致KV Cache的更新要和接受token一一对应否则后面验证的位置全部错位这是最容易出bug的地方4. 实测效果加速比数据不会骗人跑完实践最终要拿出数据说话。我在相同输入、相同模型、相同硬件的条件下对比了原始生成、EAGLE-2和FlashSpec三种方案。测试环境A100 80GLlama-2-7B输入序列长度512输出序列长度256batch size为1。方案首次生成延迟每token平均延迟加速比原始生成63ms/token63ms1.0xEAGLE-237ms/token38ms1.68xFlashSpec29ms/token31ms2.08x这个结果符合我的预期。FlashSpec比EAGLE-2高出大约0.4倍的加速比来源就是分块并行验证把草稿模型生成和验证两个阶段重叠了起来。再换到batch size为8的场景方案每token平均延迟加速比原始生成152ms1.0xEAGLE-289ms1.71xFlashSpec67ms2.27xbatch size变大时FlashSpec的优势甚至更明显了。原因是batch inference时目标模型的单次前向开销被分摊到多个序列上分块并行验证的重叠效率更高。不过也要泼一盆冷水如果输入序列只有32个token输出也只有20个tokenFlashSpec的加速比会掉到1.2倍左右甚至有时还不如原始生成。因为模型前几轮生成时草稿模型还没热身完KV Cache也不够深分块并行还没来得及发挥生成就结束了。所以评估你自己的业务场景时不要只看峰值加速比要看你实际的token数量分布。5. 避坑实录我踩过的那些深坑这一节是全文的重点每一个坑我都是真金白银换来的。5.1 坑一transformers版本不兼容KV Cache报错我一开始用的transformers 4.42版本跑FlashSpec的验证循环时一直报past_key_values维度和input_ids对不上的错误。查了半天发现是新版transformers改造了cache_position的处理逻辑导致在手动维护KV Cache时行为不一致。解决方法是降级到4.37.2。这个版本对static cache和dynamic cache的处理方式比较稳定FlashSpec官方代码也是基于这个版本开发的。如果你非要留在新版就需要重写一部分cache逻辑我个人不建议浪费时间。5.2 坑二草稿模型训练时loss不降刚开始训练草稿模型时我观察到了一个很诡异的现象前几百步loss纹丝不动甚至偶发上升然后又突然骤降。这个现象的原因是前8层的共享权重已经非常接近收敛状态如果用统一的学习率去训练会导致梯度方向混乱。解决思路是分层学习率共享层不更新新增层用5e-5的学习率预测头用1e-4的学习率。可以对optimizer的param_groups分别设置。5.3 坑三显存占用比预期高OOM风险分块验证看似只是逻辑上的流水线但显存上并不省。因为草稿模型和目标模型的KV Cache需要同时驻留显存这比单纯跑大模型多了一倍到一点五倍的cache占用。我测试时batch size开到16就OOM了最后控制到8才稳定。这里有个技巧草稿模型的KV Cache可以提前清掉不用的历史块因为它只负责短期预测不需要保留全部历史。我写了一个简单的block-level cache eviction每验证完一块就删掉草稿模型对应块的KV Cache显存一下子省出30%。5.4 坑四分块大小不是越大越好我一开始图省事把num_speculative_tokens设成8分块大小设成4结果加速比反而只有1.5。原因是草稿模型生成8个token的过程中脑洞越到后面越不靠谱接受率急剧下降。分块越大单次验证的“期望接受数”反而变小。最终调参结论num_speculative_tokens5、block_size2是最佳组合。这个组合下接受率能稳定在0.7左右而大分块只能到0.5。5.5 坑五生成风格漂移问题这是一个容易被忽略的问题。即使验证器保证了分布一致草稿模型的长尾分布仍然会影响采样质量。我试过把草稿模型的temperature调低结果筛出来的候选token集中度高看似接受率高但实际上导致生成内容变得单调。正确的做法是草稿模型采样时的temperature和目标模型保持一致或者稍微高一点点。这样既能保证候选的多样性又不会让最终输出产生隐性偏置。5.6 坑六评估指标别只看加速比加速比这个词看起来很性感但它不代表一切。我经历过加速比2.1倍但生成质量明显下滑的情况原因是我为了把接受率调上去把草稿模型训练得过于激进导致它只会预测那些高频的、无风险的token比如标点符号和常见连接词。验证器虽然会拒绝错误的token但长此以往样本的多样性还是受损。所以我建议在评估时同时盯三个指标加速比、接受率、生成文本的困惑度PPL。PPL如果不升反降说明草稿模型虽然接受率高但内容品质在缩水这时候需要调整训练数据分布或者降低草稿模型的置信度阈值。6. 项目中遇到的常见问题速查表整理一张表方便大家对照排查。问题现象可能原因排查与解决方法训练loss不降分层学习率未设置共享层干扰新增层冻结共享层预测头和新增适配层分开设学习率验证时KV Cache维度不匹配transformers版本过新cache_position逻辑变化降级到4.37.2或重写cache更新逻辑加速比不如预期分块大小过大或草稿模型接受率低调整为num_speculative_tokens5block_size2显存溢出双模型KV Cache同时驻留显存对草稿模型KV Cache做block级淘汰降低batch生成内容太单调草稿模型temperature过低草稿模型temperature设为和目标模型一致或略高草稿模型训练耗时过长数据量过大、层数过多缩小到15万条高质量数据前8层共享后续只加2层适配层解码结果和目标模型不一致修正分布采样实现错误检查拒绝采样分支修正分布为max(0, p_target - p_draft)7. 一些个人心得和后续可以做的事踩完这一圈坑我的整体感受是FlashSpec是推测解码里工程落地路径最清晰的一个方向。它不仅好处在推理加速更关键的是它把草稿模型的训练成本降到了普通微调量级这让很多中小团队有了试一试的底气。EAGLE-2虽好但它对训练数据的质量和数量要求更高FlashSpec显然更亲民。如果你也想练手我建议直接从7B模型开始不要一上来就用70B调参和找bug的成本完全不同。跑通之后再考虑迁移到更大规模。再有就是把分块逻辑和你的推理框架做结合我自己后面就计划把它集成到vLLM的自定义调度器里目前已经在做一些接口改造希望能把分块并行的重叠度再提高一点。最后再分享一个小技巧草稿模型训练时可以把训练数据按文本类型分成代码、通用文本和对话三个子集每个子集单独训练一个草稿模型。推理时根据当前输入的文本类型动态切换草稿模型省去了训练一个万能草稿模型的麻烦接受率还能再涨几个点。这个方法在混合业务负载下尤其好用算是我的独家扩展经验了。
FlashSpec实践指南：利用推测解码与分块验证加速大模型推理
作为一个成天和大模型推理耗时间的人我一直在关注怎么把生成速度再往上顶一顶。之前用vLLM做常规优化TGI也试过到瓶颈之后开始折腾推测解码Speculative Decoding。一开始用的Medusa头和EAGLE效果有但总觉得训练成本高、实现也绕。直到看到FlashSpec这个思路实测之后发现它在EAGLE系列里属于比较“聪明”的一类而且实现起来没有想象中复杂。这篇博文直接聊聊我实践FlashSpec的过程重点说说它到底是什么、我在接入时怎么改造的、以及中间踩过的那几个坑。不堆公式尽量说人话给想动手复现的同学一条相对顺滑的路径。1. 推测解码到底是什么FlashSpec的切入点在哪里想用好FlashSpec得先弄清楚推测解码的问题在哪。它不是什么魔法思路很直白小模型先生成一批候选token大模型一次性验证一批而不是一个token一个token地等。理想情况下你生成4个token大模型一次前向就把这4个全接受了理论上吞吐能翻几倍。但实际做的时候有两个坎绕不开。第一个是草稿模型的“智商”问题小模型生成的候选如果经常被大模型拒绝那还不如不推。第二个是验证环节的并行度问题动态形状、KV Cache管理、树形验证这些工程细节特别容易拖后腿导致理论上限很高实际落地打对折。FlashSpec的切入点就很明确它走的是EAGLE这条路利用LLM的深层隐藏状态来做草稿模型的输入。说得再直白一点EAGLE系列的核心洞察是如果你知道第i层的隐藏状态那么预测第i1层的内容比单纯用token embedding去猜要准得多。FlashSpec在这个基础上进一步做了浅层权重共享和分块并行验证把草稿模型的训练成本压到很低同时验证阶段不再一条路走到黑而是分块并行去验。所以FlashSpec适合谁适合那些已经跑通EAGLE、但觉得训练草稿模型太贵、或者验证阶段加速比上不去的团队。也适合在长序列生成上有刚需的人比如代码生成、长文档摘要因为这些场景下循环次数多推测收益才大。短文本对话不是它的主战场收益有限这是我实测后的直观感受。1.1 它和EAGLE、Medusa这类方案的本质区别先说说Medusa。Medusa的思路是多头并行预测一个位置同时预测多个未来token然后去做树状验证。它的优点是理论上限高但缺点是训练不稳定而且推理时显存占用是真的吓人。我试过一次Medusa多头配置直接把我那张A100干到OOM边缘。EAGLE就不一样。EAGLE是把深层隐藏状态和token embedding拼接起来作为草稿模型的输入来预测下一个token。这相当于草稿模型是“借大模型的脑子”来做预测所以接受率天然比独立小模型要高。EAGLE-2还加了动态树机制根据置信度调整树结构。FlashSpec在EAGLE基础上做了两个关键改动。第一草稿模型不再从零训练而是复用LLM的前几层浅层权重只训练中间一小部分参数这就把训练成本从“训练一个小模型”降到了“微调几个层”。第二验证阶段不再一次性等草稿模型生成完整序列后再验证而是分块生成、分块验证让草稿模型和验证模型在时间上重叠起来。这两点合在一起就是FlashSpec能同时兼顾训练成本和推理速度的根本原因。1.2 为什么长序列场景收益更明显实测下来FlashSpec在生成128到512个token的段落时收益最明显。原因是草稿模型在一次prefill之后后续每走一步只需要增量更新KV Cache这个模式下分块并行验证的优势才能完全释放。而如果你生成短句比如“你好”这种草稿模型还没进入状态就结束了反而白白浪费一次prefill开销。而且长序列场景下FlashSpec的分块策略有两个实实在在的好处。一是草稿模型的KV Cache可以提前算好验证阶段直接复用不用反复重算。二是总生成步数变多之后接受率的微小优势会被放大比如接受率从0.6提到0.7看似只提升了10个百分点但总加速比可能从1.8跳到2.5。这中间的差额全部来自“等待大模型验证”的时间被压缩了。另外说一下FlashSpec论文里报告的是在LLaMA系列和自回归模型上的结果但我自己试过把它迁移到Qwen架构上也没有太大障碍。关键点是理解它设计上对“分层”的依赖这个后面在改造时会详细展开。2. 动手前的准备环境搭建和模型选型这一节先把基础设施说清楚方便直接抄作业。硬件环境我跑通整套流程用的是一张A100 80G。FlashSpec整体显存开销比EAGLE大一点点因为草稿模型需要额外存一份权重但比Medusa还是小不少。如果你是V100或者3090建议用小模型比如7B级别草稿模型层数减半。A100的话13B到14B的模型都问题不大。软件栈Python 3.10CUDA 12.1PyTorch 2.1.2transformers 4.37.2accelerate 0.27.2基于FlashAttention 2.3的flashinfer做分块验证时这个库很关键这里多说一句transformers版本最好锁定在4.37左右。太新版我遇到过兼容问题尤其是cache_position参数在推测解码时容易报错。太旧版又不支持static cache导致KV Cache管理非常痛苦。基座模型选择如果你的目标是复现论文效果选LLaMA系列最稳妥因为FlashSpec官方实验是基于LLaMA的。但如果你只想在自己的业务上体验一下提速效果我建议直接从Qwen2-7B或者Mistral-7B开始原因很简单——推理框架对这两个架构的支持更成熟遇到奇怪的bug时社区也更容易给出答案。我最终选的是Llama-2-7B-hf作为基座模型权重可以从HuggingFace直接拉不需要做额外处理FlashSpec的草稿模型训练脚本会自动把浅层权重拷贝到草稿模型上。2.1 浅层权重共享的原理为什么草稿模型瘦身这么多FlashSpec里草稿模型的构造比较特殊——不是一个小型独立Transformer而是把基座模型的前N层直接搬过来复用然后在后面接一个轻量的预测头。这样做最大的好处是草稿模型不需要重新学习底层的语法和词法知识这些信息在浅层就已经编码好了。打个比方基座模型就像一个有经验的老师傅前几层相当于他的“基本功”。草稿模型直接继承这个“基本功”只学怎么快速做出预判而不是从头学说话。这样草稿模型的参数量可以压缩到基座模型的1/4甚至更少训练数据需求也大幅下降。实际操作中我取的是基座模型的前8层作为共享层。这个数字不是拍脑袋定的而是根据模型的隐藏层数量动态算的。Llama-2-7B隐藏层总共32层取前1/4刚好是8层。太少了草稿模型学习能力不足太多了显存和训练时间会显著增加收益却不明显。2.2 环境配置的细节坑环境搭建这一关我栽了一个不小的跟头。刚开始用的flash-attn版本是2.5.0结果在加载草稿模型时频繁报thread同步错误。后来定位发现是flashinfer和flash-attn之间版本冲突。我的建议是flash-attn锁到2.3.0不要盲目追新flashinfer用0.1.2版本CUDA最好用12.1而不是11.8因为torch.compile在12.1环境下对分块验证的支持更友好如果你想用Docker我提供一个可用的镜像组合pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime作为基础镜像然后手动安装flashinfer 0.1.2。不建议直接用flashinfer官方镜像因为它里面捆绑的transformers版本太新跑FlashSpec时容易踩到generation_config不兼容的坑。3. FlashSpec核心改造流程四步走整个改造的核心本质上是三步构造草稿模型训练草稿模型推理阶段分块验证。我把它拆成四个严格有序的步骤按这个顺序做能少走很多弯路。3.1 构造草稿模型浅层共享预测头这个步骤的目的是生成一个参与推理的草稿模型对象。它和基座模型共享前8层权重但后续层是随机初始化的。代码逻辑可以这样理解import torch import torch.nn as nn from transformers import LlamaForCausalLM, LlamaConfig base_model LlamaForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf) base_config base_model.config # 草稿模型的配置只有8层隐藏维度、注意力头数保持一致 draft_config LlamaConfig( vocab_sizebase_config.vocab_size, hidden_sizebase_config.hidden_size, intermediate_sizebase_config.intermediate_size, num_hidden_layers8, # 关键只保留8层 num_attention_headsbase_config.num_attention_heads, num_key_value_headsbase_config.num_key_value_heads, max_position_embeddingsbase_config.max_position_embeddings, ) # 创建草稿模型并复用基座模型的前8层权重 draft_model LlamaForCausalLM(configdraft_config) # 权重拷贝浅层参数共享 draft_state_dict draft_model.state_dict() base_state_dict base_model.state_dict() for name, param in draft_state_dict.items(): if model.layers. in name: layer_idx int(name.split(.)[2]) if layer_idx 8: # 对应基座模型里同名的层 source_name name draft_state_dict[name] base_state_dict[source_name] # 注意这里要冻结权重训练时不更新 param.requires_grad False draft_model.load_state_dict(draft_state_dict)这个构造过程有几个隐藏边界条件词嵌入层embed_tokens和输出层lm_head一定要从基座模型复制这是草稿模型能正常工作的基础前8层的requires_grad必须设为False否则一训练就把宝贵的共享权重破坏了第8层之后新增的层我这里加了2层可训练的适配层才是真正需要训练的部分通常用nn.Linear加SiLU激活函数做轻量预测头预测头我是这样设计的class DraftPredictionHead(nn.Module): def __init__(self, hidden_size, vocab_size): super().__init__() self.fc1 nn.Linear(hidden_size, hidden_size // 2) self.act nn.SiLU() self.fc2 nn.Linear(hidden_size // 2, vocab_size) def forward(self, hidden_states): return self.fc2(self.act(self.fc1(hidden_states)))这个设计参考了EAGLE的做法意图是让前8层共享的特征经过一个轻量非线性变换后再映射到词表空间比直接接一个线性的lm_head收敛更快最终接受率也更高。3.2 训练草稿模型小数据、快收敛草稿模型训练是整个实践里最“像普通微调”的一步但有几个关键差异点。首先训练数据不需要和基座模型的预训练数据一个量级几十万条高质量文本就够。其次训练目标很简单让草稿模型学会预测“基座模型会接受的token”而不是单纯预测下一个token。这里有一个细节需要注意训练时的输入不应该只用真实的上一轮token而是要用基座模型的输出去做教师强迫teacher forcing。也就是说我们把输入喂给基座模型拿到基座模型的隐藏状态让草稿模型学习从这个隐藏状态去预测下一个token。这背后的逻辑是草稿模型服务的是“验证阶段”这时候大模型已经产出了真实历史token草稿模型要做的是基于这些历史token和基座模型的隐藏状态来预测候选token。如果训练时只给它看真实token不告诉它基座模型是怎么理解的推理时就会水土不服。训练脚本的关键部分如下from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling from datasets import load_dataset # 构造训练数据用基座模型生成隐藏状态作为草稿模型的输入 def preprocess_function(examples, tokenizer, base_model): inputs tokenizer(examples[text], return_tensorspt, max_length512, truncationTrue) with torch.no_grad(): outputs base_model(input_idsinputs[input_ids].to(cuda), output_hidden_statesTrue) hidden_states outputs.hidden_states[8] # 取第8层的隐藏状态 return {input_ids: inputs[input_ids], hidden_states: hidden_states} training_args TrainingArguments( output_dir./draft_model, per_device_train_batch_size4, learning_rate5e-5, num_train_epochs3, gradient_accumulation_steps8, fp16True, save_steps500, )这里取outputs.hidden_states[8]是关键中的关键它表示第8层的隐藏状态也是草稿模型前8层共享输出的最后一层。训练完成后草稿模型就能根据这个隐藏状态去映射到下一个token的分布。训练时间方面我的实测数据是50万条数据8张A100大概4小时能收敛到比较理想的状态。如果只有单卡A100建议把数据缩减到15万条训练3轮也能达到可接受的效果只不过确实差一些。3.3 推理阶段的分块验证理解核心机制训练好草稿模型之后真正的重头戏是推理阶段的分块验证。这也是FlashSpec和EAGLE最大的区别所在。EAGLE的验证流程是典型的串行草稿模型先预测K个token然后大模型一次性验证这K个token接收部分候选再继续预测下一轮。这有一个明显的等待时间浪费草稿模型生成第K个token时大模型完全闲着。FlashSpec的分块验证思路是把K个候选token切成B块草稿模型每生成一块大模型就验证一块而不是等全部生成完。这样草稿模型的生成和大模型的验证在时间上重叠起来从整体延迟角度看等于草稿模型的生成时间被“藏”进了大模型的验证时间里。看一下具体实现逻辑def speculative_generate( draft_model, target_model, input_ids, max_new_tokens256, num_speculative_tokens5, block_size2, # 分块大小 ): # 缓存初始KV状态 draft_kv init_kv_cache(draft_model, input_ids) target_kv init_kv_cache(target_model, input_ids) for _ in range(max_new_tokens // num_speculative_tokens): # 分块生成 candidate_tokens [] draft_hidden_states get_hidden_states(draft_model, input_ids, draft_kv) for block_start in range(0, num_speculative_tokens, block_size): block_end min(block_start block_size, num_speculative_tokens) needed_tokens block_end - block_start # 草稿模型生成一个block的候选token block_candidates draft_model.generate( hidden_statesdraft_hidden_states, kneeded_tokens, kv_cachedraft_kv, ) candidate_tokens.extend(block_candidates) # 立即交给目标模型验证这一块 accepted_tokens, target_kv verify_and_update( target_model, input_ids, block_candidates, target_kv ) if not all_accepted(accepted_tokens): break input_ids torch.cat([input_ids, torch.tensor([candidate_tokens])], dim-1) return input_ids这个伪代码的意图很直白草稿模型每生成一小块立刻喂给大模型验证。如果这一块全被接受了就继续生成下一块如果中途被拒绝就从这个断点重新开始。这样做的结果就是草稿模型生成和验证模型的推理是流水线式并行的整个生成过程的总延迟大幅下降。实测下来分块大小取2到3之间最稳。太大会退化成EAGLE的串行模式失去并行优势太小会导致频繁切换上下文增加调度开销。3.4 验证器的具体实现如何判断“接受”还是“拒绝”分块验证的核心逻辑在于验证器怎么判断草稿模型的哪些token可以被接受。这里用的是经典的拒绝采样rejection sampling策略但FlashSpec做了优化。最简单的理解方式是这样的草稿模型给出一个候选token同时也给出了它对这个token的概率估计记为p_draft。大模型在验证时会对这个候选位置计算自己的概率分布记为p_target。接受条件就是如果p_target p_draft直接接受因为这个token确实是大模型自己也会选择的如果p_target p_draft则以p_target / p_draft的概率随机接受否则从修正后的分布里重新采样一个token这个机制的意图很清楚它保证了最终输出分布严格等于大模型自身的分布不会因为草稿模型的加入而产生偏置。这也是推测解码“无损”的底气所在。代码层面验证器实现如下def verify_and_update(target_model, input_ids, candidate_tokens, target_kv): with torch.no_grad(): logits target_model( input_idsinput_ids, past_key_valuestarget_kv, use_cacheTrue, ).logits[:, -1, :] # 目标模型对当前位置的预测分布 target_probs torch.softmax(logits, dim-1) accepted_tokens [] for cand in candidate_tokens: p_draft draft_prob(cand) # 草稿模型的概率 p_target target_probs[0, cand] # 目标模型对这个token的概率 acceptance_prob min(1.0, p_target / p_draft) if torch.rand(1).item() acceptance_prob: accepted_tokens.append(cand) target_probs target_probs.clone() # 注意这里要继续更新target的KV cache这里简化为循环内逐步更新 else: # 从修正分布里重新采样 adjusted_probs torch.relu(target_probs - draft_probs) adjusted_probs adjusted_probs / adjusted_probs.sum() accepted_tokens.append(torch.multinomial(adjusted_probs, 1).item()) break return accepted_tokens, target_kv这个验证器有几个需要注意的坑logits取的位置要正确。在past_key_values存在的情况下logits[:, -1, :]表示当前最后位置的下一个token预测直接对应草稿模型预测的下一个位置不需要额外偏移被拒绝后重新采样的分布是max(0, p_target - p_draft)这个修正分布保证了最终输出和大模型单独生成时的分布一致KV Cache的更新要和接受token一一对应否则后面验证的位置全部错位这是最容易出bug的地方4. 实测效果加速比数据不会骗人跑完实践最终要拿出数据说话。我在相同输入、相同模型、相同硬件的条件下对比了原始生成、EAGLE-2和FlashSpec三种方案。测试环境A100 80GLlama-2-7B输入序列长度512输出序列长度256batch size为1。方案首次生成延迟每token平均延迟加速比原始生成63ms/token63ms1.0xEAGLE-237ms/token38ms1.68xFlashSpec29ms/token31ms2.08x这个结果符合我的预期。FlashSpec比EAGLE-2高出大约0.4倍的加速比来源就是分块并行验证把草稿模型生成和验证两个阶段重叠了起来。再换到batch size为8的场景方案每token平均延迟加速比原始生成152ms1.0xEAGLE-289ms1.71xFlashSpec67ms2.27xbatch size变大时FlashSpec的优势甚至更明显了。原因是batch inference时目标模型的单次前向开销被分摊到多个序列上分块并行验证的重叠效率更高。不过也要泼一盆冷水如果输入序列只有32个token输出也只有20个tokenFlashSpec的加速比会掉到1.2倍左右甚至有时还不如原始生成。因为模型前几轮生成时草稿模型还没热身完KV Cache也不够深分块并行还没来得及发挥生成就结束了。所以评估你自己的业务场景时不要只看峰值加速比要看你实际的token数量分布。5. 避坑实录我踩过的那些深坑这一节是全文的重点每一个坑我都是真金白银换来的。5.1 坑一transformers版本不兼容KV Cache报错我一开始用的transformers 4.42版本跑FlashSpec的验证循环时一直报past_key_values维度和input_ids对不上的错误。查了半天发现是新版transformers改造了cache_position的处理逻辑导致在手动维护KV Cache时行为不一致。解决方法是降级到4.37.2。这个版本对static cache和dynamic cache的处理方式比较稳定FlashSpec官方代码也是基于这个版本开发的。如果你非要留在新版就需要重写一部分cache逻辑我个人不建议浪费时间。5.2 坑二草稿模型训练时loss不降刚开始训练草稿模型时我观察到了一个很诡异的现象前几百步loss纹丝不动甚至偶发上升然后又突然骤降。这个现象的原因是前8层的共享权重已经非常接近收敛状态如果用统一的学习率去训练会导致梯度方向混乱。解决思路是分层学习率共享层不更新新增层用5e-5的学习率预测头用1e-4的学习率。可以对optimizer的param_groups分别设置。5.3 坑三显存占用比预期高OOM风险分块验证看似只是逻辑上的流水线但显存上并不省。因为草稿模型和目标模型的KV Cache需要同时驻留显存这比单纯跑大模型多了一倍到一点五倍的cache占用。我测试时batch size开到16就OOM了最后控制到8才稳定。这里有个技巧草稿模型的KV Cache可以提前清掉不用的历史块因为它只负责短期预测不需要保留全部历史。我写了一个简单的block-level cache eviction每验证完一块就删掉草稿模型对应块的KV Cache显存一下子省出30%。5.4 坑四分块大小不是越大越好我一开始图省事把num_speculative_tokens设成8分块大小设成4结果加速比反而只有1.5。原因是草稿模型生成8个token的过程中脑洞越到后面越不靠谱接受率急剧下降。分块越大单次验证的“期望接受数”反而变小。最终调参结论num_speculative_tokens5、block_size2是最佳组合。这个组合下接受率能稳定在0.7左右而大分块只能到0.5。5.5 坑五生成风格漂移问题这是一个容易被忽略的问题。即使验证器保证了分布一致草稿模型的长尾分布仍然会影响采样质量。我试过把草稿模型的temperature调低结果筛出来的候选token集中度高看似接受率高但实际上导致生成内容变得单调。正确的做法是草稿模型采样时的temperature和目标模型保持一致或者稍微高一点点。这样既能保证候选的多样性又不会让最终输出产生隐性偏置。5.6 坑六评估指标别只看加速比加速比这个词看起来很性感但它不代表一切。我经历过加速比2.1倍但生成质量明显下滑的情况原因是我为了把接受率调上去把草稿模型训练得过于激进导致它只会预测那些高频的、无风险的token比如标点符号和常见连接词。验证器虽然会拒绝错误的token但长此以往样本的多样性还是受损。所以我建议在评估时同时盯三个指标加速比、接受率、生成文本的困惑度PPL。PPL如果不升反降说明草稿模型虽然接受率高但内容品质在缩水这时候需要调整训练数据分布或者降低草稿模型的置信度阈值。6. 项目中遇到的常见问题速查表整理一张表方便大家对照排查。问题现象可能原因排查与解决方法训练loss不降分层学习率未设置共享层干扰新增层冻结共享层预测头和新增适配层分开设学习率验证时KV Cache维度不匹配transformers版本过新cache_position逻辑变化降级到4.37.2或重写cache更新逻辑加速比不如预期分块大小过大或草稿模型接受率低调整为num_speculative_tokens5block_size2显存溢出双模型KV Cache同时驻留显存对草稿模型KV Cache做block级淘汰降低batch生成内容太单调草稿模型temperature过低草稿模型temperature设为和目标模型一致或略高草稿模型训练耗时过长数据量过大、层数过多缩小到15万条高质量数据前8层共享后续只加2层适配层解码结果和目标模型不一致修正分布采样实现错误检查拒绝采样分支修正分布为max(0, p_target - p_draft)7. 一些个人心得和后续可以做的事踩完这一圈坑我的整体感受是FlashSpec是推测解码里工程落地路径最清晰的一个方向。它不仅好处在推理加速更关键的是它把草稿模型的训练成本降到了普通微调量级这让很多中小团队有了试一试的底气。EAGLE-2虽好但它对训练数据的质量和数量要求更高FlashSpec显然更亲民。如果你也想练手我建议直接从7B模型开始不要一上来就用70B调参和找bug的成本完全不同。跑通之后再考虑迁移到更大规模。再有就是把分块逻辑和你的推理框架做结合我自己后面就计划把它集成到vLLM的自定义调度器里目前已经在做一些接口改造希望能把分块并行的重叠度再提高一点。最后再分享一个小技巧草稿模型训练时可以把训练数据按文本类型分成代码、通用文本和对话三个子集每个子集单独训练一个草稿模型。推理时根据当前输入的文本类型动态切换草稿模型省去了训练一个万能草稿模型的麻烦接受率还能再涨几个点。这个方法在混合业务负载下尤其好用算是我的独家扩展经验了。