KV Cache 优化（四）：Speculative Decoding 投机解码

前面三篇都在缩小或筛选 KV Cache。这一篇优化的是 Decode 的串行性这一瓶颈。

普通自回归生成一次只产生一个 token：

xₜ → xₜ₊₁ → xₜ₊₂ → xₜ₊₃

每一步都要调用一次大模型，而大模型单 token Decode 的计算规模很小，却要读取大量权重和 KV Cache。Speculative Decoding（投机解码）采用“先猜、后验”的办法：

让一个便宜的小模型先提出多个候选，再让大模型一次并行验证；验证通过的连续前缀直接接受。

它不是把大模型变成了真正并行生成器，而是把多个候选位置的验证组织成一次前向计算。

Speculative Decoding 的草稿生成、并行验证和 KV Cache 提交流程

1. 普通 Decode 为什么慢？

假设当前上下文是：

今 天 天 气

大模型要生成 4 个新 token，需要反复执行：

上下文             → 大模型 → “很”
上下文 +“很”       → 大模型 → “适”
上下文 +“很适”     → 大模型 → “合”
上下文 +“很适合”   → 大模型 → “散步”

每次只有一个新 token，线性层形状接近：

$[1,H]@[H,H]\rightarrow[1,H]$

这类小批量计算的算术强度不高，模型权重和历史 KV 的读取占比很大。

2. Draft Model 先提出候选

Speculative Decoding 引入一个更小、更快的 Draft Model。它仍然是自回归生成，但因为模型小，连续生成几步的成本较低。

例如 Draft Model 给出：

第 1 步：很
第 2 步：适
第 3 步：合
第 4 步：散步

这里假定“很、适、合、散步”分别是四个 token，实际边界由 tokenizer 决定。标准双模型方案通常采用兼容的词表与 tokenizer。它们不是四个互相独立的预测，而是依次基于前一个草稿 token：

$x_{t+1}^d\sim q(\cdot\mid x_{\le t}),\quad x_{t+2}^d\sim q(\cdot\mid x_{\le t},x_{t+1}^d)$

其中 q 是 Draft Model 的概率分布。

3. Target Model 一次验证多个位置

接下来把原上下文和整段候选一起交给大模型：

原上下文：[今 天 天 气]
候选：    [很 适 合 散步]

Target Model 使用因果 Mask，一次得到多个位置的 logits：

$[k,H]@[H,H]\rightarrow[k,H]$

这里 k=4。四个位置在 GPU 上可以并行计算，但它们的逻辑依赖仍然是因果的。验证分布与输入位置之间有一个错位：

原上下文最后位置的分布 p₁：验证“很”
输入“很”后得到的分布 p₂：验证“适”
输入“适”后得到的分布 p₃：验证“合”
输入“合”后得到的分布 p₄：验证“散步”

实现可能复用上一轮已有的上下文末尾 logits，也可能在验证调用中一并组织这些输出。无论缓存接口怎样设计，数学对应关系都是“某位置的输出分布预测下一个 token”。核心是：一趟 Target 前向同时获得连续候选位置的条件分布。

这次计算更像一个短 Prefill，而不是普通单 token Decode：

单 token Decode：序列维度约为 1
Speculative Verify：候选序列维度约为 k

4. 接受和拒绝

最简单的直觉是：如果大模型给候选 token 的概率足够高，就接受，否则拒绝。但为了保持 Target Model 的原始分布，标准 Speculative Sampling 使用接受-拒绝采样。

对草稿 token x，设：

p(x)：Target Model 的概率；
q(x)：Draft Model 的概率。

接受概率为：

$a(x)=\min\left(1,\dfrac{p(x)}{q(x)}\right)$

如果拒绝，需要从修正分布中重新采样：

$p'(x)=\dfrac{\max(0,p(x)-q(x))}{\sum_y\max(0,p(y)-q(y))}$

这一步的意义是：Draft Model 可以大胆猜，但最终接受的结果仍然服从 Target Model 的分布，而不是把小模型的答案悄悄当成大模型答案。

更严谨地说，是“被接受的草稿 + 拒绝后的修正采样”共同构成 Target 分布，不是只收集接受样本就自然得到 Target 分布。举一个只有 A、B 两种候选的例子：

token	Draft 概率 q	Target 概率 p	接受概率
A	0.8	0.6	0.75
B	0.2	0.4	1

Draft 过于偏爱 A。接受的 A 概率质量是 0.8×0.75=0.6；接受的 B 概率质量是 0.2×1=0.2。剩余 0.2 是拒绝概率，修正分布正好全部给 B，所以最终是 P(A)=0.6、P(B)=0.4。

每一步的 p 和 q 都要基于相同的已接受前缀计算；实际温度、top-k/top-p 等处理后的采样分布也必须纳入接受规则。分布等价不意味着使用相同随机种子会逐字生成相同文本。浮点误差、近似接受策略或工程实现差异也需要另行评估。Greedy 模式则可比较 Target 的 argmax 与草稿是否一致，遇到首个不一致就采用 Target 选择，不需要套用随机采样例子。

5. 一个完整例子

