LLM 推理加速（四）：新范式从 Speculative Decoding 到 DFlash

输入: "The cat sat on the"
→ 大模型推理 → "mat"
→ 大模型推理 → "."
→ 大模型推理 → "The"
→ 大模型推理 → "end"

输入: "The cat sat on the"
→ 草稿模型推理 → 猜测: ["mat", ".", "The", "end"]
→ 大模型一次性验证 4 个 Token
→ 接受前 3 个，拒绝第 4 个
→ 从拒绝位置继续

α = min(1, P_large(x_t | context) / P_draft(x_t | context))

阶段 1：冻结大模型主干，只训练额外的解码头
阶段 2：解冻最后几层，微调整个模型

步骤 1：对目标块添加噪声，得到 x_T（完全随机）
步骤 2：逐步去噪 x_T → x_{T-1} → ... → x_0
步骤 3：每一步使用草稿模型预测去噪方向
步骤 4：用大模型验证关键步骤的采样质量
步骤 5：当置信度达到阈值时，直接输出

from typing import List, Tuple, Optional
import numpy as np
import time


class SpeculativeDecoder:
    """
    推测解码引擎实现
    
    核心流程:
    1. 草稿模型生成 γ 个候选 Token
    2. 大模型并行验证所有候选 Token
    3. 使用接受-拒绝采样确定输出
    """
    
    def __init__(
        self,
        draft_model,
        target_model,
        gamma: int = 4,
        temperature: float = 0.8,
    ):
        self.draft_model = draft_model
        self.target_model = target_model
        self.gamma = gamma  # 每次推测的 Token 数
        self.temperature = temperature
        self.stats = {
            "total_draft_tokens": 0,
            "total_accepted_tokens": 0,
            "total_rounds": 0,
        }
    
    def _sample_with_temperature(
        self, logits: np.ndarray, temperature: float
    ) -> int:
        """带温度采样的 Token 选择"""
        if temperature == 0:
            return int(np.argmax(logits))
        scaled = logits / temperature
        probs = np.exp(scaled - np.max(scaled))
        probs = probs / probs.sum()
        return int(np.random.choice(len(probs), p=probs))
    
    def draft_tokens(
        self, input_ids: List[int], num_tokens: int
    ) -> List[int]:
        """草稿模型生成候选 Token 序列"""
        tokens = []
        current = list(input_ids)
        for _ in range(num_tokens):
            logits = self.draft_model.forward(current)
            next_token = self._sample_with_temperature(
                logits[-1], self.temperature * 1.2  # 草稿模型用更高温度增加多样性
            )
            tokens.append(next_token)
            current.append(next_token)
        return tokens
    
    def verify_tokens(
        self, input_ids: List[int], draft_tokens: List[int]
    ) -> List[Tuple[int, bool]]:
        """
        大模型并行验证草稿 Token
        
        返回: [(token, accepted), ...]
        """
        # 构建完整输入: input_ids + draft_tokens
        full_input = input_ids + draft_tokens
        
        # 大模型一次性前向传播
        all_logits = self.target_model.forward(full_input)
        
        results = []
        for i, draft_tok in enumerate(draft_tokens):
            # 获取大模型在该位置的分布
            target_logits = all_logits[len(input_ids) + i - 1]
            target_probs = self._softmax(target_logits)
            draft_probs = self._softmax(
                self.draft_model.forward(
                    input_ids + draft_tokens[:i]
                )[-1]
            )
            
            # 接受概率: min(1, P_target / P_draft)
            if draft_probs[draft_tok] > 0:
                acceptance_ratio = (
                    target_probs[draft_tok] / draft_probs[draft_tok]
                )
            else:
                acceptance_ratio = 1.0
            
            accepted = np.random.random() < min(1.0, acceptance_ratio)
            results.append((draft_tok, accepted))
            
            if not accepted:
                break
        
        return results
    
    @staticmethod
    def _softmax(logits: np.ndarray) -> np.ndarray:
        exp_logits = np.exp(logits - np.max(logits))
        return exp_logits / exp_logits.sum()
    
    def generate(
        self,
        input_ids: List[int],
        max_new_tokens: int = 100,
    ) -> List[int]:
        """
        使用推测解码生成文本
        
        Args:
            input_ids: 输入 Token IDs
            max_new_tokens: 最多生成的新 Token 数
            
        Returns:
            生成的 Token IDs
        """
        output = list(input_ids)
        total_generated = 0
        
        while total_generated < max_new_tokens:
            self.stats["total_rounds"] += 1
            
            remaining = max_new_tokens - total_generated
            gamma = min(self.gamma, remaining)
            
            # Step 1: 草稿模型生成
            draft = self.draft_tokens(output, gamma)
            self.stats["total_draft_tokens"] += len(draft)
            
            # Step 2: 大模型验证
            verification = self.verify_tokens(output, draft)
            
            # Step 3: 处理验证结果
            accepted_count = 0
            for token, accepted in verification:
                if accepted:
                    output.append(token)
                    accepted_count += 1
                    total_generated += 1
                    self.stats["total_accepted_tokens"] += 1
                else:
                    # 拒绝时，用大模型重新采样
                    logits = self.target_model.forward(output)
                    new_token = self._sample_with_temperature(
                        logits[-1], self.temperature
                    )
                    output.append(new_token)
                    total_generated += 1
                    break
            
            if total_generated >= max_new_tokens:
                break
        
        return output[len(input_ids):]
    
    def get_stats(self) -> dict:
        """获取统计信息"""
        acceptance_rate = (
            self.stats["total_accepted_tokens"] / 
            max(1, self.stats["total_draft_tokens"])
        )
        return {
            **self.stats,
            "acceptance_rate": f"{acceptance_rate:.1%}",
            "avg_accepted_per_round": (
                self.stats["total_accepted_tokens"] / 
                max(1, self.stats["total_rounds"])
            ),
        }


