作为一个成天和大模型推理耗时间的人，我一直在关注怎么把生成速度再往上顶一顶。之前用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层的隐藏状态，那么预测第i+1层的内容比单纯用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.10
- CUDA 12.1
- PyTorch 2.1.2
- transformers 4.37.2
- accelerate 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_size=base_config.vocab_size, hidden_size=base_config.hidden_size, intermediate_size=base_config.intermediate_size, num_hidden_layers=8, # 关键：只保留8层 num_attention_heads=base_config.num_attention_heads, num_key_value_heads=base_config.num_key_value_heads, max_position_embeddings=base_config.max_position_embeddings, ) # 创建草稿模型，并复用基座模型的前8层权重 draft_model = LlamaForCausalLM(config=draft_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_tensors="pt", max_length=512, truncation=True) with torch.no_grad(): outputs = base_model(input_ids=inputs["input_ids"].to("cuda"), output_hidden_states=True) 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_size=4, learning_rate=5e-5, num_train_epochs=3, gradient_accumulation_steps=8, fp16=True, save_steps=500, )这里取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_tokens=256, num_speculative_tokens=5, block_size=2, # 分块大小 ): # 缓存初始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_states=draft_hidden_states, k=needed_tokens, kv_cache=draft_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_ids=input_ids, past_key_values=target_kv, use_cache=True, ).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 80G，Llama-2-7B，输入序列长度512，输出序列长度256，batch size为1。
| 方案 | 首次生成延迟 | 每token平均延迟 | 加速比 |
|---|---|---|---|
| 原始生成 | 63ms/token | 63ms | 1.0x |
| EAGLE-2 | 37ms/token | 38ms | 1.68x |
| FlashSpec | 29ms/token | 31ms | 2.08x |
这个结果符合我的预期。FlashSpec比EAGLE-2高出大约0.4倍的加速比，来源就是分块并行验证把草稿模型生成和验证两个阶段重叠了起来。
再换到batch size为8的场景：
| 方案 | 每token平均延迟 | 加速比 |
|---|---|---|
| 原始生成 | 152ms | 1.0x |
| EAGLE-2 | 89ms | 1.71x |
| FlashSpec | 67ms | 2.27x |
batch size变大时，FlashSpec的优势甚至更明显了。原因是batch inference时，目标模型的单次前向开销被分摊到多个序列上，分块并行验证的重叠效率更高。
不过也要泼一盆冷水：如果输入序列只有32个token，输出也只有20个token，FlashSpec的加速比会掉到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_tokens=5、block_size=2是最佳组合。这个组合下接受率能稳定在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_tokens=5，block_size=2 |
| 显存溢出 | 双模型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的自定义调度器里，目前已经在做一些接口改造，希望能把分块并行的重叠度再提高一点。
最后再分享一个小技巧：草稿模型训练时，可以把训练数据按文本类型分成代码、通用文本和对话三个子集，每个子集单独训练一个草稿模型。推理时根据当前输入的文本类型动态切换草稿模型，省去了训练一个万能草稿模型的麻烦，接受率还能再涨几个点。这个方法在混合业务负载下尤其好用，算是我的独家扩展经验了。
