LLM推理的"投机取巧"：推测解码如何榨干GPU算力

为什么LLM生成这么慢？

当你用 ChatGPT 或 Claude 生成一段文字时，有没有注意到输出是"一个字一个字"蹦出来的？这不是产品设计选择，而是 Transformer 架构的本质约束——自回归解码（Autoregressive Decoding） 的串行特性让每一步都在等上一步。

让我们先理解这个问题有多严重。假设用 A100 GPU 运行 Llama-3-70B：

Decode 阶段之所以慢，是因为 Transformer 的核心是 Self-Attention：生成第 N 个 token 时，需要attend到前面 N-1 个 token。这意味着：

1. 第 N 步的计算量 ≈ 第 N-1 步的计算量（KV Cache 帮助下）

2. 但每次计算只能产生 1 个 token

3. GPU 的矩阵乘法单元（Tensor Core）严重吃不饱——它们设计用来处理大批量并行运算，现在却被绑在单个 token 上

这就是为什么 A100 跑 70B 模型时，GPU 利用率常常低于 30%。

推测解码：用一个便宜模型预测，多个贵模型验证

推测解码（Speculative Decoding）的核心思想来自一个简单观察：

> 与其每次让大模型猜 1 个字，不如让小模型先猜 一串字，然后让大模型 并行验证。

数学直觉

假设：

• 大模型 M_big 生成 1 token 需要时间 T_big
• 小模型 M_small 生成 1 token 需要时间 T_small，且 T_small ≈ T_big / k（k 通常 5-20x）
• 小模型预测的接受率为 α（约 70-90%）

传统自回归方式：T_big per token

推测解码：

1. 小模型一次生成 γ 个 token：时间 ≈ γ × T_small

2. 大模型并行验证这 γ 个 token：时间 ≈ T_big（矩阵乘法，并行处理）

3. 实际被接受的 token 数 ≈ α × γ

有效每 token 时间 ≈ (γ × T_small + T_big) / (α × γ + 1)

当 T_big >> T_small, α > 0.7, γ = 32-64 时，加速比可达 2-5x。

实际算法流程

`

大模型 M_big (目标模型)

小模型 M_small (推测模型，与 M_big 通常同结构)

温度 T > 0 用于生成

# 阶段1：小模型"猜测"

draft_tokens = []

for i in range(gamma):

token = M_small.forward(draft_tokens)

draft_tokens.append(token)

if token == EOS: break

# 阶段2：大模型"验证"（KVCache 复用，并行）

logits_big = M_big.forward_parallel(draft_tokens) # 一次性 forward

# 阶段3：自适应接受

accepted = 0

for i, token in enumerate(draft_tokens):

# 方法1：Greedy - 只看最大概率 token

if i == 0:

accepted_token = argmax(logits_big[i])

else:

accepted_token = argmax(logits_big[i])

if token == accepted_token or random.random() < alpha:

accepted += 1

else:

# 拒绝，从这里重新采样

draft_tokens = draft_tokens[:accepted+1]

draft_tokens.append(sampling_from(logits_big[accepted]))

break

return draft_tokens # accepted + 1 tokens

`

实现细节：为什么你的推测解码跑不起来？

KVCache 的陷阱

大多数推测解码实现会在这里翻车。小模型生成的 token 序列，在大模型眼里是 全新的 token，没有 KVCache 可用。这意味着大模型的验证阶段仍然需要计算 attention——但好消息是，所有 γ 个位置的 attention 可以并行计算，而不是串行。

`python

# 伪代码：大模型并行验证（关键优化）

def verify_large_model(batch_draft_tokens: List[Token], kv_cache: KVCache):

"""

一次性处理所有 draft tokens，利用矩阵并行的优势

不同于自回归的单 token forward，这里是批量矩阵乘法

"""

# 输入形状: [batch_size=gamma, seq_len]

# 输出形状: [batch_size=gamma, vocab_size]

logits = large_model(batch_draft_tokens, kv_cache=kv_cache)

# 注意：这里 large_model 需要支持 packed batch inference

return logits

`

如何选择小模型？

不是随便找个小模型就行。关键要求：

1. 分布对齐：小模型和大模型的预测分布要足够接近

2. draft 长度：γ 越大，加速比越高，但接受率会下降

3. 延迟差距：T_small / T_big 越大，整体收益越高

最佳实践：

• **同结构不同 size**：如 Llama-3-70B + Llama-3-8B，接受率高
• **同 size 不同量化**：Q4 大模型 + FP16 小模型
• **蒸馏模型**：用大模型数据微调过的小模型，接受率可到 95%+

拒绝采样策略

最简单的 Greedy（只接受最大概率 token）效果一般。更好的方法：

方法1：基于温度的接受

