让大模型跑得快一点：Speculative Decoding 实战与加速比分析

一、为什么 Decode 阶段那么慢？

大模型推理其实就两件事：Prefill（把提示词塞进去）和 Decode（一个字一个字往外吐）。

Prefill 阶段是计算密集型，GPU 吃得饱饱的；但到了 Decode 阶段，情况就变了。每生成一个字，GPU 都得把整个模型的权重从显存里读一遍。以 70B 模型为例，一次 Decode 的计算量只有 0.01 TFLOPS，但光读权重就要 140GB 带宽。这就导致 GPU 大部分时间都在等数据，计算单元空转。

更麻烦的是自回归的串行特性：第 N 个字必须等第 N-1 个字生成完才能开始。不管 GPU 算力多强，Decode 阶段的延迟始终被内存带宽和串行步数卡着。

Speculative Decoding（投机解码）就是为了解决这个问题：它试图打破这种串行依赖，一次吐出多个字。

二、投机解码是怎么工作的？

核心逻辑很简单：找个“小机灵鬼”（Draft Model）先猜几个字，然后让“老大哥”（Target Model）一次性验证这些猜测对不对。

sequenceDiagram
    participant Draft as Draft Model<br/>(小模型, 快速)
    participant Target as Target Model<br/>(大模型, 精确)
    participant Output as 输出

    Note over Draft: Step 1: Draft 模型快速生成 K 个候选 Token
    Draft->>Draft: 生成 t1, t2, t3, t4, t5
    Draft->>Target: 传入候选序列

    Note over Target: Step 2: Target 模型单次前向验证
    Target->>Target: 并行计算 P(t1), P(t2), P(t3), P(t4), P(t5)
    Target->>Target: 同时生成 Target 的 P(t|prefix)

    Note over Target: Step 3: 逐个验证候选 Token
    Target->>Target: t1: P_draft=0.8, P_target=0.9 → 接受
    Target->>Target: t2: P_draft=0.7, P_target=0.6 → 接受
    Target->>Target: t3: P_draft=0.5, P_target=0.3 → 拒绝

    Note over Output: 输出 t1, t2 + Target 修正的 t3'
    Target->>Output: t1, t2, t3'
    Note over Draft: 从 t3' 重新开始投机

具体流程分三步：

1. Draft 模型快速生成 K 个候选 Token：小模型速度快，几秒钟就能吐出几个字。
2. Target 模型单次前向验证：大模型把 Draft 生成的序列一次性跑一遍，并行计算每个位置的概率。
3. 逐个验证候选 Token：对比 Draft 和大模型的概率，决定接受还是拒绝。

数学上的保证：通过特定的接受概率计算，Speculative Decoding 能保证输出分布和原始自回归解码完全一致。也就是说，它只是加速，不会改变生成质量。

接受概率的计算逻辑：

• 如果 Draft 的概率 $p_d(t_i) \leq p_t(t_i)$，直接接受。
• 如果 $p_d(t_i) > p_t(t_i)$，以 $p_t(t_i) / p_d(t_i)$ 的概率接受。
• 拒绝时，从修正分布 $\max(0, p_t - p_d) / \sum \max(0, p_t - p_d)$ 中采样一个 Token 作为修正。

三、工程实现代码

# speculative_decoding.py — 投机解码引擎
import time
from dataclasses import dataclass, field
from typing import Optional

import numpy as np


@dataclass
class DraftResult:
    """Draft 模型的生成结果"""
    tokens: list[int]
    log_probs: list[float]       # 每个 Token 的 log 概率
    latency_ms: float


@dataclass
class VerifyResult:
    """Target 模型的验证结果"""
    accepted_count: int          # 被接受的 Token 数
    rejected_at: int             # 拒绝位置（-1 表示全部接受）
    corrected_token: int         # 拒绝时修正的 Token
    corrected_log_prob: float    # 修正 Token 的 log 概率
    target_log_probs: list[float]  # Target 模型对每个位置的 log 概率
    latency_ms: float


