投机采样加速大模型推理：原理、EAGLE与草稿树实战调优

1. 为什么投机采样值得单独拿出来讲大模型推理优化这个话题能聊的东西太多了。量化、KV Cache 管理、算子融合、连续批处理、PagedAttention每一个方向都能写一整篇。但投机采样Speculative Sampling是个特别的存在——它不改变模型结构不损失输出质量却能在特定场景下把解码速度拉高两到三倍。我第一次在生产环境里跑通这套方案的时候看着 tokens/s 从 35 跳到 90 多说实话有点不太敢相信。这个系列的第一篇聊的是量化主要解决显存占用和计算密度的问题。但量化有个天花板它优化的是单次前向传播的效率而自回归解码的本质瓶颈在于——每生成一个 token 就要完整跑一遍模型。你量化做得再好7B 模型生成 100 个 token 还是要跑 100 次前向。投机采样的思路完全不同它想办法让这 100 次前向变成 30 次甚至更少。适合谁来读这篇如果你已经在本地或者服务器上部署过大模型用过 vLLM、llama.cpp、TensorRT-LLM 这类推理框架对 KV Cache、batch size、temperature 这些概念不陌生那这篇的内容你可以直接拿去用。如果你还在选模型、配环境的阶段建议先把基础推理跑通再回来看不然有些细节会缺少体感。我下面会从原理讲到实操包括 EAGLE 和草稿树这两种主流方案的对比、参数怎么调、踩过哪些坑。代码层面以 Python 和 HuggingFace 生态为主但思路是通用的你用什么框架都能套。2. 投机采样的核心思路拆解2.1 自回归解码到底慢在哪要理解投机采样为什么有效得先看清楚标准解码流程的瓶颈。假设你有一个 7B 的模型生成一个 token 的流程是把当前所有 token 的 KV 送进模型做一次完整的前向传播得到下一个 token 的概率分布采样出一个 token追加到序列末尾然后重复。这个过程有两个特点。第一计算是串行的——第 n1 个 token 必须等第 n 个 token 生成完才能开始。第二每次前向传播的计算量其实很小因为 batch size 通常是 1单用户场景GPU 的并行计算能力远远没有被吃满。换句话说GPU 大部分时间在等内存带宽而不是在算。这里有个常见的误解很多人以为大模型推理慢是因为计算量大。实际上在 decode 阶段瓶颈往往是显存带宽——每生成一个 token 都要把整个模型的权重从显存读一遍。7B 模型 FP16 精度下大约 14GB 权重每次前向都要读一遍这个带宽开销是固定的。所以核心矛盾是串行解码 低 GPU 利用率。投机采样就是冲着这个矛盾去的。2.2 用一个小模型来加速大模型投机采样的核心思想可以用一句话概括用一个便宜的小模型先猜出接下来几个 token然后让大模型一次性验证这些猜测对不对。打个比方。你是一个资深翻译大模型每翻译一个词都要停下来想很久。现在给你配了一个实习生小模型实习生先把接下来五个词都翻出来你只需要扫一眼确认哪些是对的、哪个错了、错的应该改成什么。如果实习生猜对了四个你一次前向就确认了四个 token效率直接翻四倍。这里的关键在于验证过程只需要一次前向传播。你把原始序列加上实习生猜的 k 个 token 一起送进大模型大模型会输出每个位置的概率分布。然后从第一个猜测位置开始比对如果大模型采样出的 token 和实习生猜的一致就接受遇到第一个不一致的就拒绝后面的所有猜测用大模型自己的结果替换。数学上可以证明这种接受-拒绝机制产生的输出分布和直接用大模型逐个生成是完全一样的。这不是近似是无损加速。这一点非常重要——意味着你不用担心投机采样会降低输出质量。2.3 投机采样、EAGLE、草稿树的区别在哪这三个概念经常被混在一起说但它们的层次不太一样。投机采样是框架性的思路核心就是猜-验证这个范式。至于用什么来猜、怎么猜、猜多少有不同的实现方案。EAGLE 是一种具体的投机采样实现它的特点是不用一个独立的小模型来猜而是在大模型本身的特征上训练一个轻量的 draft head。这个 head 复用大模型的隐藏层状态预测下一个 token 的特征表示再通过大模型的 LM head 输出 token。因为复用了大模型已经算出来的特征EAGLE 的猜测质量比独立小模型高很多接受率通常能到 0.7-0.8 甚至更高。草稿树Draft Tree则是另一种思路。它不只猜一条序列而是同时猜多条候选路径形成一棵树。验证的时候大模型可以并行验证树上的多个分支一次前向能接受更多 token。代价是验证时的计算量会增加因为要处理树结构上的注意力掩码。方案猜测来源典型接受率额外显存实现复杂度独立小模型同系列小模型0.4-0.6小模型权重低EAGLE大模型特征轻量head0.7-0.85head权重特征缓存中草稿树多分支候选0.6-0.8树结构KV高选哪个我的经验是如果你用的是 llama.cpp 或者 vLLM它们已经内置了投机采样支持直接用独立小模型方案最省事。如果你在追求极致性能愿意花时间训练 draft headEAGLE 的收益最大。草稿树目前在生产环境用得相对少主要是实现复杂度高调试起来比较痛苦。3. 实操环境搭建与模型准备3.1 硬件和基础环境我这边的测试环境是一台单卡 A100 80G 的机器CUDA 12.1PyTorch 2.1.0。如果你用的是消费级显卡比如 4090 24G跑 7B 模型加一个 1B 左右的 draft 模型也是够的。显存估算大概是主模型 FP16 占 14Gdraft 模型 FP16 占 2GKV Cache 看序列长度2048 上下文大概 1-2G。总共 18G 左右24G 卡能跑。软件层面我建议用 vLLM 或者 llama.cpp 作为推理后端。vLLM 对投机采样的支持比较成熟配置也简单。llama.cpp 更适合本地部署场景CPUGPU 混合推理也能用。下面我以 vLLM 为主来讲llama.cpp 的配置我会在后面的章节补充。pip install vllm0.4.2 pip install transformers4.40.0 pip install torch2.1.0版本这块要注意vLLM 不同版本对投机采样的 API 有变化。0.4.x 之后接口基本稳定了用speculative_model参数指定 draft 模型就行。3.2 选一对合适的模型组合这是最容易踩坑的地方。不是随便找个小模型就能当 draft 用的。几个原则同系列优先。比如主模型用 Qwen2.5-7B-Instructdraft 就用 Qwen2.5-0.5B 或者 1.5B。同系列模型的 tokenizer 一致词表分布接近猜测的接受率会高很多。我试过用 Llama 的小模型给 Qwen 的大模型做 draft接受率直接掉到 0.2 以下基本没有加速效果。词表必须一致。这个不用多解释词表不一样的话draft 模型猜出来的 token id 在主模型里对应的完全是另一个词验证根本没法做。draft 模型不能太小。0.5B 是个下限再小的话猜测质量太差接受率低到没有意义。1B-1.5B 是比较甜的点猜测质量够用推理开销也不大。from vllm import LLM, SamplingParams llm LLM( modelQwen/Qwen2.5-7B-Instruct, speculative_modelQwen/Qwen2.5-1.5B-Instruct, num_speculative_tokens5, tensor_parallel_size1, gpu_memory_utilization0.9, )num_speculative_tokens这个参数控制每次猜多少个 token。设太小加速不明显设太大验证时浪费的计算多。后面我会详细讲怎么调。3.3 验证环境是否跑通跑通之后先做个简单的基准测试确认加速确实生效了。import time prompts [请用中文解释什么是机器学习要求通俗易懂。] sampling_params SamplingParams(temperature0.7, max_tokens256) # 预热 llm.generate(prompts, sampling_params) # 计时 start time.time() outputs llm.generate(prompts, sampling_params) elapsed time.time() - start total_tokens sum(len(o.outputs[0].token_ids) for o in outputs) print(f生成 {total_tokens} tokens耗时 {elapsed:.2f}s速度 {total_tokens/elapsed:.1f} tokens/s)把speculative_model去掉再跑一遍对比一下速度。正常情况下应该有 1.5-2.5 倍的提升。如果没有提升甚至更慢大概率是 draft 模型选得不对或者num_speculative_tokens设得太激进。4. 关键参数调优与性能分析4.1 num_speculative_tokens 怎么定这个参数直接决定了每次猜测的 token 数量也决定了验证时的计算量。设成 k意味着每次前向要处理 k1 个位置原始位置加上 k 个猜测位置。理论上如果接受率是 α那么每次验证期望接受的 token 数是 (1-α^(k1))/(1-α)。当 k 增大时这个期望值会趋近于 1/(1-α)但验证成本也线性增长。所以存在一个最优的 k。我实测下来对于 EAGLE 方案k5 到 7 比较合适。对于独立小模型方案k3 到 5 就够了因为接受率本身低一些猜太多浪费。# 不同 k 值的对比测试 for k in [1, 3, 5, 7, 10]: llm LLM( modelQwen/Qwen2.5-7B-Instruct, speculative_modelQwen/Qwen2.5-1.5B-Instruct, num_speculative_tokensk, ) # ... 跑基准测试我跑出来的一组数据A100Qwen2.5-7B 1.5B drafttemperature0.7k 值接受率速度 (tokens/s)相对加速比0基线-381.0x10.72521.37x30.68782.05x50.65912.39x70.61882.32x100.55792.08x可以看到 k5 是个拐点再往上加收益就下降了。当然这个数字跟具体模型、任务类型都有关系你需要在自己的场景下测。4.2 温度参数对接受率的影响temperature 越高采样越随机draft 模型猜中的概率就越低。这个很好理解——如果每次都是从概率分布里随机抽那猜中的概率自然低。实测数据temperature0 时接受率能到 0.85 以上temperature0.7 时降到 0.65 左右temperature1.0 时只有 0.5 出头。这意味着什么如果你的应用场景是代码生成、结构化输出这种确定性强的任务投机采样的加速效果会非常好。如果是创意写作、头脑风暴这种需要高 temperature 的场景加速比会打折扣但通常也还有 1.5 倍左右的提升。有个小技巧如果你的场景允许把 temperature 设低一点比如 0.3-0.5既能保持一定的多样性又能让投机采样发挥更好的效果。当然这要看具体需求不能为了加速牺牲输出质量。4.3 显存和吞吐的权衡投机采样会额外占用显存主要是 draft 模型的权重和它自己的 KV Cache。在显存紧张的情况下你可能需要调低gpu_memory_utilization或者减小max_model_len。另一个容易被忽略的点是投机采样在 batch size 较大的时候收益会下降。因为大 batch 本身就能把 GPU 利用率拉高decode 阶段的瓶颈从显存带宽转移到了计算。这时候投机采样带来的额外计算反而可能成为负担。我的建议是batch size 在 1-8 的时候用投机采样收益最明显。如果 batch size 超过 16可以考虑关掉投机采样用连续批处理来提升吞吐。5. 常见问题与排查实录5.1 加速效果不明显甚至变慢这是最常见的问题。排查思路按优先级来第一检查 draft 模型和主模型的 tokenizer 是否一致。这个用AutoTokenizer加载两个模型对比vocab_size和几个特殊 token 的 id 就能确认。第二看接受率。vLLM 的日志里会输出接受率相关的指标。如果接受率低于 0.3基本可以判断是模型组合的问题。第三检查num_speculative_tokens是否设得太大。设成 10 以上验证成本会吃掉大部分收益。第四确认你的场景是不是 batch size 太大。前面说过大 batch 下投机采样收益会下降。5.2 输出质量出现异常理论上投机采样是无损的输出分布和标准解码完全一致。但实际中如果实现有 bug可能会出现重复、乱码、截断等问题。我遇到过一次draft 模型和主模型的 padding token id 不一致导致验证时位置对齐出错输出全是乱码。解决办法是显式指定pad_token_id确保两个模型一致。还有一种情况是 KV Cache 的索引对不上。这个在自定义实现的时候容易出用 vLLM 这类成熟框架基本不会遇到。5.3 显存溢出OOM投机采样的显存开销比标准推理大主要是多了一个模型。如果 OOM 了可以按这个顺序调降低gpu_memory_utilization比如从 0.9 降到 0.85减小max_model_len比如从 4096 降到 2048换更小的 draft 模型用 FP8 或者 INT8 量化 draft 模型注意draft 模型量化后猜测质量会下降接受率可能降低 10-15 个百分点。需要权衡一下。5.4 常见问题速查表现象可能原因排查方法解决方式无加速效果draft 模型不匹配对比 tokenizer换同系列模型加速效果差k 值设置不当测试不同 k 值调到 3-5输出乱码padding id 不一致检查 pad_token_id显式指定一致的值OOM显存不足监控显存占用降 gpu_memory_utilization接受率低temperature 太高查看日志接受率适当降低 temperature大 batch 下变慢计算瓶颈转移对比不同 batch sizebatch16 时关闭投机采样6. 进阶EAGLE 方案的实际落地6.1 EAGLE 和独立小模型的本质区别前面说的独立小模型方案draft 模型和大模型是独立的它自己完整地跑一遍前向输出 token 给大模型验证。EAGLE 不一样它不跑完整的前向而是复用大模型已经算出来的隐藏层特征用一个很轻的 head 来预测下一个 token 的特征。这个区别带来的好处是draft 的开销极小几乎可以忽略不计。同时因为复用了大模型的特征猜测质量更高。代价是你需要额外训练这个 draft head不能直接拿现成的模型用。EAGLE 的论文里报告的数据是在 Vicuna 7B 上EAGLE 能达到 2.7-3.5 倍的加速接受率 0.75 以上。我自己的实测数据略低一些大概 2.2-2.8 倍但确实比独立小模型方案好不少。6.2 训练 draft head 的流程如果你决定走 EAGLE 路线训练流程大概是这样的第一步准备数据。用你的目标场景数据或者通用的指令跟随数据让主模型生成一批输出保存每步的隐藏层状态。数据量大概 10万-50万条就够。第二步定义 draft head 结构。EAGLE 的 head 通常是一个单层或者两层的 Transformer decoder输入是主模型最后一层的隐藏状态输出是预测的下一个 token 的特征。第三步训练。损失函数是预测特征和真实特征的 MSE加上 token 预测的交叉熵。学习率用 1e-4 左右训练 1-2 个 epoch。第四步集成到推理框架。vLLM 从 0.4.x 开始支持 EAGLE需要把训练好的 head 权重转成对应的格式。# EAGLE head 的简化结构示意 import torch.nn as nn class EagleHead(nn.Module): def __init__(self, hidden_size, vocab_size, num_layers1): super().__init__() self.layers nn.ModuleList([ nn.TransformerDecoderLayer(d_modelhidden_size, nhead8, batch_firstTrue) for _ in range(num_layers) ]) self.lm_head nn.Linear(hidden_size, vocab_size, biasFalse) def forward(self, hidden_states): x hidden_states for layer in self.layers: x layer(x, x) return self.lm_head(x)实际训练的时候还有很多细节比如特征归一化、位置编码的处理、多步预测的展开方式等等。这块展开讲能写一整篇我这里先给个框架后面有机会单独聊。6.3 草稿树的实现思路草稿树的核心是在 draft 阶段不只生成一条序列而是生成一棵树。比如每个位置生成 2-3 个候选形成一棵深度为 k 的树。验证的时候大模型一次前向处理整棵树用树形注意力掩码确保每个节点只能看到自己的祖先。这样做的好处是即使某条路径猜错了其他分支可能猜对整体接受率更高。代价是验证时的计算量增加因为要处理更多的位置。实现上草稿树的难点在于注意力掩码的构造和 KV Cache 的管理。vLLM 目前对草稿树的支持还在实验阶段生产环境用得不多。如果你对这个方向感兴趣建议先看 Medusa 和 SpecInfer 这两篇论文它们对树形验证的工程实现讲得比较清楚。7. 我踩过的坑和几条实用建议第一个坑是模型版本不匹配。我有一次用 Qwen2.5-7B 做主模型draft 用了 Qwen2-1.5B看着名字差不多实际上 tokenizer 有细微差异导致接受率极低。后来换成同版本的 Qwen2.5-1.5B 就正常了。所以选 draft 模型的时候一定要确认是同一个系列、同一个版本。第二个坑是忽略了 warmup。第一次推理的时候draft 模型和主模型都要加载权重、编译 CUDA kernel速度会慢很多。做基准测试的时候一定要先跑几轮预热不然数据会很难看。第三个坑是盲目追求高 k 值。我一开始觉得 k 越大越好设成 10结果速度反而比 k5 慢。后来才想明白验证成本是线性增长的而接受率的提升是递减的存在一个最优值。几条实用建议先用独立小模型方案跑通确认加速效果再考虑上 EAGLE接受率低于 0.5 的时候优先检查模型匹配问题而不是调参数生产环境部署前一定要在你的真实数据上测接受率和加速比公开 benchmark 的数字只能参考如果显存紧张draft 模型可以量化但主模型尽量保持 FP16不然输出质量会受影响投机采样和连续批处理可以叠加使用但要注意 batch size 大了之后收益会下降最后分享一个我常用的调试技巧在 vLLM 启动的时候加上--disable-log-requests之外的日志参数把接受率、平均接受长度这些指标打出来。观察这些指标的变化比单纯看 tokens/s 更能定位问题。接受率突然下降通常是输入数据的分布变了或者 draft 模型遇到了不擅长的任务类型。这个方向后续还可以往几个方向扩展一是结合量化把 draft 模型压到 INT4进一步降低显存开销二是针对特定领域微调 draft 模型把接受率再往上推三是探索自适应 k 值根据当前上下文动态调整猜测长度。这些我还在折腾有结果了再写一篇。