Notes/机器学习/深度学习模型/07-计算架构/3-推理优化/3-投机解码.md

投机解码（Speculative Decoding）

背景与动机

自回归生成的计算瓶颈

在自回归LLM生成中，每个token的生成依赖于之前所有token的KV Cache：

生成序列: [t1, t2, t3, ..., tN]

每个token生成: P(t_i | t_1, ..., t_{i-1})
计算量: 每步 O(seq_len × hidden) 的注意力计算

问题：即使正确预测了大部分token，每个token生成仍需完整注意力计算，GPU计算资源消耗大。

投机解码思想

投机解码（Speculative Decoding）使用一个轻量的"draft model"快速生成多个候选token，再用原模型验证：

Draft Model: 快速生成 [d1, d2, d3, ..., d_k] 候选
Target Model: 验证候选token的准确性
接受策略: 概率阈值或贪婪验证

工作原理

核心流程

Step 1: Draft Model前向
   输入: [t_1, ..., t_{i-1}]
   Draft输出: [d_1, d_2, d_3, ..., d_k]  (k个候选)
   时间: T_draft << T_target

Step 2: Target Model验证
   构造: [t_1, ..., t_{i-1}, d_1, d_2, ..., d_k]
   Target输出: P_target(d_j) 对每个d_j
   时间: T_target

Step 3: 接受/拒绝决策
   - 使用 q(d_j) / p(d_j) 概率比值
   - 接受: draft预测概率高于阈值
   - 拒绝: 采样或回退

数学原理

def speculative_decoding(draft_probs, target_probs, temperature=1.0):
    """
    draft_probs: draft model对候选的预测概率
    target_probs: target model对候选的验证概率
    """
    # 概率比值检验
    acceptance_ratio = target_probs / (draft_probs + epsilon)
    
    # 生成随机数决定接受/拒绝
    threshold = torch.rand_like(acceptance_ratio)
    
    accepted = acceptance_ratio > threshold
    first_reject = find_first_false(accepted)  # 首个拒绝位置
    
    return first_reject

加速比分析

理想情况（draft模型准确率高）:
- Draft生成k个token，验证通过
- 实际执行: k次draft + 1次target验证
- 加速比: (k × T_target) / (T_draft + T_target) ≈ k × T_target / T_target = k

实际约束:
- Draft模型质量影响接受率
- 候选序列长度k的优化
- Target和Draft共享部分计算（KV Cache）

加速比公式

def compute_speedup(k, accept_rate, T_target, T_draft):
    """
    k: draft候选数
    accept_rate: 平均接受率
    T_target / T_draft: 时间比值
    """
    # 期望生成token数（几何分布）
    expected_tokens = (1 - accept_rate**k) / (1 - accept_rate) if accept_rate < 1 else k
    
    # 总计算时间
    T_total = T_draft + T_target  # draft一次 + target一次（验证k个）
    
    # 朴素自回归时间
    T_baseline = expected_tokens * T_target
    
    speedup = T_baseline / T_total
    return speedup

Draft Model设计

轻量化模型

常见的Draft Model选择：

1. 小模型自回归：使用参数量小的LLM（如7B draft服务65B target）
2. 投机采样（Speculative Sampling）：相同模型但减少候选数
3. N-gram模型：简单但快速
4. 基于树的解码：构建token树而非线性序列

结构设计

class DraftModel:
    def __init__(self, model_path, num_layers=4):
        # 截断的深层网络
        self.transformer = load_partial_model(model_path, num_layers)
        
    def generate_candidates(self, x, k=5):
        """快速生成k个候选token"""
        candidates = []
        for _ in range(k):
            logits = self.forward(x)
            token = sample(logits, temperature=0.8)
            candidates.append(token)
            x = torch.cat([x, token.unsqueeze(0)], dim=-1)
        return candidates

与Target Model的关系

Draft Model可以是:
1. 独立的小模型（不同权重）
2. 同一模型的前N层（共享底层）
3. 量化版本的目标模型

KV Cache共享:
- Draft和Target共享prefix的KV Cache
- 只对draft生成的token计算增量

