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预测概率高于阈值
   - 拒绝: 采样或回退
```

### 数学原理

```python
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）
```

### 加速比公式

```python
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树而非线性序列

### 结构设计

```python
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设计

### 验证策略

```python
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)])
```

### 采样分布修正

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

```python
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

```python
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
```

### 拒绝回退

```python
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 | 动态调整候选数 |

## 实践建议

### 部署配置

```python
# 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 | 最优平衡 | 可控 |