# ==========================================
# 使用示例：模拟推理加速
# ==========================================

class MockModel:
    """模拟模型，用于演示"""
    
    def __init__(self, vocab_size: int = 1000, quality: float = 0.5):
        self.vocab_size = vocab_size
        self.quality = quality  # 预测准确率 (0-1)
        self._rng = np.random.RandomState(42)
    
    def forward(self, input_ids: List[int]) -> np.ndarray:
        """返回 logits"""
        logits = self._rng.randn(len(input_ids), self.vocab_size) * 0.1
        # 让正确 Token 的 logits 更高
        for i in range(len(input_ids)):
            correct = (input_ids[i] + 7) % self.vocab_size  # 简单模式
            logits[i, correct] += self.quality * 10
        return logits


def benchmark_speculative_decoding():
    """基准测试：对比标准解码和推测解码"""
    vocab_size = 5000
    input_length = 50
    
    # 不同质量的草稿模型
    for draft_quality in [0.3, 0.5, 0.7]:
        draft = MockModel(vocab_size, quality=draft_quality)
        target = MockModel(vocab_size, quality=0.9)
        
        decoder = SpeculativeDecoder(
            draft_model=draft,
            target_model=target,
            gamma=4,
            temperature=0.7,
        )
        
        input_ids = list(range(input_length))
        
        start = time.time()
        output = decoder.generate(input_ids, max_new_tokens=200)
        elapsed = time.time() - start
        
        stats = decoder.get_stats()
        print(f"草稿质量 {draft_quality:.1f}:")
        print(f"  接受率: {stats['acceptance_rate']}")
        print(f"  平均每轮接受: {stats['avg_accepted_per_round']:.1f}")
        print(f"  生成 200 Token 用时: {elapsed:.3f}s")
        print()


if __name__ == "__main__":
    benchmark_speculative_decoding()