Verifier设计

验证策略

def verify_candidates(draft_candidates, target_model, input_ids, threshold=0.5):
    """
    验证draft生成的候选序列
    """
    # 构建完整输入（包括候选）
    full_input = input_ids + draft_candidates
    
    # Target model前向传播
    with torch.no_grad():
        target_logits = target_model(full_input)
        target_probs = F.softmax(target_logits / temperature, dim=-1)
    
    # 验证每个候选
    accepted = []
    for i, draft_token in enumerate(draft_candidates):
        # 取target对draft_token的预测概率
        p_target = target_probs[input_ids.shape[1] - 1 + i, draft_token]
        p_draft = draft_probs[i, draft_token]  # draft模型的预测
        
        # 接受条件
        if p_target > p_draft * threshold:  # 放宽接受条件
            accepted.append(draft_token)
        else:
            break  # 拒绝后停止
    
    # 返回接受的候选数和回退token
    return accepted, target_model.sample_next(target_probs[input_ids.shape[1] - 1 + len(accepted)])

采样分布修正

使用修正的采样分布处理拒绝：

def corrected_sampling(target_probs, accepted_ids, draft_probs):
    """
    修正采样分布
    当draft被拒绝时，从修正分布中采样
    """
    # 构造修正分布
    # 从target概率中移除已接受的概率
    corrected = target_probs.clone()
    for token_id in accepted_ids:
        corrected[token_id] = 0
    
    # 重新归一化
    corrected = corrected / corrected.sum()
    
    return torch.multinomial(corrected, num_samples=1)

实现考量

KV Cache利用

普通生成:
  每个token生成: 需要O(seq_len)的注意力计算

投机解码:
  Draft生成: 复用已有KV Cache，只需计算新token
  Target验证: 复用完整KV Cache，一次前向传播验证多个token
  
  总体: 节省 (k-1) × O(seq_len) 的注意力计算

最优候选数k

def find_optimal_k(T_target, T_draft, accept_rate):
    """
    搜索最优k值
    """
    best_k = 1
    best_speedup = 0
    
    for k in range(1, 16):
        speedup = compute_speedup(k, accept_rate, T_target, T_draft)
        if speedup > best_speedup:
            best_speedup = speedup
            best_k = k
    
    return best_k

# 典型参数
# T_target / T_draft ≈ 10  (draft是target的1/10计算量)
# accept_rate ≈ 0.8
# 最优k ≈ 4-6

拒绝回退

def handle_rejection(accepted_count, draft_candidates, target_model, full_input):
    """
    处理拒绝的情况
    """
    if accepted_count == len(draft_candidates):
        # 所有候选被接受，继续
        return target_model.generate_next()
    else:
        # 有候选被拒绝
        # 使用target model在拒绝位置采样
        next_token = target_model.sample_at(len(accepted_count))
        return next_token

与其他技术的关系

技术	关系	组合效果
KV Cache	共享基础	提高共享效率
Continuous Batching	互补	投机解码本身支持批处理
Flash Attention	底层加速	加速验证阶段计算
量化	互补	Draft模型可用更低精度

实验结果

模型配置	加速比	备注
LLaMA-65B + 7B draft	2.5-3x	接受率约80%
LLaMA-13B + 7B draft	2-2.5x	接受率约75%
优化后（自适应k）	3-4x	动态调整候选数

实践建议

部署配置

# vLLM中的投机解码配置
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    speculative_model="meta-llama/Llama-2-7b-hf",  # draft model
    num_speculative_tokens=5,  # 候选数
    temperature=0.8
)

适用场景

适用:
- 长序列生成（token数>100）
- 高批量推理（请求间可共享KV Cache）
- 对延迟要求较高的在线服务

不适用:
- 极短序列（投机开销大于收益）
- 高质量要求场景（接受率下降影响质量）
- Draft模型质量差的情况

质量-速度权衡

配置	速度	质量损失
高接受率（阈值低）	慢	少
低接受率（阈值高）	快	可能明显
自适应k	最优平衡	可控