投机解码深度解析：LLM 推理速度翻倍的秘密

自回归解码的瓶颈

标准 LLM 推理是自回归的：每次只生成一个 token，然后把它拼回输入再生成下一个。这意味着：

# 标准自回归生成
for step in range(max_tokens):
    # 前向传播：计算整个序列
    logits = model(input_ids)           # O(N) 计算
    next_token = sample(logits[:, -1])   # 只用最后一个位置
    input_ids = concat(input_ids, next_token)

问题在于解码阶段的 GPU 利用率极低。即使 batch_size=1，生成单个 token 也需要加载全部模型权重，但只产生 1 个 token 的输出。这就是所谓的 memory-bound（内存带宽瓶颈）。

模型规模	参数量	生成1 token需读取	实际计算量	利用率
7B	7B	~14 GB	~14 GFLOP	<5%
70B	70B	~140 GB	~140 GFLOP	<3%
405B	405B	~810 GB	~810 GFLOP	<2%

核心洞察：大模型推理时，大部分时间花在搬权重，而不是做计算。如果能一次性预测多个 token，就能摊薄权重加载成本。

投机解码：核心思想

投机解码（Speculative Decoding）由 Chen et al. (2023) 和 Leviathan et al. (2023) 同时提出，核心思想是：

用一个**小模型（draft model）**快速生成 k 个候选 token
用**大模型（target model）**并行验证这 k 个 token
接受小模型猜对的 token，拒绝后重新采样

标准解码：  [大模型] → t1 → [大模型] → t2 → [大模型] → t3 → [大模型] → t4
投机解码：  [小模型] → t1', t2', t3'  (快速猜测)
            [大模型] → 并行验证 t1', t2', t3'  (一次前向传播)
            结果：    t1' ✓, t2' ✓, t3' ✗ → 从 t3 处重新采样

算法详解

import torch
import torch.nn.functional as F

def speculative_decode(target_model, draft_model, input_ids, 
                        max_tokens=256, k=4):
    """
    投机解码算法
    
    Args:
        target_model: 大模型（target）
        draft_model: 小模型（draft）
        input_ids: 输入 token ids (1, seq_len)
        max_tokens: 最大生成 token 数
        k: 每次投机猜测的 token 数
    """
    generated = input_ids.clone()
    
    while generated.shape[1] - input_ids.shape[1] < max_tokens:
        # Step 1: 小模型快速生成 k 个候选 token
        draft_tokens = []
        draft_probs = []
        temp_ids = generated.clone()
        
        for _ in range(k):
            draft_logits = draft_model(temp_ids)
            next_token = sample(draft_logits[:, -1, :])
            draft_tokens.append(next_token)
            draft_probs.append(softmax(draft_logits[:, -1, :]))
            temp_ids = torch.cat([temp_ids, next_token.unsqueeze(0)], dim=-1)
        
        draft_tokens = torch.stack(draft_tokens, dim=1)  # (1, k)
        
        # Step 2: 大模型并行验证
        # 将 generated + draft_tokens 一起送入大模型
        verify_input = torch.cat([generated, draft_tokens], dim=1)
        target_logits = target_model(verify_input)
        
        # 大模型在 draft token 位置的 logits
        target_probs = F.softmax(target_logits[:, generated.shape[1]-1:-1, :], dim=-1)
        
        # Step 3: 接受/拒绝
        accepted = 0
        for i in range(k):
            draft_token = draft_tokens[0, i].item()
            target_prob = target_probs[0, i, draft_token].item()
            draft_prob = draft_probs[i][0, draft_token].item()
            
            # 拒绝采样：如果 target_prob >= draft_prob，必定接受
            # 否则以 target_prob / draft_prob 的概率接受
            if target_prob >= draft_prob:
                accepted += 1
                continue
            else:
                ratio = target_prob / max(draft_prob, 1e-10)
                if torch.rand(1).item() < ratio:
                    accepted += 1
                    continue
                else:
                    # 拒绝！从修正分布中采样
                    residual_prob = F.relu(target_probs[0, i] - draft_probs[i][0])
                    residual_prob = residual_prob / residual_prob.sum()
                    new_token = torch.multinomial(residual_prob, 1)
                    break
        
        # 接受的 token 直接追加
        generated = torch.cat([generated, draft_tokens[:, :accepted]], dim=-1)
        
        # 如果有拒绝，追加重新采样的 token
        if accepted < k:
            generated = torch.cat([generated, new_token.unsqueeze(0)], dim=-1)
        else:
            # 全部接受，用大模型最后一个位置的预测继续
            bonus_token = sample(target_logits[:, -1, :])
            generated = torch.cat([generated, bonus_token.unsqueeze(0)], dim=-1)
    
    return generated

为什么能加速？

关键在于大模型的并行验证。大模型处理 k+1 个 token 只比处理 1 个 token 稍慢一点（因为主要是 memory-bound），但小模型生成 k 个 token 非常快。

假设：

大模型单 token 延迟：T_target = 50ms
小模型单 token 延迟：T_draft = 5ms
接受率：α = 0.7（70% 的猜测被接受）
k = 4

方式	耗时	产出 token	tokens/s
标准解码	4 × 50ms = 200ms	4	20
投机解码	4×5ms + 50ms = 70ms	~3.4 (期望)	~49

加速比约 2.5x，且接受率越高加速越明显。

Medusa：多头并行猜测

Medusa（Cai et al., 2024）不需要单独的小模型，而是在大模型最后一层添加多个猜测头：

                    ┌── Head 0 (主头) → t1
hidden state ──────┼── Head 1 (猜测) → t2'
                    ├── Head 2 (猜测) → t3'
                    └── Head 3 (猜测) → t4'

