首页›分类文章›人工智能›生成式AI与大模型›从投机采样到树状验证:大模型推测解码的技术演进与工程落地

从投机采样到树状验证:大模型推测解码的技术演进与工程落地

人工智能 · 生成式AI与大模型2026-10-020次浏览0 个评论

摘要

自回归语言模型的推理过程由一条严格的因果链驱动:每生成一个token都必须等待其前驱的产出,这使得单序列延迟与输出长度呈严格的线性正比。即便在单卡H100上运行Llama 3 70B,逐token前向传播的算子启动开销和传播延迟仍使GPU Tensor Core的有效利用率不足15%。推测解码(Speculative Decoding, 也称Speculative Sampling)通过"草稿-验证"范式打破了这一瓶颈:用一个轻量的草稿模型(或模型自身的辅助头)并行生成γ个候选token,再由目标大模型在单次前向中批量验证这些候选。只要候选分布与目标分布之间存在足够的重叠,验证阶段的批量计算开销会被线性生成节省的时间数倍覆盖。从Leviathan等人2022年的经典双模型方案,到Medusa的单模型多头投机,再到EAGLE使用第一层隐藏状态进行特征级自回归草稿生成,以及Sequoia、SpecInfer将候选序列扩展为树状结构以一次性验证多条候选路径,这条技术路线在过去三年内经历了三次明确的架构跃迁。每一次跃迁都减少了对独立草稿模型的依赖,同时提升了候选空间的并行度。本文将系统梳理推测解码从原理到落地的完整技术脉络,深入分析概率等价性保证的数学基础、特征级投机的工程优势、树状注意力掩码的设计约束,并系统考察其与vLLM、PagedAttention等生产推理系统的耦合细节。面向正在构建或优化大模型推理服务的工程师,本文力求提供一份可复现的决策框架。

第一章:自回归推理的延迟困局

1.1 逐token前向传播:Transformer解码的计算特性

Transformer在训练阶段是一个完整的矩阵乘法链:前向传播一次性处理整个序列,注意力矩阵的填充与连续矩阵乘法使GPU Tensor Core可以长时间满载运行,计算密集度极高。然而,在自回归生成阶段,输入矩阵的维度从[B,S,d][B, S, d]坍缩为[B,1,d][B, 1, d]——每次forward仅处理当前单个新token的嵌入向量。这使得算子启动(kernel launch)和中间结果在内存中的搬移时间,在延迟中占据了远超实际浮点运算的比例。

以Llama 3 70B为例,其层数L=80L=80,隐藏维度d=8192d=8192。单次前向传播的浮点运算量大约是2×L×d2≈10.7×10122 \times L \times d^2 \approx 10.7 \times 10^{12} FLOPs。在H100 SXM5(FP8 tensor core峰值约2000 TFLOPS)上,理想计算时间为10.7/2000≈5.410.7 / 2000 \approx 5.4毫秒。但实际测量中,单个token的前向延迟通常在30–60毫秒之间,实际利用率不足20%。差距主要来自:算子调度开销、KV Cache逐层读取的显存带宽瓶颈、以及张量并行下跨GPU通信的集体同步。

更关键的是,这30毫秒仅仅生成一个token。一个平均长度为500个token的回复,单序列延迟就达到15–30秒。将batch size增大可以缓解吞吐率,但对单序列延迟没有任何帮助——每个序列仍然必须等待自己的前一个token计算完成后才能启动下一步。

1.2 带宽瓶颈与算力浪费的定量分析

自回归解码阶段的计算强度(arithmetic intensity)极低。计算强度定义为单位数据搬运所对应的计算量。对于单次前向,

Arithmetic Intensity=FLOPsBytes of memory traffic\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{Bytes of memory traffic}}

在解码阶段,模型权重和KV Cache的读取量与正向浮点运算的比例远低于训练阶段。以FP8权重为例,权重大小约2×70×109≈352 \times 70 \times 10^9 \approx 35 GB(含优化器状态和激活值暂存后实际显存占用更高),而每次正向仅对极低维度的输入做一次矩阵乘法,计算/读取比处于显存带宽敏感区而非计算敏感区。NVIDIA的公开分析表明,在消费级到数据中心级GPU上,当batch中每层的序列长度较小且batch size适中时,解码延迟的60–80%消耗在显存带宽上,Tensor Core的实际占用时间往往不足20%。

