通过序列蒙特卡洛加速LLM推理

# 通过序列蒙特卡洛加速 LLM 推理

来源: https://arxiv.org/html/2604.15672

Yahya Emara∗1,2 Mauricio Barba da Costa∗3 Chi-Chih Chang1 Cameron Freer3 Tim Vieira4 Ryan Cotterell4 Mohamed S. Abdelfattah1,2

1康奈尔大学 2Makora 3麻省理工学院 4苏黎世联邦理工学院

###### 摘要

推测解码(SD)通过从廉价的提议模型草拟令牌，并通过拒绝采样针对昂贵的目标模型进行验证来加速语言模型推理。由于拒绝在第一个错误处截断草稿块，当草稿和目标发散时吞吐量会下降。与其直接拒绝草稿令牌，我们建议对其重新加权。为此，我们引入了序列蒙特卡洛推测解码(SMC-SD)，它用重要性加权重采样替代令牌级拒绝。SMC-SD是一个原则性的近似推理方案，用精确性换取额外的速度，同时保持每步近似误差的理论界。由于LLM推理受内存带宽限制，并行草拟和评分粒子所需的算术运算几乎免费——SMC-SD使用闲置计算将验证转变为向量化、固定大小的操作，无需回滚。经验上，SMC-SD相比推测解码实现了2.36×的加速，相比自回归解码实现了5.2×的加速，同时在推理、指令跟随和编码基准上保持在目标模型精度的3%以内。

https://github.com/abdelfattah-lab/smcsd

## 1 引言

神经语言模型的自回归生成本质上是顺序的：每个令牌依赖于所有之前的令牌，因此生成单个序列需要与输出中令牌数量一样多的通过模型的串行前向传递。这个顺序瓶颈是更快推理的主要障碍。

| SGLang AR | SGLang SD | SSD | SMC-SD(本文) |
|---|---|---|---|
| 0.0 tok/s | 14.0 (1.4×) | 40.8 (2.2×) | 65.0 (1.0×) |
| | | 225.2 (3.5×) | 342.0 (5.2×) |

**图1**: SMC-SD在Llama 1B→70B草稿-目标对上相对于自回归基线的吞吐量加速，优化的基于树的SD (SGLang)、推测解码(SSD; Kumar等人, 2026)在ShareGPT数据集上。AR、SGLang SD、SMC-SD在4块H100 GPU上运行，SSD在5块H100 GPU上运行。

推测解码(SD; Leviathan等人, 2023)通过摊销目标模型调用的成本来解决这一瓶颈。其核心是一个拒绝采样器：一个小型**草稿**模型——选择使其评估成本显著更低——提议K个令牌，较大的**目标**模型在单个前向传递中验证所有K个令牌，接受一个前缀并拒绝其余部分。这个拒绝准则确保生成的字符串完全按照目标分布，而加速因子是随机的，取决于草稿和目标模型的对齐程度。

本文采取了一种根本不同的方法来加速自回归生成。与其寻找更快的精确采样器，我们设计了一个近似采样方案，其保真度可以与速度相权衡。关键修改是用重要性加权重采样替代SD的令牌级拒绝步骤，使验证步骤成为序列蒙特卡洛(SMC)的一个实例。我们的方法**反转**了SD的保证——近似质量是随机的，加速因子是确定性的。我们称我们的方法为**序列蒙特卡洛推测解码(SMC-SD)**，并在§3中具体描述它。

我们使用GPU的屋线模型(Roofline model)推导了SMC-SD的每秒令牌数(TPS)加速因子的闭式表达式，在§3.1中。我们展示了屋线模型在内存带宽限制和计算限制两个制度中给出了加速因子的简单表征。

从理论角度，因为SMC-SD最终基于重要性重采样，非渐近近似误差界可以使用已知技术获得。对于SMC-SD，我们展示了L2偏差和均方误差线性衰减于粒子数量，L1偏差在粒子数量的平方根处衰减，常数由草稿和目标模型之间的χ²散度控制(§3.2)。

从系统角度，SMC-SD呈现了一个新的计算模式，需要围绕它设计推理引擎。在§3.3中，我们使用硬件和算法交叉点的若干观察来指导高吞吐量SMC-SD推理引擎的设计。SMC-SD的草拟和验证阶段都使用可向量化的执行结构，自然映射到GPU硬件上，实现了比标准推测解码更大的并行性。对于KV缓存管理，我们观察到重采样步骤中所有KV缓存数据移动可以用高效的原地指针交换替代，我们通过利用粒子间的共享前缀来减少缓存大小。