class MedusaHead(nn.Module):
    def __init__(self, hidden_size, vocab_size):
        super().__init__()
        self.linear = nn.Linear(hidden_size, hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, hidden):
        return self.lm_head(F.silu(self.linear(hidden)))

class MedusaModel(nn.Module):
    def __init__(self, base_model, num_heads=4):
        super().__init__()
        self.base_model = base_model
        self.medusa_heads = nn.ModuleList([
            MedusaHead(base_model.config.hidden_size, 
                       base_model.config.vocab_size)
            for _ in range(num_heads)
        ])
    
    def forward(self, input_ids):
        # 基础模型前向传播
        outputs = self.base_model(input_ids, output_hidden_states=True)
        hidden = outputs.hidden_states[-1]
        
        # 主头 + 猜测头
        main_logits = self.base_model.lm_head(hidden)
        medusa_logits = [head(hidden) for head in self.medusa_heads]
        
        return main_logits, medusa_logits

Medusa 的优势：

无需训练独立小模型
猜测头只需轻量训练（LoRA 即可）
使用 tree attention 验证多个候选路径

Tree Attention

         t1
        / | \
     t2a t2b t2c
     |   |   |
     t3a t3b t3d

通过构建候选 token 树，用注意力掩码让它们共享 KV Cache：

def tree_attention_mask(tree_structure, total_nodes):
    """
    构建树形注意力掩码
    每个节点只能看到自己的祖先路径
    """
    mask = torch.zeros(total_nodes, total_nodes)
    for node_idx, parent_idx in enumerate(tree_structure):
        # 节点可以看到自己和所有祖先
        ancestor = node_idx
        while ancestor != -1:
            mask[node_idx, ancestor] = 1
            ancestor = tree_structure[ancestor] if ancestor > 0 else -1
    return mask.bool()

EAGLE：基于特征的猜测

EAGLE（Li et al., 2024）是目前最先进的投机解码方法之一。它观察到：

隐藏特征比 token 本身更适合预测未来 token
用大模型的倒数第二层特征训练一个轻量猜测网络

class EAGLEDecoder(nn.Module):
    def __init__(self, hidden_size, vocab_size, num_layers=2):
        super().__init__()
        # 输入：当前层的 hidden state + embedding
        self.input_proj = nn.Linear(hidden_size + vocab_size, hidden_size)
        
        # 轻量 Transformer 解码器
        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(hidden_size, nhead=8,
                                        dim_feedforward=hidden_size*4),
            num_layers=num_layers
        )
        
        # 输出层
        self.lm_head = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, hidden_states, embeddings, k=4):
        """
        hidden_states: 大模型最后一层隐藏状态 (1, seq, hidden)
        embeddings: token embeddings (1, seq, emb)
        """
        x = self.input_proj(torch.cat([hidden_states, embeddings], dim=-1))
        
        # 自回归生成 k 个候选 token 的特征
        draft_features = []
        for _ in range(k):
            x = self.decoder(x, x)  # 自注意力递推
            draft_features.append(x[:, -1:, :])
        
        # 预测 token
        draft_logits = [self.lm_head(f) for f in draft_features]
        return draft_logits

EAGLE 的接受率可达 85%+，加速比 3-4x。

各方法对比

方法	需要额外模型	训练成本	接受率	加速比	适用场景
Vanilla Spec	✅ 小模型	中	~60%	2x	有现成小模型
Medusa	❌ 猜测头	低	~50%	2.2x	通用
Medusa-2	❌ 猜测头	中	~65%	2.8x	通用
EAGLE	✅ 轻量网络	中	~85%	3.5x	追求极致
EAGLE-2	✅ 轻量网络	中	~88%	4x	动态上下文
Lookahead	❌	无	~45%	1.8x	零成本

工程实践建议

1. 选择合适的 draft model

# 经验法则：draft model 参数量应为 target 的 1/10 ~ 1/5
# 例如 Llama-3-70B → Llama-3-8B 或 Qwen-2.5-1.5B
GOOD_PAIRS = [
    ("Llama-3-70B",  "Llama-3-8B"),     # 8.75x 比例
    ("Qwen-2.5-72B", "Qwen-2.5-7B"),    # 10x 比例
    ("Llama-3-70B",  "Llama-3.5-1.5B"), # 46x，接受率较低但速度快
]

2. 动态调整 k 值

def adaptive_k(acceptance_history, base_k=4, min_k=2, max_k=8):
    """根据历史接受率动态调整猜测长度"""
    recent_rate = sum(acceptance_history[-20:]) / max(len(acceptance_history[-20:]), 1)
    if recent_rate > 0.8:
        return min(max_k, base_k + 2)
    elif recent_rate < 0.4:
        return max(min_k, base_k - 2)
    return base_k

3. vLLM 中的投机解码

# vLLM 已内置投机解码支持
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="meta-llama/Llama-3-8B",  # draft model
    num_speculative_tokens=5,                     # k 值
    use_v2_block_manager=True,
)

sampling = SamplingParams(temperature=0.7, max_tokens=512)
output = llm.generate("Explain quantum computing", sampling)

参考文献

Leviathan, Y. et al. (2023). Fast Inference from Transformers via Speculative Decoding
Cai, T. et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
Li, Y. et al. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

自回归解码的瓶颈#

投机解码：核心思想#

算法详解#

为什么能加速？#

Medusa：多头并行猜测#

Tree Attention#

EAGLE：基于特征的猜测#

各方法对比#

工程实践建议#

1. 选择合适的 draft model#

2. 动态调整 k 值#

3. vLLM 中的投机解码#

参考文献#

加入讨论#

📌 相关推荐