class SpeculativeDecoder:
    """投机解码引擎"""

    def __init__(self, draft_model_fn, target_model_fn,
                 speculate_length: int = 5,
                 temperature: float = 1.0):
        self._draft_fn = draft_model_fn
        self._target_fn = target_model_fn
        self.speculate_length = speculate_length
        self.temperature = temperature

        # 统计信息
        self._stats = {
            "total_tokens_generated": 0,
            "total_draft_tokens": 0,
            "total_accepted_tokens": 0,
            "total_target_calls": 0,
            "total_draft_calls": 0,
        }

    def generate(self, prompt_tokens: list[int],
                 max_tokens: int = 256) -> list[int]:
        """使用投机解码生成文本"""
        generated = list(prompt_tokens)

        while len(generated) - len(prompt_tokens) < max_tokens:
            remaining = max_tokens - (len(generated) - len(prompt_tokens))
            k = min(self.speculate_length, remaining)

            # Step 1: Draft 模型快速生成 K 个候选 Token
            draft_result = self._draft_generate(generated, k)
            self._stats["total_draft_calls"] += 1
            self._stats["total_draft_tokens"] += len(draft_result.tokens)

            # Step 2: Target 模型单次前向验证
            verify_result = self._verify(generated, draft_result)
            self._stats["total_target_calls"] += 1

            # Step 3: 根据验证结果更新生成序列
            if verify_result.rejected_at == -1:
                # 全部接受
                generated.extend(draft_result.tokens)
                self._stats["total_accepted_tokens"] += len(draft_result.tokens)

                # Target 模型在最后一个位置也生成了一个 Token
                # 可以额外获取 1 个 Token
                bonus_token = self._sample_from_logits(
                    verify_result.target_log_probs[-1]
                )
                generated.append(bonus_token)
                self._stats["total_accepted_tokens"] += 1
            else:
                # 部分接受
                accepted_tokens = draft_result.tokens[:verify_result.rejected_at]
                generated.extend(accepted_tokens)
                generated.append(verify_result.corrected_token)
                self._stats["total_accepted_tokens"] += len(accepted_tokens) + 1

            self._stats["total_tokens_generated"] = \
                len(generated) - len(prompt_tokens)

        return generated

    def _draft_generate(self, prefix: list[int],
                        k: int) -> DraftResult:
        """Draft 模型生成 K 个候选 Token"""
        start = time.time()
        tokens = []
        log_probs = []

        current = list(prefix)
        for _ in range(k):
            # 调用 Draft 模型获取下一个 Token 的分布
            logits = self._draft_fn(current)
            probs = self._softmax(logits / self.temperature)
            token = np.random.choice(len(probs), p=probs)
            log_prob = float(np.log(probs[token] + 1e-10))

            tokens.append(int(token))
            log_probs.append(log_prob)
            current.append(token)

        latency = (time.time() - start) * 1000
        return DraftResult(tokens=tokens, log_probs=log_probs,
                           latency_ms=latency)

    def _verify(self, prefix: list[int],
                draft: DraftResult) -> VerifyResult:
        """Target 模型验证候选 Token"""
        start = time.time()

        # 构造验证输入：prefix + draft tokens
        verify_input = prefix + draft.tokens

        # Target 模型单次前向传播
        # 返回每个位置的 logits（包括 draft tokens 的位置）
        all_logits = self._target_fn(verify_input)

        # 提取 draft tokens 位置的 logits
        # 位置 i 对应 prefix 长度 + i - 1 处的 logits
        # 因为位置 i 的 Token 是由前 i-1 个 Token 预测的
        prefix_len = len(prefix)
        target_log_probs = []

        for i in range(len(draft.tokens)):
            pos = prefix_len + i - 1
            if pos < 0:
                pos = 0
            logits = all_logits[pos]
            probs = self._softmax(logits / self.temperature)
            token = draft.tokens[i]
            target_log_probs.append(float(np.log(probs[token] + 1e-10)))

        # 逐个验证候选 Token
        rejected_at = -1
        corrected_token = -1
        corrected_log_prob = 0.0
        accepted_count = 0

        for i in range(len(draft.tokens)):
            draft_lp = draft.log_probs[i]
            target_lp = target_log_probs[i]

            # 计算接受概率
            # p_accept = min(1, p_target / p_draft)
            # 在 log 空间: log(p_accept) = min(0, target_lp - draft_lp)
            log_accept_ratio = target_lp - draft_lp

            if log_accept_ratio >= 0:
                # p_target >= p_draft，一定接受
                accepted_count += 1
            else:
                # p_target < p_draft，以概率 p_target/p_draft 接受
                accept_prob = np.exp(log_accept_ratio)
                if np.random.random() < accept_prob:
                    accepted_count += 1
                else:
                    # 拒绝：从修正分布中采样
                    rejected_at = i
                    pos = prefix_len + i - 1
                    if pos < 0:
                        pos = 0
                    corrected_token = self._sample_corrected(
                        all_logits[pos],
                        draft.tokens[i],
                        draft_lp,
                        target_lp,
                    )
                    corrected_log_prob = target_log_probs[i]
                    break

        latency = (time.time() - start) * 1000
        return VerifyResult(
            accepted_count=accepted_count,
            rejected_at=rejected_at,
            corrected_token=corrected_token,
            corrected_log_prob=corrected_log_prob,
            target_log_probs=target_log_probs,
            latency_ms=latency,
        )

    def _sample_corrected(self, logits: np.ndarray,
                          draft_token: int,
                          draft_lp: float,
                          target_lp: float) -> int:
        """从修正分布中采样 Token"""
        probs = self._softmax(logits / self.temperature)
        draft_prob = np.exp(draft_lp)

        # 修正分布: max(0, p_target - p_draft) / Z
        corrected = np.maximum(0, probs - draft_prob)
        total = np.sum(corrected)

        if total < 1e-10:
            # 修正分布退化为均匀分布，使用 Target 分布
            return int(np.random.choice(len(probs), p=probs))

        corrected /= total
        return int(np.random.choice(len(corrected), p=corrected))

    def _softmax(self, logits: np.ndarray) -> np.ndarray:
        """数值稳定的 Softmax"""
        shifted = logits - np.max(logits)
        exp_vals = np.exp(shifted)
        return exp_vals / (np.sum(exp_vals) + 1e-10)

    def _sample_from_logits(self, logits) -> int:
        """从 logits 中采样"""
        if isinstance(logits, (int, float)):
            return int(logits)
        probs = self._softmax(np.array(logits) / self.temperature)
        return int(np.random.choice(len(probs), p=probs))

    def get_stats(self) -> dict:
        """获取加速统计"""
        total_draft = self._stats["total_draft_tokens"]
        total_accepted = self._stats["total_accepted_tokens"]
        acceptance_rate = (
            total_accepted / total_draft * 100
            if total_draft > 0 else 0
        )

        # 加速比 = 平均每次 Target 调用生成的 Token 数
        target_calls = self._stats["total_target_calls"]
        avg_tokens_per_call = (
            total_accepted / target_calls
            if target_calls > 0 else 1.0
        )

        return {
            **self._stats,
            "acceptance_rate": round(acceptance_rate, 1),
            "avg_tokens_per_target_call": round(avg_tokens_per_call, 2),
            "effective_speedup": round(avg_tokens_per_call, 2),
        }