经验上，SMC-SD在GSM8K、MATH500、AlpacaEval和DS1000上跨Llama和Qwen模型族实现了推测解码优化实现吞吐量的高达2.5×，同时保持在目标模型精度的3%以内。在多GPU设置中，SMC-SD相比最先进的推测解码实现了2.36×加速，相比自回归基线实现了5.2×加速。

图1总结了这些端到端吞吐量增益。

## 2 语言模型和推测解码

### 2.1 语言建模背景

令 𝒳 为有限字母表，一个非空有限集合。我们称 𝒳 的元素为**令牌**。我们用 𝒳* 表示 𝒳 上所有有限字符串的集合，包括空字符串 ε，并用 x 表示单个令牌，用 **x** = x₁···xₜ ∈ 𝒳* 表示字符串。我们使用缩写 x₁:ₖ 表示 x₁···xₖ。

语言模型是一个函数 pθ: 𝒳* → [0,1]，为任何有限字符串分配概率，使得 Σₓ ∈ 𝒳 pθ(x | **x**₁:ₜ₋₁) = 1 对所有 **x** ∈ 𝒳* 和所有 t。参数 θ 通常通过在大型文本语料库上训练神经网络获得。在本工作中，我们假设我们可以有效评估条件概率 pθ(x | **x**)。

**自回归采样。** 令 p 为语言模型。要从 p(x₁···xₜ) 采样长度为 T 的序列，我们顺序采样：对于 t = 1,2,...,T，样本 xₜ ~ p(xₜ | x₁:ₜ₋₁)。这需要 T 个前向传递，是推理中顺序瓶颈的来源。

### 2.2 推测解码

推测解码通过使用较小的**草稿**模型 q 来提议 K 个令牌，然后用较大的**目标**模型 p 验证它们，从而加速推理。该方案保证输出按照目标模型分布，同时减少目标模型的调用次数。

**标准推测解码算法。** 标准推测解码(SD)工作如下：

1. 从草稿模型采样 K 个令牌：x̂₁:ₖ ~ q(· | **x**)
2. 对目标模型的输出进行单次前向传递，获得所有提议令牌的目标概率
3. 对于 t = 1,2,...,K：
   - 如果 pₜ(x̂ₜ | **x**₁:ₜ₋₁) ≥ qₜ(x̂ₜ | **x**₁:ₜ₋₁)：接受 x̂ₜ，继续
   - 否则：使用拒绝采样接受替代令牌或拒绝 x̂ₜ，停止

这个过程保证输出根据目标分布分布，因为它实现了有效的拒绝采样。

## 推测解码的扩展

我们现在讨论推测解码的几种自然扩展，这些可能对特定应用有用。

### 软约束

许多应用需要对生成的输出施加软约束。一个自然的方法是引入加权函数 Ψ: 𝒳* → ℝ₊，并采样以下形式的序列：

π(x) ∝ p(x) · Ψ(x)

这对应于目标分布的改写。当可以计算 Ψ(x) 而不改变渐近复杂度时，SMC-SD可以直接适应。例子包括：

- **约束生成**(Ψ = 𝟙[x ∈ 𝒜])。可以在保持以约束条件的目标分布的同时执行硬约束——语法有效性、格式合规性、长度界限。先前工作(Lew等人, 2023; Loula等人, 2025; Xefteri等人, 2025)表明此类约束可以改进编码基准的精度；将它们纳入SMC-SD可以同时增加下游性能同时加速推理。

- **奖励加权解码**(Ψ = exp(β·R(x))，目标 π ∝ p·e^(βR))。这以KL正则化的最优策略为目标，来自RLHF而无需微调，与需要用同一基础模型朝向不同目标的代理设置相关。标准推测解码无法表达这些中的任何一个：它需要知道目标的归一化常数，这对该类分布不可用。

### 系统面

在系统方面，还有很多工作要做以实现SMC-SD所演示的理论屋线模型加速(如图3所示)。首先，有机会将重采样与草拟重叠以完全消除其性能开销。这可以通过预先推测ESS阈值的结果，并相应地开始下一轮草拟来完成。此外，草稿和目标模型可以跨硬件分解，并异步分派草拟和评分调用。

其次，通过允许 N 和 K 动态变化，SMC-SD可以适应底层硬件和工作负载特征，实现其可调可帕累托边界的运行时探索。