graph LR
    subgraph 标准自回归解码
        X0[Token 0] -->|Forward| Y0[Output 0]
        Y0 --> X1[Token 1]
        X1 -->|Forward| Y1[Output 1]
        Y1 --> X2[Token 2]
        X2 -->|Forward| Y2[Output 2]
        Y2 --> X3[Token 3]

        style X0 fill:#FFE4B5
        style Y0 fill:#ADD8E6
        style X1 fill:#FFE4B5
        style Y1 fill:#ADD8E6
        style X2 fill:#FFE4B5
        style Y2 fill:#ADD8E6
    end

    subgraph GPU状态
        G0[Tensor Core: 5ms计算] --> G1[等待KV Cache加载: 25ms]
        G1 --> G2[Tensor Core: 5ms计算]
        G2 --> G3[等待KV Cache加载: 25ms]
        G3 --> G4[Tensor Core: 5ms计算]
        G4 --> G5[等待KV Cache加载: 25ms]
    end

上图展示了标准自回归解码的串行因果链。每一个token的生成严格依赖前一个token的输出结果,形成一条不可缩短的依赖序列。在GPU层面,这种串行意味着每次forward之间有25毫秒以上的等待时间消耗在显存带宽上,Tensor Core在短暂的5毫秒计算后再次进入等待。即便模型参数量继续增大,只要每次只处理单个token,带宽瓶颈不会自动缓解——模型越大,读取权重和KV Cache的显存搬运量反而越大,延迟更可能恶化。

1.3 现有吞吐优化的天花板

业界为缓解解码延迟已尝试了多条路径。

批处理(Batching) 将多个用户的请求合并为一个batch,在一次forward中同时处理多个序列。这提高了矩阵乘法的维度,使GPU利用率从[1,d][1, d]提升到[B,d][B, d]。但批处理解决的是"系统总吞吐量",而非"单序列延迟"。每个序列仍需等待自己的前驱token完成才能进入下一步,而且batch中不同序列长度不一致时,短的序列必须等待长序列的当前步完成,产生气泡。

连续批处理(Continuous Batching, 如vLLM / Orca / TGI) 通过动态调度解决了静态批处理的长度不一致问题:只要batch中有一条序列完成了自己的当前步,它就可以立即被同一条或另一序列的下一步替换,减少了空闲等待。但即便如此,单序列延迟仍然受限于GPU上单个token forward的物理延迟。

KV Cache压缩(如GQA、MQA、MLA)通过减少需要存储的Key和Value头数量来降低显存读取带宽。这在长序列场景下有效,但核心矛盾并未改变:每次forward的矩阵维度仍然是[B,1,d][B, 1, d],计算/读取比依然低下。

量化(INT8/INT4/FP8)减少了权重体积,降低了显存带宽消耗,但同时也引入了精度溢出风险和在长序列上的误差累积。

这些方案的共性在于:它们都在优化单次forward的开销,而没有减少forward的次数。如果能在一个forward中生成或验证多个token,则可以在不增加多少额外计算的前提下,成倍数缩短序列延迟。这正是推测解码的核心洞察。

第二章:推测解码的核心机制

2.1 核心思想:用"小步快跑"换"大步验证"

推测解码的基本范式由两个异构参与者组成:一个轻量的草稿模型(draft model,记作qq)和一个目标大模型(target model,记作pp)。

步骤如下:

  1. 草稿模型qq以自回归方式快速生成一个长度为γ\gamma的候选token序列 x1,x2,…,xγx_1, x_2, \ldots, x_\gamma。由于qq的参数量通常比pp小一个数量级以上(例如7B级别对70B级别),其每token延迟可能只有1–5毫秒,生成完整γ\gamma个token的总延迟仍在可接受范围内。
  2. 目标模型pp在一次前向传播中并行计算这γ\gamma个候选token的条件概率分布 p1(x),p2(x),…,pγ(x)p_1(x), p_2(x), \ldots, p_\gamma(x),以及草稿模型已经计算出的对应概率 q1(x),q2(x),…,qγ(x)q_1(x), q_2(x), \ldots, q_\gamma(x)。
  3. 验证阶段逐位决策:对于第ii个候选token,设草稿分布在该位置对候选token的概率为 qi(xi)q_i(x_i),目标分布为 pi(xi)p_i(x_i)。以概率
αi=min⁡(1,pi(xi)qi(xi))\alpha_i = \min\left(1, \frac{p_i(x_i)}{q_i(x_i)}\right)

