推测解码加速原理：Draft Model 验证范式

1. Draft Model 自回归快速生成 k 个 token: [t1, t2, ..., tk]
2. Target Model 一次性对 [prompt + t1, ..., tk] 做前向传播
3. Target Model 的输出可以验证 t1..tk 的正确性
4. 接受匹配的 token，拒绝不匹配的，从第一个拒绝点重新生成

import torch
import torch.nn.functional as F

def speculative_sampling(draft_logits, target_logits, draft_token):
    """推测解码的接受/拒绝采样"""
    draft_prob = F.softmax(draft_logits, dim=-1)
    target_prob = F.softmax(target_logits, dim=-1)
    
    # 接受概率
    accept_prob = torch.minimum(
        torch.ones_like(draft_prob),
        target_prob / (draft_prob + 1e-10)
    )
    
    # 随机决定是否接受
    rand = torch.rand_like(draft_prob)
    accepted = rand < accept_prob
    
    if accepted[draft_token]:
        return draft_token  # 接受 Draft Model 的 token
    else:
        # 从残差分布重新采样
        residual = torch.clamp(target_prob - draft_prob, min=0)
        residual = residual / residual.sum()
        return torch.multinomial(residual, 1).squeeze()

Target Model 隐状态
    ├── LM Head → t0 (当前 token)
    ├── Medusa Head 1 → t1 (下一个 token)
    ├── Medusa Head 2 → t2
    └── Medusa Head 3 → t3

class MedusaHead(nn.Module):
    """Medusa 预测头"""
    def __init__(self, d_model, vocab_size):
        super().__init__()
        self.linear = nn.Linear(d_model, d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)
        
    def forward(self, hidden_state):
        x = F.silu(self.linear(hidden_state))
        return self.lm_head(x)

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]
        
        # 基座预测当前 token
        base_logits = self.base_model.lm_head(hidden)
        
        # Medusa 头预测未来 token
        medusa_logits = [head(hidden) for head in self.medusa_heads]
        return base_logits, medusa_logits

                    t0
                   / | \
              t1a  t1b  t1c
              / \   |    |
          t2a t2b t2c  t2d

from vllm import LLM, SamplingParams

# 配置推测解码
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    speculative_model="meta-llama/Llama-2-7b-hf",  # Draft Model
    num_speculative_tokens=5,  # 推测长度
    use_v2_block_manager=True,
)

sampling_params = SamplingParams(
    temperature=0.0,  # 贪婪解码接受率最高
    max_tokens=1000,
)

outputs = llm.generate(["请解释量子计算的基本原理"], sampling_params)

┌─────────────────────────────────────────┐
│           vLLM 推理引擎                  │
│                                         │
│  ┌──────────┐    ┌──────────────────┐  │
│  │ Draft     │───→│ Target Model     │  │
│  │ Model     │    │ (单次前向验证)    │  │
│  │ (快速生成) │    │                  │  │
│  └──────────┘    └──────────────────┘  │
│       │                   │             │
│       ▼                   ▼             │
│   [t1,t2,t3,t4,t5]  ← 验证 & 接受/拒绝  │
│                                │         │
│                                ▼         │
│                    [accepted tokens]     │
│                                │         │
│                                ▼         │
│                    更新 KV Cache         │
└─────────────────────────────────────────┘