"""
Medusa 多解码头训练实现

展示如何在大模型上添加额外的解码头并进行训练。
这是 Medusa 方法的核心：无需独立草稿模型，
而是通过多解码头实现推测解码。
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional
from dataclasses import dataclass


@dataclass
class MedusaConfig:
    """Medusa 配置"""
    num_heads: int = 4           # 额外解码头数量
    hidden_size: int = 4096      # 隐藏层维度
    vocab_size: int = 32000      # 词表大小
    head_layers: int = 2         # 每个 Head 的 MLP 层数
    training_steps: int = 10000  # 训练步数
    lr: float = 1e-4            # 学习率


class MedusaHead(nn.Module):
    """单个 Medusa 解码头"""
    
    def __init__(self, config: MedusaConfig, lookahead: int):
        super().__init__()
        self.lookahead = lookahead  # 预测未来几个 Token
        
        layers = []
        for _ in range(config.head_layers):
            layers.extend([
                nn.Linear(config.hidden_size, config.hidden_size),
                nn.ReLU(),
            ])
        layers.append(nn.Linear(config.hidden_size, config.vocab_size))
        self.net = nn.Sequential(*layers)
    
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.net(hidden_states)


class MedusaModel(nn.Module):
    """
    带 Medusa 头的大语言模型
    
    架构:
    - 基础模型 (LLaMA/Mistral 等)
    - 多个 Medusa 解码头
    """
    
    def __init__(self, base_model: nn.Module, config: MedusaConfig):
        super().__init__()
        self.base_model = base_model
        self.config = config
        
        # 创建多个 Medusa 解码头
        self.medusa_heads = nn.ModuleList([
            MedusaHead(config, lookahead=i + 1)
            for i in range(config.num_heads)
        ])
    
    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> dict:
        # 基础模型前向传播
        outputs = self.base_model(input_ids)
        hidden_states = outputs.last_hidden_state  # [batch, seq, hidden]
        
        # 每个 Medusa Head 预测
        head_logits = []
        for head in self.medusa_heads:
            logits = head(hidden_states)
            head_logits.append(logits)
        
        result = {"head_logits": head_logits}
        
        if labels is not None:
            # 计算损失
            loss = self._compute_loss(head_logits, labels)
            result["loss"] = loss
        
        return result
    
    def _compute_loss(
        self,
        head_logits: List[torch.Tensor],
        labels: torch.Tensor,
    ) -> torch.Tensor:
        """计算多 Head 的联合损失"""
        total_loss = 0
        
        for i, logits in enumerate(head_logits):
            # Head i 预测未来 i+1 个 Token
            shift_labels = labels[:, i + 1:]
            shift_logits = logits[:, :-(i + 1)]
            
            loss = F.cross_entropy(
                shift_logits.reshape(-1, self.config.vocab_size),
                shift_labels.reshape(-1),
                ignore_index=-100,
            )
            total_loss += loss / len(head_logits)
        
        return total_loss
    
    @torch.no_grad()
    def speculative_generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 0.7,
    ) -> torch.Tensor:
        """使用 Medusa 头进行推测解码生成"""
        generated = list(input_ids.squeeze().tolist())
        
        for _ in range(max_new_tokens):
            current_ids = torch.tensor([generated]).to(
                input_ids.device
            )
            
            # 获取隐藏状态
            outputs = self.base_model(current_ids)
            hidden = outputs.last_hidden_state[:, -1:]
            
            # 用所有 Medusa Head 预测
            candidates = []
            for head in self.medusa_heads:
                logits = head(hidden)
                probs = F.softmax(logits[0, 0] / temperature, dim=-1)
                candidates.append(int(probs.argmax()))
            
            # 基础 Head 预测
            base_logits = self.base_model.lm_head(hidden)
            base_probs = F.softmax(
                base_logits[0, 0] / temperature, dim=-1
            )
            
            # 验证并接受
            accepted = []
            for i, candidate in enumerate(candidates):
                # 简化的接受逻辑
                if base_probs[candidate] > 0.01:
                    accepted.append(candidate)
                else:
                    break
            
            if accepted:
                generated.extend(accepted)
            else:
                generated.append(int(base_probs.argmax()))
        
        return torch.tensor(generated)


# ==========================================
# 训练 Medusa Head
# ==========================================

def train_medusa(
    base_model: nn.Module,
    config: MedusaConfig,
    train_loader,
    device: str = "cuda",
):
    """
    训练 Medusa 解码头
    
    两阶段训练:
    阶段 1: 冻结 base_model，只训练 Head
    阶段 2: 解冻最后 N 层，微调整个模型
    """
    medusa = MedusaModel(base_model, config).to(device)
    
    # 阶段 1: 只训练 Head
    for param in medusa.base_model.parameters():
        param.requires_grad = False
    
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, medusa.parameters()),
        lr=config.lr,
    )
    
    print("=== 阶段 1: 训练 Medusa Head ===")
    for step, batch in enumerate(train_loader):
        if step >= config.training_steps:
            break
        
        optimizer.zero_grad()
        result = medusa(
            input_ids=batch["input_ids"].to(device),
            labels=batch["labels"].to(device),
        )
        loss = result["loss"]
        loss.backward()
        optimizer.step()
        
        if step % 100 == 0:
            print(f"Step {step}: loss = {loss.item():.4f}")
    
    # 阶段 2: 解冻最后 N 层
    print("\n=== 阶段 2: 微调 ===")
    for param in medusa.base_model.parameters():
        param.requires_grad = True
    
    # 只解冻最后 4 层
    num_layers = len(
        [m for m in medusa.base_model.children() 
         if isinstance(m, nn.TransformerEncoderLayer)]
    )
    for i, param in enumerate(medusa.base_model.parameters()):
        if i < num_layers - 4:
            param.requires_grad = False
    
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, medusa.parameters()),
        lr=config.lr * 0.1,  # 更低的学习率
    )
    
    for step, batch in enumerate(train_loader):
        optimizer.zero_grad()
        result = medusa(
            input_ids=batch["input_ids"].to(device),
            labels=batch["labels"].to(device),
        )
        loss = result["loss"]
        loss.backward()
        optimizer.step()
    
    return medusa


if __name__ == "__main__":
    # 配置
    config = MedusaConfig(
        num_heads=4,
        hidden_size=4096,
        vocab_size=32000,
    )
    
    # 模拟基础模型
    base_model = nn.Transformer(
        d_model=config.hidden_size,
        nhead=32,
        num_encoder_layers=4,
    )
    
    print(f"Medusa 配置:")
    print(f"  解码头数量: {config.num_heads}")
    print(f"  隐藏层维度: {config.hidden_size}")
    print(f"  词表大小: {config.vocab_size}")
    
    medusa = MedusaModel(base_model, config)
    total_params = sum(p.numel() for p in medusa.parameters())
    medusa_params = sum(p.numel() for m in medusa.medusa_heads 
                        for p in m.parameters())
    
    print(f"  总参数量: {total_params:,}")
    print(f"  Medusa Head 参数量: {medusa_params:,}")
    print(f"  额外开销占比: {medusa_params/total_params:.1%}")

# 1. 安装支持推测解码的 vLLM
pip install vllm>=0.4.0

# 2. 启动服务（指定草稿模型）
vllm serve meta-llama/Llama-3-70B \
    --speculative-model meta-llama/Llama-3-8B \
    --num-speculative-tokens 5

# 3. 客户端请求（无需额外配置）
curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-3-70B",
        "messages": [{"role": "user", "content": "解释推测解码"}]
    }'

class SpeculativeMonitoring:
    """推测解码性能监控"""
    
    def __init__(self):
        self.rounds = 0
        self.draft_tokens = 0
        self.accepted_tokens = 0
        self.generation_times = []
    
    def record_round(
        self, num_draft: int, num_accepted: int, time_ms: float
    ):
        self.rounds += 1
        self.draft_tokens += num_draft
        self.accepted_tokens += num_accepted
        self.generation_times.append(time_ms)
    
    @property
    def acceptance_rate(self) -> float:
        return self.accepted_tokens / max(1, self.draft_tokens)
    
    @property
    def speedup(self) -> float:
        """相对标准解码的加速比"""
        avg_accepted = self.accepted_tokens / max(1, self.rounds)
        return 1 + avg_accepted  # 每轮多生成 avg_accepted 个 Token
    
    @property
    def avg_latency(self) -> float:
        return sum(self.generation_times) / max(1, len(self.generation_times))
    
    def report(self) -> str:
        return (
            f"推测解码监控报告:\n"
            f"  总轮数: {self.rounds}\n"
            f"  接受率: {self.acceptance_rate:.1%}\n"
            f"  加速比: {self.speedup:.2f}x\n"
            f"  平均延迟: {self.avg_latency:.1f}ms/轮"
        )