接受该token。如果接受,继续验证第i+1i+1个;如果拒绝,则停止验证,并根据修正分布重新采样一个token作为该位置的确定输出,后续token全部丢弃。

  1. 无论接受了多少个候选token,目标模型都只需要一次forward即可完成验证。如果所有γ\gamma个token被接受,则草稿模型继续生成下一批γ\gamma个token;如果第kk个token被拒绝(k≤γk \le \gamma),则从该位置开始重新启动草稿生成。

2.2 概率等价性:输出分布与原始自回归完全一致

这是推测解码的理论基石。表面上看,草稿模型的介入似乎会改变最终输出token的分布,因为草稿模型qq和目标模型pp并非同一分布。然而,Leviathan等人(2022, "Fast Inference from Transformers via Speculative Decoding")和Chen等人(2023, "Accelerating Large Language Model Decoding with Speculative Sampling")独立证明:通过上述接受-修正规则,最终生成的token序列的联合分布与目标模型pp独立自回归生成的分布完全一致。

严格证明的核心逻辑如下。对于单个候选token xx,草稿模型采样 x∼q(⋅)x \sim q(\cdot)。验证阶段有两种可能:

  • 接受:以概率 α=min⁡(1,p(x)/q(x))\alpha = \min(1, p(x)/q(x)) 接受。该事件发生的概率为 q(x)⋅min⁡(1,p(x)/q(x))=min⁡(q(x),p(x))q(x) \cdot \min(1, p(x)/q(x)) = \min(q(x), p(x))。
  • 拒绝:以概率 1−α1 - \alpha 拒绝,并从修正分布 p′(x′)=max⁡(0,q(x′)−p(x′))/Zp'(x') = \max(0, q(x') - p(x')) / Z 中重新采样。具体来说,修正分布定义为:
p′(x′)=p(x′)−min⁡(p(x′),q(x′))1−∑zmin⁡(p(z),q(z))p'(x') = \frac{p(x') - \min(p(x'), q(x'))}{1 - \sum_z \min(p(z), q(z))}

但一个更简洁而等价的表达是:当拒绝时,从如下分布中采样 x∗∼p∗x^* \sim p^*,其中

p∗(x∗)=max⁡(0,p(x∗)−q(x∗))1−∑zmin⁡(p(z),q(z))p^*(x^*) = \frac{\max(0, p(x^*) - q(x^*))}{1 - \sum_z \min(p(z), q(z))}

那么,输出token最终来自p(x)p(x)的总概率可以证明严格成立:

P(output=x)=min⁡(p(x),q(x))+∑y:p(y)>q(y)p(x)−q(x)1−∑zmin⁡(p(z),q(z))⋅max⁡(0,p(y)−q(y))\mathbb{P}(\text{output} = x) = \min(p(x), q(x)) + \sum_{y: p(y) > q(y)} \frac{p(x) - q(x)}{1 - \sum_z \min(p(z), q(z))} \cdot \max(0, p(y) - q(y))

经过代数化简,上式恒等于 p(x)p(x)。这意味着,对于每一个token位置,推测解码的输出分布与目标模型pp的分布完全相同,没有任何偏差。这一点至关重要:它意味着推测解码不是一个近似加速技术,而是一个精确加速技术——在获得速度提升的同时,输出质量不会有任何下降。

2.3 双模型投机的工程约束

经典方案要求草稿模型qq与目标模型pp在tokenizer、词汇表和上下文长度上兼容,且qq必须与pp在分布上有足够重叠。在实践中,草稿模型的选择面临以下约束:

  • 同系列小模型:如用Llama 3 8B为Llama 3 70B生成草稿。这种方法实现简单,分布重叠度较高,但显存上需要同时加载两个模型。
  • 层数裁剪/深度蒸馏:将目标模型的前若干层取出作为qq,仅用于生成候选token,验证时仍使用完整模型。这节省了额外加载,但分布质量通常弱于同系列小模型。
  • 独立蒸馏小模型:训练一个专门为推测解码优化的超小Transformer(如124M参数级别),通过蒸馏目标模型的输出分布。成本更高,但延迟最低。

在vLLM等生产推理引擎中,双模型方法通常要求草稿模型驻留在与目标模型相同的张量并行组中。如果显存允许,两张模型可以各占用一半的GPU;如果显存紧张,则可能需要模型cpu-offload或交替加载,反而引入额外延迟。

graph TD
    subgraph Draft阶段
        D0[当前上下文] --> D1[草稿模型 q 生成]
        D1 --> |逐个自回归| D2[候选序列 1..γ]
    end

    subgraph 验证阶段 - 单次Target Forward
        T0[当前上下文 + 候选序列] --> T1[目标模型 p 并行计算]
        T1 --> T2[获取 p_1..p_γ 与 q_1..q_γ]
    end

    subgraph 逐位决策
        A[候选 token 1] --> B{α = min(1, p₁/q₁)}
        B -->|接受| C[候选 token 2]
        C --> D{α = min(1, p₂/q₂)}
        D -->|接受| E[候选 token 3]
        D -->|拒绝| F[从修正分布重采样]
        B -->|拒绝| F
        E --> G{α = min(1, p₃/q₃)}
        G -->|拒绝| F
        G -->|接受| H[全部接受, 启动下一轮]
        F --> I[以确定输出为起点, 回到Draft阶段]
        H --> J[候选序列成为正式输出, 回到Draft阶段]
    end

    D2 --> T0
    T2 --> A
    I --> D0
    J --> D0

上图描绘了双模型推测解码的完整流程。上半部分为草稿阶段:轻量模型qq以常规自回归方式串行生成长度为γ\gamma的候选序列。中间为验证阶段:目标模型pp对上下文拼接候选序列进行一次完整的前向传播,同时得到每个位置的目标分布和草稿分布概率。下半部分为逐位决策阶段:系统从左到右依次比对每个候选token,以α=min⁡(1,pi/qi)\alpha = \min(1, p_i/q_i)的概率接受;一旦出现拒绝事件,则从该位置按照修正分布重新采样一个token,后续所有候选被丢弃,整个流程重新从新的上下文启动下一轮草稿生成。

在理想情况下,如果草稿模型的准确率达到80%,取γ=5\gamma=5,则期望接受长度为∑i=150.8i≈3.36\sum_{i=1}^{5} 0.8^i \approx 3.36个token。目标模型一次forward的平均产出约为3.36个有效token,等效速度提升约3.36倍。这是理论上限;实践中,由于分布差异和回退开销,实际加速比通常在1.5–2.5倍之间。

第三章:从单模型到树状验证的演进

3.1 单模型多头投机:Medusa

双模型方案的一个明显缺点是草稿模型需要额外的显存与加载成本。2023年,Medusa提出:不引入独立草稿模型,而是在目标大模型的顶部添加多个并行的解码头,每个头独立负责预测未来不同偏移的token。

具体来说,在标准语言模型输出头的基础上,Medusa增加KK个额外的头部,分别预测当前位置向后第1、第2、...、第KK个token的概率分布。这些头部在微调阶段被训练,输入是主模型最后一层的隐藏状态。在推理时,主模型正常完成一次前向后,不仅获得当前位置的token概率,还同时获得未来KK个位置的独立预测。这些KK个预测被组织成树状结构,并通过树状注意力机制在一次目标模型forward中批量验证。

Medusa的优势在于零额外模型加载和高度集成的特性。缺点是需要对目标模型进行微调,且当KK较大时,额外头部的参数量不可忽略。后续Medusa-2进一步提出了基于拒绝采样的训练目标,使分布质量进一步提升。

3.2 Lookahead Decoding:Jacobi迭代视角

2023年末,Lookahead Decoding从一个完全不同的数学视角重新形式化了投机问题。它将token序列生成视为一个非线性方程的定点迭代问题。

在一个自回归模型中,位置tt的输出 yty_t 是历史 y<ty_{<t} 的函数: yt=f(y1,y2,…,yt−1)y_t = f(y_1, y_2, \ldots, y_{t-1})。这可以看作一个非线性动力系统。Lookahead Decoding提出使用Jacobi/Gauss-Seidel迭代来并行生成token:给定当前对整条序列的猜测,用模型并行地修正每个位置的猜测值,反复迭代直至收敛。这与线性方程组的Jacobi迭代完全同构。

具体实现中,系统维护一组n-gram窗口和局部缓存,利用历史序列中出现过的token片段作为自然候选(而非由独立草稿模型生成),并通过lookahead tree快速构建候选集合。这一方案的独特优势是无需任何额外模型或微调:它完全依赖目标模型自身以及历史序列的重复模式来产生候选。缺点在于候选质量对文本重复性的依赖较强,在生成多样性高的内容(如代码、数学推理)时加速效果有限。

3.3 特征级自回归投机:EAGLE与EAGLE-2

从双模型(Speculative Decoding)到单模型多头(Medusa)的演进减少了一个独立模型的成本,但仍然在token层面进行自回归生成草稿。2024年,EAGLE("Extrapolation Acceleration via Lookahead Decoding with Layer Embedding")引入了一个更本质的改进:

第tt个token的下一个token概率,可以被第tt个token的第一层隐藏状态以极高精度预测。

这意味着,你不需要等待目标模型跑完全部80层来获得token tt的输出,而只需要在第1层得到隐藏状态ht(1)h_t^{(1)},便可以用一个极小的辅助网络以极高的准确率预测token t+1t+1。与Medusa不同,EAGLE的草稿网络不是接在最后一层,而是接在第一层。

更具体地说,在深度Transformer中,第1层的隐藏状态已经捕获了词嵌入和局部上下文的大部分信息,足以用以预测下一个token。EAGLE的实验表明,一个仅有原模型0.5%参数量的轻量FFN网络接在第1层,其对未来1个token的预测精度可达92%以上(对比Token-level Medusa的约85%)。当扩展到预测2–3个未来token时,EAGLE的精度衰减也远小于Medusa。

这是因为Medusa的辅助头需要从最后一层预测未来token,而最后一层的状态已经高度"特化"于当前位置,包含大量与下一位置无直接关联的信息;相比之下,第1层的状态更接近下一token的"自然输入",预测任务更简单,所需辅助网络也更小。

EAGLE-2进一步引入了一个语境感知的草稿策略:在每次草稿生成之前,系统先评估当前上下文的"可预测性"——如果上下文的确定性高(如格式化文本、代码块、常见短语),则增大投机深度γ\gamma;如果上下文高度不确定(如开放式创意文本),则缩小γ\gamma甚至关闭投机。这可以通过一个轻量分类器实现,将平均接受率再提升8–12%。

3.4 树状注意力验证:候选空间的指数级扩展

无论是双模型投机还是单模型多头投机,早期方案都按一条线性序列生成和验证候选:草稿生成一条候选链,目标模型从左到右验证,遇到拒绝即终止。这浪费了目标模型的并行验证能力:一次forward可以并行计算整个候选序列的每个位置的概率,但验证逻辑只接受最长前缀。

树状验证将候选空间从一条链推广为一棵token tree。系统不再只生成一种候选走线,而是在每个节点分叉多个可能的下一token,形成一棵分支树。目标模型在验证阶段使用Tree Attention机制——一种经过改造的因果注意力掩码,允许不同的叶子节点共享共同的前缀路径,同时保证各叶子节点之间不产生跨路径的信息泄漏。

Tree Attention的核心设计如下:假设草稿阶段生成了一棵token树,每个节点代表一个候选token,从根到叶是一条候选路径。在目标模型的前向传播中,系统将所有树中节点按拓扑序平铺为一条序列,但注意力掩码的设计使得每个叶子只可以看到从根到自身路径上的节点,而不能看到其他分支上的兄弟节点。这保证了一次forward可以同时计算多条候选路径每个位置的概率。

Sequoia("Sequoia: Scalable, Robust, and Hardware-aware Speculative Decoding")将这一思路推进到硬件感知层面:它根据当前GPU的内存带宽、Tensor Core算力和KV Cache状态,动态构造最优的树拓扑(分支因子和深度),使一次目标模型forward验证的候选token数最大化。在H100上,Sequoia的实测加速比达到2.8–3.4倍,相比链式验证的1.5–2.0倍有显著提升。

graph TD
    subgraph 单链投机 2018-2022
        S0[T0] --> S1[T1a]
        S1 --> S2[T2a]
        S2 --> S3[T3a]
        S3 --> S4[T4a]
    end

    subgraph 多头投机 2023
        M0[T0] --> M1[T1a]
        M0 --> M12[T1b]
        M1 --> M2[T2a]
        M12 --> M22[T2b]
        M1 -.共享验证-> V1[1 Target Forward]
    end

    subgraph 树状验证 2024
        T0[Root] --> T1a[T1a]
        T0 --> T1b[T1b]
        T0 --> T1c[T1c]
        T1a --> T2a[T2a]
        T1a --> T2b[T2b]
        T1b --> T2c[T2c]
        T2a --> T3a[T3a]
        T2b --> T3b[T3b]
        T3a --> T4a[Leaf 1]
        T3b --> T4b[Leaf 2]
        T2c --> T4c[Leaf 3]
        T1c --> T3d[Leaf 4]
        T0 --> V2[1 Target Forward 验证全部叶子]
    end

    style V1 fill:#90EE90
    style V2 fill:#90EE90

上图对比了三种投机验证范式在候选空间并行度上的本质差异。左侧的单链投机中,一次目标模型forward最多只能按前缀顺序验证一条线性链,一旦遇到拒绝则后续全部作废。中间的多头投机通过引入并行的未来预测头,使验证阶段可以比较多个候选路径,但每个路径仍独立验证。右侧的树状验证将候选组织为一棵分支树,借助Tree Attention掩码在一次forward中并行评估所有叶子节点的候选序列,极大提升了目标模型每步计算的有效产出。并行候选token数量从O(γ)O(\gamma)线性增长跃迁为O(Bd)O(B^d)的树状指数增长,其中BB是分支因子,dd是树深度。Sequoia的实验表明,当树中有效叶子数达到8–12时,相比单链的平均加速比从约1.8倍提升到约3.2倍。

第四章:工程落地的关键权衡

4.1 草稿模型选择的三角困境

部署推测解码时,第一个工程决策是草稿来源。三种主流方案各有明确的折中:

独立同系列小模型(如8B → 70B):接受率通常最高(60–80%),因为分布重叠度好;显存代价最大(需要同时容纳两个模型);维护成本高(小模型需跟随大模型版本同步更新)。

模型自带多头(Medusa):无需额外模型,显存代价低;但需要微调(通常需数千GPU小时);接受率依赖于微调质量(通常50–70%)。

第一层特征网络(EAGLE):无额外模型加载,草稿网络极小(通常占原模型0.3–0.5%参数);接受率可高达75–90%(尤其在格式化文本上);但对特征质量的假设(第1层隐藏状态足以预测下一token)在部分模型架构上可能不成立,例如在使用了复杂嵌入层或前置卷积/位置变换的变体上精度会下降。

无模型方案(Lookahead / 检索增强草稿):零额外开销和零维护;但接受率对文本类型极度敏感(30–70%波动大)。

4.2 与连续批处理和分页注意力的兼容性

推测解码并非孤立于推理引擎之外的技术。在生产环境中,它必须与以下机制深度耦合:

分页注意力(PagedAttention):vLLM通过块表将KV Cache分块存储到非连续物理显存中。在推测解码中,当草稿序列被部分接受、部分拒绝时,KV Cache需要能够回退(rollback)到接受前缀的最后一个位置。PagedAttention的分块存储天然适合这种操作:只需将已接受的块保留,将后续块标记为空闲。但调度器必须能够追踪每条序列当前处于"草稿深度"dd(即当前正在验证的候选长度)和"已确认长度"ll,在回退时高效重建设备上的注意力状态。

连续批处理的调度复杂性:在静态批处理中,所有序列在验证阶段结束后同步回退或前进,调度简单。但在连续批处理中,batch中的不同序列可能在不同时间点完成验证——一条序列可能全部接受并立即需要下一轮草稿生成,而另一条序列在第2个token就被拒绝并需立即回退重采样。调度器必须动态将完成验证的序列从batch中移出或重新填充,并保证回退态分配不产生碎片。SpecServe("SpecServe: Serving Speculative Decoding on vLLM")是vLLM生态中对这一问题的系统性解答,它引入了一个两阶段调度器:第一阶段并行执行所有序列的草稿生成阶段,第二阶段将所有就绪序列合并为一个大的验证batch,统一做一次目标模型forward,然后各自回退或确认。

Sarathi-Serve的Chunked-Prefills与推测解码的冲突:Sarathi-Serve通过将prefill阶段也chunk化来提高调度利用率,但这与推测解码的"一次验证整段候选"设计存在冲突:如果prefill被chunk分割,目标模型在验证时只能看到部分上下文,候选的接受率会下降。当前主流方案是在decode阶段启用推测解码,prefill阶段不启用。

4.3 接受率的实测光谱与延迟模型

接受率(acceptance rate)是推测解码性能的唯一核心指标。它并非模型本身的静态属性,而是随任务类型、序列阶段、上下文不确定性而动态变化。

在代码生成任务中,由于代码具有高度结构化且强依赖前文(括号匹配、API调用链、缩进结构),草稿模型的预测质量很高。GitHub Copilot的内部数据显示,推测解码在代码补全场景下的(token-level)接受率可达85–92%,对应有效加速比约3.0–3.5倍。

在开放式对话或创意写作中,由于下一token的熵很高,接受率通常降至45–65%,对应加速比约1.4–1.8倍。

在数学推理或逻辑推问答任务中,由于token选择对精度高度敏感,一个错误的token会导致整段逻辑崩溃,目标模型倾向于更频繁地拒绝草稿候选,接受率通常低至35–55%,加速比仅1.2–1.5倍。

理论上,设单步接受率为α\alpha,最大候选深度为γ\gamma。如果各步独立,期望接受长度

E[L]=1−αγ+11−αE[L] = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha}

目标模型每步有效产出token数为 E[L]/1E[L] / 1(因为目标模型仍只forward一次),加速比大约≈E[L]\approx E[L]。若同时考虑草稿生成的延迟TdraftT_{draft}和目标模型forward延迟TtargetT_{target},总加速比修正为:

Speedup=E[L]⋅TseqTdraft(γ)+Ttarget\text{Speedup} = \frac{E[L] \cdot T_{seq}}{T_{draft}(\gamma) + T_{target}}

其中TseqT_{seq}为标准自回归生成E[L]E[L]个token的总延迟。因为TtargetT_{target}通常为主导项,上式近似为≈E[L]/(1+Tdraft/Ttarget)\approx E[L] / (1 + T_{draft}/T_{target})。若草稿模型足够小,Tdraft≪TtargetT_{draft} \ll T_{target},则加速比趋近E[L]E[L]。

4.4 回退策略与计算浪费

当一条候选链长度为γ=5\gamma=5,但仅前2个token被接受、第3个被拒绝时,第3–5个token的草稿计算和对应的KV Cache写入被完全浪费。在树状验证中浪费更复杂:若某分支的叶子被拒绝,其该分支路径上的部分KV Cache虽然被共享于其他分支而未必完全浪费,但叶子节点自身的计算仍被丢弃。

工程上控制浪费的核心机制:

  1. 动态γ\gamma调整:根据实时测量的接受率动态调整候选深度。如果最近10步的平均接受率低于阈值,立即缩小γ\gamma。
  2. 置信度剪枝:草稿生成阶段引入提前退出机制。如果草稿模型对当前token的置信度低于阈值,立即停止后续生成,将已生成的短链送交验证。这减少了低质量长链被大面积拒绝的浪费。
  3. 三级内存池:为草稿阶段的KV Cache分配独立于主推理的内存池。回退时直接释放整个草稿池,无需逐token恢复,降低管理复杂度。
graph LR
    subgraph 输入层
        REQ[用户请求队列] --> SCHED[推测解码调度器]
    end

    subgraph 调度核心
        SCHED --> |解码阶段| DEC{批量验证准备}
        DEC --> DRAFT[草稿生成阶段\n(轻量模型/多头/特征网络)]
        DRAFT --> TREE[构建候选树/链]
        TREE --> VERIFY[目标模型并行验证\n(Tree Attention / 序列前缀验证)]
        VERIFY --> ACCEPT{逐位/逐叶决策}
        ACCEPT -->|全接受| NEXT[推进序列, 启动下一轮]
        ACCEPT -->|部分接受| PART[截断到接受前缀, 重启动]
        ACCEPT -->|全拒绝| REJ[以修正采样点重启]
        PART --> ROLL[KV Cache回退\n+ 分块内存重分配]
        REJ --> ROLL
    end

    subgraph 输出层
        ROLL --> SCHED
        NEXT --> OUT[输出流聚合]
        OUT --> USR[返回用户/下游]
    end

上图勾勒了生产级推理系统中推测解码与调度器深度耦合的架构。用户请求进入调度器后,系统首先判断当前序列是否处于解码阶段——若是,则启动草稿生成(可能使用轻量模型、多头预测头或EAGLE特征网络),并将候选序列组织为链或树状结构。随后,目标模型对这些候选执行一次并行的批量验证。验证结果通过接受-修正规则逐位(链式)或逐叶(树状)决策,输出被确认的token前缀,并将拒绝位置之后的状态全部丢弃。KV Cache的回退与分页内存的重分配由内存管理层统一处理,回退后的序列重新进入调度器的就绪队列,等待下一轮调度。整个循环的调度效率、回退速度和内存碎片管理,决定了推测解码在实际生产环境能否兑现其理论加速。

第五章:前沿与展望

5.1 与KV Cache压缩的协同

推测解码与MLA/GQA/MQA等KV Cache压缩技术并非互斥。事实上,两者可以叠加获得复合收益:KV Cache压缩降低了单次forward的显存带宽瓶颈,使目标模型的验证forward更快;而推测解码减少了forward的总次数。两者协同可以将单序列端到端延迟压缩到原来的1/4–1/6。DeepSeek-V3的推理栈便同时采用了MLA和投机解码,在128K长上下文场景下实现了业界领先的单用户延迟。

5.2 硬件感知投机:利用张量并行中的闲置算力

在多GPU张量并行(TP)部署中,目标模型forward过程中的all-reduce通信会引入同步等待。在这段时间内,GPU的计算资源处于空闲。一种前沿思路是:在等待all-reduce完成时,让各GPU上的rank利用闲置算力运行本地草稿小模型或辅助头。这样,草稿生成与目标模型forward中的通信间隙可以被重叠,提升端到端GPU利用率。这是DistServe、SpecServe等系统在开发中的方向。

5.3 Prefill阶段的投机可行性

当前推测解码几乎完全聚焦于自回归解码阶段(decode phase)。但预填充阶段(prefill phase,即处理用户输入的前向传播)同样存在优化空间:输入token序列之间的依赖关系是已知的(全部在用户提交时即已确定),因此可以被并行处理。已有工作如"Prompt Lookup Decoding"利用用户输入文本中的重复n-gram,在prefill阶段构建候选并验证,从而缩短首token时间(Time to First Token, TTFT)。将prefill投机与decode投机统一为一个框架,是vLLM社区正在讨论的方向。

5.4 测试时计算与解码时并行化的交汇

自2024年以来,"测试时计算"(Test-Time Compute)或"延长推理链"(Long CoT / Chain-of-Thought)的方法论通过让模型在回答前生成更多中间思考token来提升输出质量,如OpenAI的o1/o3系列和DeepSeek-R1的强化学习推理链。这一方向与推测解码之间存在张力:更多的中间token意味着推理链更长,延迟上升;但中间思考token的结构化程度通常高于自由生成文本(如"思考...结论..."),这恰好为投机解码提供了高接受率的土壤。事实上,DeepSeek-R1的推理链在启用EAGLE级投机后,端到端延迟被压缩了约2.2倍,接受率达78%。

另一条交汇方向是自修正投机:系统在生成一个初步回答后,用草稿模型快速生成多个修正版本,然后由目标模型选择或综合最佳答案。这进入了一个介于投机解码与多采样聚合(best-of-N)之间的灰色地带,可能会在2025–2026年的生产系统中出现。

第六章:结论

推测解码的核心贡献在于,它通过改变验证的组织方式而非改变模型本身,在不牺牲输出分布精度的前提下,将自回归语言模型的单序列延迟降低了1.5到3倍。从早期的双模型方案到单模型多头投机,再到特征级自回归(EAGLE)和树状验证(Sequoia/SpecInfer),每一步改进都在减少额外依赖的同时提升候选并行度。

对于正在规划部署的工程团队,以下决策树可以提供参考:

  • 显存充裕且维护团队足够:采用同系列小模型作为草稿,直接复用成熟的vLLM + SpecServe方案,预期加速比2.0–2.5倍。
  • 显存紧张且不愿维护额外模型:采用Medusa或EAGLE方案。EAGLE在格式化/代码/推理链场景下接受率更高,Medusa集成度更好,vLLM社区已有初步支持。
  • 对延迟极度敏感且接受率波动可容忍:采用无模型Lookahead方案或Prompt Lookup,零额外成本,但需接受场景依赖的加速波动。
  • 追求极致加速且愿意投入自研:采用树状验证+硬件感知调度(如Sequoia),预期加速比3.0–3.5倍,但Tree Attention实现和调度逻辑复杂度较高。

无论选择哪条路径,操控的核心变量始终是接受率。任何投机系统的工程优化,最终都应围绕提升和稳定接受率展开。这与推测解码的理论保证一样清晰:只要接受率足够高,速度提升就近乎确定。

评论

0

评论加载中…

发表评论

0/2000