`python

def accept_with_temperature(logits_draft, logits_big, temperature=0.8):

# 计算小模型和大模型在 draft token 上的概率比

p_small = softmax(logits_draft / temperature)

p_big = softmax(logits_big)

for i, token in enumerate(draft_tokens):

ratio = p_big[token] / (p_small[token] + 1e-8)

if ratio > 1.0 or random.random() < ratio:

accepted.append(token)

else:

# 拒绝，重新采样

accepted.append(sample_from(p_big))

break

return accepted

`

方法2：树状验证（Tree Verification）

一次验证多个 draft 路径，而不是线性序列。Google 的 Medusa 和 HuggingFace 的 EDSD 都用了这个思路：

`python

# Medusa 风格：多个 draft head 同时预测

class MedusaHead(nn.Module):

def __init__(self, hidden_size, vocab_size, depth=5):

super().__init__()

self.layers = nn.ModuleList([

nn.Sequential(

nn.Linear(hidden_size, hidden_size),

nn.ReLU(),

nn.Linear(hidden_size, vocab_size)

) for _ in range(depth)

])

def forward(self, hidden_states):

# 同时预测 5 个未来位置的 token

return [layer(hidden_states) for layer in self.layers]

`

实战：HuggingFace Speculative Decoding 详解

HuggingFace Transformers 从 4.36 开始支持推测解码：

`python

from transformers import AutoModelForCausalLM, AutoTokenizer

from transformers.generation import SpeculativeDecoding

model_id = "meta-llama/Llama-3.1-70B-Instruct"

small_model_id = "meta-llama/Llama-3.1-8B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)

big_model = AutoModelForCausalLM.from_pretrained(

model_id,

device_map="auto",

torch_dtype=torch.bfloat16

)

small_model = AutoModelForCausalLM.from_pretrained(

small_model_id,

device_map="auto"

)

# 关键参数

speculative_decoding = SpeculativeDecoding(

main_model=big_model,

speculative_model=small_model,

num_speculative_tokens=32, # γ，越大越省但越挑模型

threshold=0.8, # 接受率阈值

)

prompt = "解释为什么天空是蓝色的，用物理原理说明："

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with SpeculativeDecoding.speculative_decoding_context(speculative_decoding):

outputs = big_model.generate(

**inputs,

max_new_tokens=200,

do_sample=True,

temperature=0.7,

)

result = tokenizer.decode(outputs[0], skip_special_tokens=True)

`

Benchmark 数据（Llama-3-70B + Llama-3-8B，A100 80GB）：

超越推测解码：未来方向

推测解码只是 LLM 推理优化的冰山一角。更激动人心的方向：

1. 投石机解码（Rockpile Decoding）

用 KVCache 预测下一个 token 的位置，直接跳转到那里计算，避免无意义的 attention。

2. 前向验证（Forward Verification）

不仅是预测下一个 token，而是预测下一个 N 个 token 的完整 KV 向量，大模型直接用这些预计算的 KV 进行验证。

3. 混合推测（Hybrid Speculation）

根据内容难度自适应选择：小模型负责简单句子的预测，遇到复杂推理时自动退化为大模型直接生成。

`python

class AdaptiveSpeculativeDecoder:

def __init__(self, big_model, small_models: List):

self.big = big_model

self.smalls = small_models # 多级模型：8B, 3B, 1B

def generate(self, prompt, difficulty_hint=None):

if difficulty_hint == "complex":

return self.big.generate(prompt) # 直接用大模型

small = self.smalls[0] # 默认用最大的小模型

return self.speculative_decode(prompt, small)

`

4. 推测解码 + 量化协同

Q4 量化的大模型 + INT8 的小模型，减少内存带宽压力，进一步放大加速效果。

总结：什么场景适合推测解码？

优点：

• 显著提升 token 生成吞吐量（2-4x）
• 不改变模型输出分布（数学上等效）
• 易于集成到现有推理框架

缺点：

• 增加了内存占用（小模型也要在显存里）
• 对小模型质量要求高
• 不适合流式输出场景（需要等 γ 个 token 才能开始验证）

最佳场景：

• 批量推理（batch inference）
• 对延迟要求不高但对吞吐要求高的场景（客服机器人、文案生成）
• 部署时显存足够放下两个模型

不太适合：

• 实时交互流式输出（每次等 γ 个 token 才开始吐字）
• 单个请求延迟敏感场景（首 token 时间不变）

推测解码不是银弹，但它聪明地利用了"小模型预测 + 大模型验证"的范式，在不损失质量的前提下榨干 GPU 并行算力。如果你正在优化 LLM 推理服务，值得把它加入工具箱。

---

*附：主流实现参考*

• HuggingFace Transformers: `generation.SpeculativeDecoding`
• vLLM: `--speculative-decoding` flag
• TensorRT-LLM: `SpeculativeDecodingPlugin`
• Medusa (多 draft head): https://github.com/FasterDecoding/Medusa