Draft Model 生成：

[很, 适, 合, 散步]

Target Model 逐位置验证后：

很     ✓ 接受
适     ✓ 接受
合     ✓ 接受
散步   ✗ 拒绝

那么本轮直接提交：

[很, 适, 合]

拒绝位置之后的候选不能继续沿用，因为它们是基于被拒绝的草稿路径计算的。标准算法会在拒绝处从修正分布采样一个替代 token，输出三个已接受 token 加这个替代 token。

如果四个候选都被接受，还可以利用 Target 在最后一个候选位置产生的分布，多采样一个 bonus token。标准算法中，草稿接受数为 0 到 k，加上修正 token 或 bonus token，本轮一般推进 1 到 k+1 个 token；EOS、长度上限会提前终止。

6. KV Cache 如何配合？

投机解码会产生一个重要工程问题：候选 token 的 K/V 不能在验证前直接永久提交。

验证前：

正式 Cache：[原上下文 KV]

Draft Model 会维护自己的草稿缓存；Target Model 在验证这批候选时，会为候选位置计算另一套属于 Target 的临时 K/V。两者参数不同，因此缓存不能混用。下面画的是 Target Cache：

临时 Cache：[原上下文 KV | 很 | 适 | 合 | 散步]

验证后，如果只接受前三个：

提交 Cache：[原上下文 KV | 很 | 适 | 合]
回滚：      丢弃“散步”及其后续临时状态

工程上通常使用预留槽位和提交长度指针，不会为了拒绝一个 token 复制整块缓存。每层都必须保持一致的提交边界，否则下一轮 Attention 会读到错误的历史状态。

如果拒绝位置通过修正分布采样出了替代 token，这个替代 token 的 Target K/V 还没有包含在已提交前缀中。系统需要在后续 Decode 中处理它，再把对应 K/V 写入正式 Cache；具体实现也可能利用验证得到的额外状态减少重复工作。

投机解码中 KV Cache 的临时写入、接受提交和拒绝回滚

7. 为什么它能提高大模型利用率？

普通 Decode 的大模型输入序列维度接近 1：

$[1,H]@[H,H]$

投机验证则把多个候选位置合成一个短序列：

$[k,H]@[H,H]$

一次加载大模型权重，就能为多个候选位置计算 logits。序列维度变大后，GEMM 的并行度和算术强度通常更好，GPU 计算单元也更容易被利用。

但它没有消除 Attention 对历史 KV 的读取。验证阶段还要维护草稿分支对应的临时 K/V，实际收益取决于 kernel、缓存布局和调度。

8. 加速取决于接受率

设每次 Draft 提出 k 个 token，连续接受长度为 A。粗略地看，Target 调用次数大约从“每生成一个 token 调用一次”变成“每轮验证推进 A 个左右”。

但不能只看 k：

候选很长 + 第一个 token 就被拒绝 → 验证工作浪费
候选较短 + 接受率高           → 更容易获得稳定加速

代码、格式化文本和模型高置信度区域往往更容易得到长的连续接受前缀；开放式、分布发散的内容接受率可能较低。

实际收益还要扣除：

Draft Model 自身的计算时间；
Candidate tree 或临时 KV Cache 管理开销；
Target 验证 kernel 的效率；
动态 batching 下不同请求的候选长度差异。

9. 它和 KV Cache 优化是什么关系？

Speculative Decoding 不直接减少每个历史 token 的 KV 大小，也不一定减少历史 token 的数量。它优化的是：

普通 Decode：每个 token 都触发一次大模型调用
投机解码：多 token 草稿由一次大模型验证，成功前缀一起提交

因此它与前面三类方法正交，可以组合：

GQA/MLA：减少 KV Cache 的宽度
Sliding Window/Pruning：减少 KV Cache 的长度
Speculative Decoding：减少大模型串行调用次数

10. 常见变体

10.1 独立 Draft Model

使用一个较小的语言模型作为草稿模型，结构清晰，部署时需要额外加载一个模型。

10.2 N-gram 或前缀匹配

不引入额外神经网络，而是从 Prompt、历史生成文本或检索结果中寻找可能的连续片段，适合重复度高的内容。

10.3 Self-Speculative / 多头草稿

让同一个大模型增加轻量预测头，或使用中间层表示提出候选。这样减少了额外模型成本，但需要专门的结构和训练。

10.4 Tree-based Draft

草稿模型提出多个可能分支，Target Model 一次验证一棵候选树。它能提高找到长接受路径的概率，但会增加树的管理和临时缓存复杂度。

11. 总结

投机解码的核心流水线是：

当前正式上下文
      ↓
Draft Model 逐 token 猜 k 个候选
      ↓
Target Model 一次并行验证候选
      ↓
接受连续前缀，拒绝位置执行修正采样
      ↓
提交接受 token 的 KV Cache，回滚其余临时状态

最重要的三点是：

Draft Model 负责提出候选，Target Model 负责决定最终分布；
大模型并行的是多个候选位置的验证，不是无条件独立生成多个 token；
KV Cache 需要支持临时写入、提交和回滚，才能保证候选被拒绝后状态仍然正确。