最后，SMC-SD反映了推理算法与非对称硬件缩放的协同设计的趋势——计算增长快于内存带宽(Zadouri等人, 2026)。例如，Blackwell GPU架构相对于前一代Hopper GPU显著增加了FLOPS，并保持内存带宽大致相等；这个增加的计算预算使其成为SMC-SD的理想目标。

## 作者贡献

- **Yahya Emara**([email protected])：主导系统设计/推理引擎开发、实验设计、形式分析(算术强度)、可视化、写作。

- **Mauricio Barba da Costa**([email protected])：研究构想、形式分析(加速、算术强度、近似误差)、实验设计(原型、幂采样)、可视化、写作。

- **Chi-Chih Chang**([email protected])：系统设计/推理引擎开发、可视化、写作。

- **Cameron Freer**([email protected])：高级项目领导、项目建议和指导、技术建议。

- **Tim Vieira**([email protected])：高级项目领导、项目建议和指导、技术建议。

- **Ryan Cotterell**([email protected])：形式分析(近似误差)、高级项目领导、项目建议和指导、技术建议、写作。

- **Mohamed Abdelfattah**([email protected])：高级项目领导、形式分析(加速、算术强度分析)、系统设计、项目叙事发展、项目建议和指导、技术建议、可视化、写作。

## 参考文献

- 重要性采样：内在维度和计算成本。统计科学32，第405-431页。外部链接：https://projecteuclid.org/journals/statistical-science/volume-32/issue-3/Importance-Sampling-Intrinsic-Dimension-and-Computational-Cost/10.1214/17-STS611.pdf
  引用于：§3.2

- S. Azizi, E. B. Potraghloo, M. Ahmadi, S. Kundu, 和 M. Pedram (2026) Power-SMC：用于无训练LLM推理的低延迟序列级幂采样。外部链接：2602.10273, https://arxiv.org/abs/2602.10273
  引用于：附录F

- G. Bachmann, S. Anagnostidis, A. Pumarola, M. Georgopoulos, A. Sanakoyeu, Y. Du, E. Schönfeld, A. Thabet, 和 J. Kohler (2025) Judge解码：更快的推测采样需要超越模型对齐。arXiv预印本 arXiv:2501.19309。外部链接：https://arxiv.org/abs/2501.19309
  引用于：§5

- T. Cai, Y. Li, Z. Geng, H. Peng, J. D. Lee, D. Chen, 和 T. Dao (2024) Medusa：具有多个解码头的简单LLM推理加速框架。外部链接：2401.10774, https://arxiv.org/abs/2401.10774
  引用于：§5

- C. Chen, S. Borgeaud, G. Irving, J. Lespiau, L. Sifre, 和 J. Jumper (2023) 用推测采样加速大语言模型解码。外部链接：2302.01318, https://arxiv.org/abs/2302.01318
  引用于：附录A，§1，§2.2

- J. Chen, Y. Liang, 和 Z. Liu (2026) DFlash：用于闪存推测解码的块扩散。外部链接：2602.06036, https://arxiv.org/abs/2602.06036
  引用于：§5

- K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, 和 J. Schulman (2021) 训练验证器解决数学词问题。外部链接：2110.14168, https://arxiv.org/abs/2110.14168
  引用于：§1，§4

- P. Del Moral, A. Doucet, 和 A. Jasra (2006) 序列蒙特卡洛采样器。皇家统计学会期刊B系列：统计方法68(3)，第411-436页。
  引用于：§1，§3.2，§3，§6

- A. Doucet, N. De Freitas, N. J. Gordon，等人 (2001) 序列蒙特卡洛方法实践。卷1，Springer出版社。
  引用于：§1，§3

- Y. Dubois, B. Galambosi, P. Liang, 和 T. B. Hashimoto (2025) 长度控制的AlpacaEval：一种消除自动评估器偏差的简单方法。外部链接：2404.04475, https://arxiv.org/abs/2404.04475
  引用于：§1，§4

- M. Holsman, Y. Huang, 和 B. Dhingra (2025) 模糊推测解码用于可调精度-运行时权衡。外部链接：2502.20704, https://arxiv.org/abs/2502.20704
  引用于：§5

- K. Huang, X. Guo, 和 M. Wang (2025) SpecDec++：通过自适应候选长度提升推测解码。外部链接：2405.19715, https://arxiv.org/abs/2405.19715
  引用于：§5

- J. H. Huggins (2014) 序列蒙特卡洛中重采样的信息论分析。博士学位论文，麻省理工学院。