四、实际加速效果与 Draft 模型选择

投机解码能不能提速，主要看两点：Draft 模型的接受率，以及 Target 模型的验证开销。

接受率是关键：Draft 模型和大模型的输出分布越接近，接受率越高。实测中，如果用同一模型系列的 Draft（比如 LLaMA-7B 给 LLaMA-70B 当小弟），5-Token 投机的接受率大概在 70%-80%，平均加速比 2-2.5x。如果 Draft 模型和大模型完全不搭界，接受率可能掉到 40%-50%，加速比只有 1.3-1.5x，甚至不如直接跑自回归解码。

验证开销不能忽视：Target 模型验证时需要一次前向传播处理 K+1 个 Token（K 个候选 + 1 个额外位置）。虽然比 K 次单步 Decode 高效（KV Cache 可以复用），但 Prefill 阶段的计算量是随序列长度增长的。如果 K 设得太大，验证的 Prefill 开销可能会把投机带来的加速抵消掉。经验上，K 值在 4-7 之间比较合适。

Draft 模型怎么选：有两个硬指标——推理速度至少比 Target 快 3x，且输出分布要接近。最佳实践是用同一模型系列的较小版本。比如 LLaMA-70B 的 Draft 模型，用 LLaMA-7B 或 LLaMA-13B 都不错。如果 Draft 太小（比如 1B），速度快但接受率低；如果太大（比如 30B），接受率高但速度优势不明显。

适用场景：投机解码适合单请求低延迟场景（比如实时对话），这时候 Decode 阶段的串行性是主要瓶颈。如果是高并发批量推理，GPU 本来就能通过批量并行把算力吃满，投机解码的加速效果就有限了。

五、总结

Speculative Decoding 通过 Draft-Verify 机制打破了自回归解码的串行瓶颈，在保证输出分布无偏的前提下，通常能实现 2-2.5x 的加速。

加速比的核心决定因素是 Draft 模型的接受率，而接受率取决于 Draft 与 Target 模型的分布匹配度。建议用同一模型系列的较小版本作为 Draft 模型，投机长度设为 5。投机解码最适用于单请求低延迟场景，在高并发批量场景下收益有限。