MAGICDEC: BREAKING THE LATENCY-THROUGHPUT TRADEOFF FOR LONG CONTEXT GENERATION WITH SPECULATIVE DECODING

发表时间: 2024-08 · arXiv:2408.11049

原文: https://arxiv.org/abs/2408.11049

作者/机构: Ranajoy Sadhukhan1∗, Jian Chen1∗, Zhuoming Chen1, Vashisth Tiwari1, Ruihang Lai1, Jinyuan Shi2, Ian En-Hsu Yen2, Avner May3, Tianqi Chen1, Beidi Chen1
1卡内基梅隆大学 2Moffett AI 3Together AI

速读

一句话结论 本文提出了 MagicDec 框架，通过为草稿模型引入 KV Cache 压缩技术，打破了长上下文和大批量推理中投机解码的效率瓶颈，在不损失生成质量的前提下，实现了最高 2.51 倍的吞吐量与延迟双重提升。

要解决什么问题 传统的投机解码（Speculative Decoding, SD）通过小模型起草、大模型验证来降低延迟，但业界普遍认为它无法胜任高吞吐量场景。卡点在于机制上的计算资源竞争：当批量大小（batch size）增加时，算术强度提高，线性层的计算会占满 GPU 算力，导致大模型验证草稿 token 的计算成本急剧上升。如果草稿的命中率不高，大模型就会把宝贵的算力浪费在验证错误 token 上，使得投机解码效率不增反降。现有的批处理技术（如 vLLM）靠增大批量来提高吞吐量，但这会牺牲每个 token 的延迟；而量化等模型压缩手段又会掉精度。因此，如何在不牺牲输出质量的前提下，同时兼顾长序列大批量推理的低延迟和高吞吐量，成为了一个亟待解决的核心卡点。

怎么做的 核心思路是放弃压缩草稿模型的权重，转而压缩草稿模型的 KV Cache。该设计能绕开原有卡点，源于作者发现的“临界序列长度”现象。在中长序列和大批量的场景下，推理瓶颈会发生转移：KV Cache 的内存占用会超过模型参数，并随批量大小线性增长。此时，现代 GPU 加载海量 KV Cache 的内存耗时远超计算耗时，推理从受计算限制转变为受内存限制。因为验证和解码共享相同的 KV 加载成本，所以在超过临界长度后，验证步骤的额外计算开销就被掩盖了，投机解码重新变得高效。为了最大化该优势，MagicDec 的关键设计由三个部件构成：草稿模型大小、草稿 KV 预算、以及 KV 压缩算法。首先，在长上下文中，使用目标模型自身配合压缩 KV 进行“自投机”，比使用独立的小型草稿模型能获得显著更高的 token 接受率（轻松超过 90%）。其次，系统需要动态分配草稿 KV 预算，以在草稿生成成本和接受率之间找到平衡。最后是 KV 选择策略，相比于每次生成都要重新搜索的动态算法（如 PQCache），预先收集稀疏 KV 的静态算法（如 SnapKV）消除了搜索开销，综合收益更高。方法的加速比由以下定义性公式决定。假设目标模型和草稿模型解码一个 token 的时间分别为 $T_T$ 和 $T_D$，验证 $\gamma$ 个 token 的时间为 $T_V(\gamma)$，预期生成长度为 $\Omega(\gamma,\alpha)$（$\alpha$ 为接受率），则投机解码的加速比为：$$ \text{Speedup} = \frac{T_T \cdot \Omega(\gamma,\alpha)}{T_D \cdot \gamma + T_V(\gamma)} $$ 当引入 KV 压缩策略时，草稿模型的总耗时 $T_D \cdot \gamma$ 会被替换为 $T_{select}(K) + T_D(K) \cdot \gamma$，其中 $T_{select}(K)$ 是保留 $K$ 个 token 的 KV 选择成本，$T_D(K)$ 是基于这 $K$ 个 token 的解码时间。通过优化该目标函数，MagicDec 确保了极低的草稿延迟和极高的接受率。

效果如何 实验在 8x A100、8x H100 和 L40 GPU 上进行，主要评估模型为 LLaMA-3.1-8B，并在 LLaMA-3.1-70B、Qwen2.5 等模型上进行了泛化测试。测试任务包括 PG-19 数据集生成及 Ruler 长文本任务。对比基线主要包括两条路线：一是标准的自回归解码（代表无投机路线），二是使用 TinyLlama-1.1B 作为草稿模型的标准投机解码（代表压缩模型权重的传统路线）。量化结果显示，在序列长度超过 4000 的大批量设置下，基于 SnapKV 压缩的 MagicDec 展现出显著优势，在 LLaMA-3.1-8B 上最高实现了 2.51 倍的加速比。在 Qwen2.5-32B 和 Mistral-7B-v0.3 上也分别测得了 1.51 倍和 2.06 倍的加速。实验证明，当序列极长且批量极大时，自投机反超了小型草稿模型，因为此时极高的接受率带来的收益远大于加载大模型权重的成本。该方法的局限性在于，它仅针对解码阶段的性能优化，并未解决长上下文在预填充阶段的计算挑战。此外，由于该方法依赖于掩盖计算开销，它在具有更高算力与内存带宽比的高端 GPU（如 H100）上能取得比普通硬件更好的加速效果。

A1 主要贡献

本文旨在解决在长上下文应用（如交互式聊天机器人、文档分析和智能体工作流）中，同时实现大型语言模型（LLM）服务的低延迟和高吞吐量这一挑战。

核心问题：
传统的推理优化技术面临一个固有的权衡：
1. 投机解码（Speculative Decoding, SD）：通过使用一个较小的草稿模型预测多个 token，再由目标模型验证，可以无损地降低延迟。然而，现有观点认为，当批量大小（batch size）增加时，验证成本会急剧上升，导致 SD 效率低下，因此不适用于高吞吐量场景。
2. 批处理（Batching）技术：如 vLLM 等通过增大批量大小来提高吞吐量，但这通常以牺牲每个 token 的延迟为代价。
3. 模型压缩技术：如量化、剪枝等可以同时改善延迟和吞吐量，但通常会牺牲模型输出的质量。

基于这些挑战，本文提出了一个核心问题：我们能否在不牺牲准确性的前提下，同时提高长序列推理的吞吐量和延迟？

研究目标与创新点：
本文对上述问题给出了肯定的回答，并指出对于中长序列的大批量推理场景，投机解码（SD）可以被有效地用来同时改善吞吐量和延迟。这一结论基于以下几个关键洞察：

1. KV Cache 成为大批量长上下文场景下的主要瓶颈：在长上下文和大批量推理中，KV Cache 的内存占用会超过模型参数，并随批量大小线性增长。尽管计算量也随之增加，但现代 GPU 极高的峰值 FLOPS 与内存带宽比导致 KV Cache 的加载时间增长远快于计算时间，使 LLM 推理变得更加受内存带宽限制（memory-bound）。
   

   

2. SD 的有效性取决于一个“临界序列长度”：现有研究认为 SD 在大批量下效率低，但这只适用于短序列。当序列长度超过一个“临界长度”后，即使在非常大的批量下，KV Cache 的加载成本也会成为主导因素。此时，验证步骤的计算开销相对于 KV 加载成本变得不那么重要，使得 SD 再次变得高效。

3. 压缩 KV Cache 是更有效的投机策略：为了在大批量处理中最小化昂贵的验证步骤，高 token 接受率至关重要。研究发现，相比压缩模型权重（使用更小的草稿模型），压缩草稿模型的 KV Cache 能在相似的内存约束下实现显著更高的接受率。如图 1c 所示，仅压缩模型权重难以达到 90% 的接受率，而 KV 压缩则可以轻松超过这一水平。

本文贡献总结：
基于以上洞察，本文提出了 MagicDec 框架，证明了与普遍认知相反，通过利用 KV 压缩，投机解码（SD）即使在大批量场景下也能实现显著的加速。

• 理论分析：本文在第 3.1 节和第 3.2 节中，理论分析了影响 SD 效率的因素，并揭示了 LLM 推理瓶颈如何随批量大小和序列长度变化，从而为 SD 在大批量场景下的有效性提供了理论基础。
• 方法论：本文在第 3.3 节和第 4.4 节中，论证了为何压缩 KV Cache 是批处理的关键，并比较了不同的静态和动态 KV 稀疏化算法，探讨了草稿成本和接受率之间的权衡。
• 综合评估：本文在第 5 节中，通过在不同硬件和任务上的全面实验，验证了理论分析的有效性。结果表明，对于 LLaMA-3.1-8B 模型，在大批量设置下最高可实现 2.51倍 的加速，显著提升了吞吐量和延迟，且不影响生成质量。

A3 背景知识/关键Observation/设计原则

本节介绍了对投机解码和 LLM 推理性能的理论分析。首先回顾投机解码加速的数学公式，并确定影响它的关键因素。其次，分析长上下文场景中的 LLM 推理，重点说明了使投机解码能够在大批量下实现加速的瓶颈转移。最后，论证了在长上下文、大批量的场景下，基于压缩 KV 的草稿模型是实现高加速比的必要条件。

3.1 投机解码加速分析

• 基本定义与公式。目标模型和草稿模型在批量大小为 B、序列长度为 S 的情况下，解码一个 token 所需的时间分别为 $T_T(B,S)$ 和 $T_D(B,S)$。目标模型验证 $\gamma$ 个 tokens 的时间为 $T_V(B,S,\gamma)$。给定草稿 token 接受率为 $\alpha \in [0,1]$，推测长度为 $\gamma$，则一次验证步骤中预期生成的 token 数为 $\Omega(\gamma,\alpha)$。根据【19, Fast inference from transformers via speculative decoding, 2022, arXiv】中的描述，预期生成 token 数可以估计为：
  

  

• 总时间与平均延迟。投机解码的总时间 $T_{SD_{Total}}$ 由以下公式给出：
  

  

  
  因此，投机解码的每 token 预期延迟为 $T_{SD_{Avg}} = T_{SD_{Total}} / \Omega(\gamma,\alpha)$。为简洁起见，除非特别说明，否则 S 是隐含的。
• 加速比及其影响因素。投机解码的加速比可以通过以下方程来理解，它揭示了调控加速比的因素：
  

  

  
  从方程2可以看出，加速比主要取决于三个因素：（a）目标模型验证与解码的成本比 $T_V(\gamma)/T_T$，（b）草稿模型与目标模型的成本比 $T_D/T_T$，以及（c）预期生成长度 $\Omega(\gamma,\alpha)$。为了获得更好的加速比，我们的目标是实现较低的 $T_V(\gamma)/T_T$（接近1）、较低的 $T_D/T_T$（接近0）和较高的 $\Omega(\gamma,\alpha)$。
  

  

3.2 KV Cache 瓶颈使投机解码加速成为可能

本节分析了随着序列长度和批量大小的增加，推理瓶颈如何转移，以及这种转移如何影响第3.1节中讨论的因素。

• 短序列场景下的挑战。对于短序列长度，投机解码对批量推理效率产生负面影响【23, Optimizing speculative decoding for serving large language models using goodput, 2024, arXiv】【33, The synergy of speculative decoding and batching in serving large language models, 2023, arXiv】。随着批量大小的增长，由于算术强度提高，线性层变得受计算限制（compute-bound）。这减少了投机解码可用于并行验证的计算资源，从而实质上增加了验证与解码的成本比。

• 中长序列场景下的转变。相比之下，对于中等到长序列，我们观察到推理过程向内存限制（memory-bound）状态转变，因为随着批量大小的增加，加载 KV Cache 的内存成本成为主导因素。这种从计算限制到内存限制的转变，使得验证成本与目标解码成本相当。由于验证和解码共享相同的 KV 预算，它们的 KV Cache 加载成本是相同的。现代 GPU 的峰值 FLOPS 与内存带宽之比很高，这导致随着批量大小的增加，KV 加载时间的增幅超过了计算时间的增幅（见图 1a）。因此，尽管受计算限制的线性层增加了验证成本，但 KV 瓶颈缓解了这一影响。

• 临界序列长度的概念。基于这种瓶颈转移，我们确定了一个临界序列长度 $S_{inflection}$，超过这个长度，投机解码即使在大批量下也能实现加速。此外，其加速比往往会随着批量大小的增加而增加。这个阈值取决于模型架构、硬件配置和草稿策略等因素。

  • 当 $S < S_{inflection}$ 时：
    在这个范围内，随着批量大小的增加，解码变得更加受计算限制。大批量会饱和可用计算资源，使得验证相对更昂贵，如图 2b 所示。对于长度为1000的 token 序列，$T_V(\gamma)/T_T$ 的成本比显著增加。如果草稿 token 的接受率低，目标模型会花费大量时间验证错误的推测，从而降低了 SD 的效率。我们在此范围内的理论估计与【23, Optimizing speculative decoding for serving large language models using goodput, 2024a, arXiv】一致。对于上下文长度低于临界序列长度的情况，投机解码的预期加速比会随着批量大小的增加而降低。

  • 当 $S \ge S_{inflection}$ 时：
    在这个范围内，投机解码可以为大批量提供加速，并且当我们使用一些智能的草稿策略时，这种加速比甚至会随着批量大小的增加而增加。这是验证与解码成本比（$T_V(\gamma)/T_T$）和草稿与目标成本比（$T_D/T_T$）随批量增加而共同演变的结果，如图 2b 和 2a 所示。

• 长序列下的成本分析。对于长序列，KV Cache 加载成为主要瓶颈，而不是计算【34, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, 2024, arXiv】【5, Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale, 2022, arXiv】，目标模型转向内存限制状态，如图 3c 所示。由于 KV 内存瓶颈随批量大小扩展，这种转变即使在大批量下也能维持。由于验证和解码阶段共享相同的 KV 加载成本，成本比 $T_V(\gamma)/T_T$ 保持接近1。

• 草稿模型成本比的作用。然而，成本比 $T_V(\gamma)/T_T$ 仍然随批量大小单调增加，这无法解释我们如何能为更大的批量实现更高的加速比。草稿与目标成本比（$T_D/T_T$）在这里扮演了重要角色。如果草稿模型的 KV Cache 大小增长慢于目标模型，那么对于更大的批量，$T_D/T_T$ 的成本比将会降低。这是因为目标模型的推理将更多地由 KV Cache 瓶颈主导，而不是草稿模型。

• 理论加速比的趋势。如图 2c 所示，在 LLaMA-3.1-8B 的情况下，投机解码的理论加速比预计会随着批量大小的增加而提高（对于较长的序列长度）。当 $S < 4000$ 时，加速比随批量大小减小；但当 $S \ge 4000$ 时，加速比随批量大小增加。

• 临界长度的影响因素。如图 3c 所示，临界序列长度 $S_{inflection}$ 取决于模型的 FLOPS 与内存比率以及 GPU 的 FLOPS 与内存带宽比率。对于具有更高 FLOPS 与内存带宽比率的设备，我们期望 $S_{inflection}$ 更低。模型也会影响这个临界序列长度。例如，像 LLaMA-3.1-8B 这样的 GQA 模型由于使用了分组查询注意力（GQA），往往具有更高的 $S_{inflection}$，因为它需要更长的序列长度才能达到相同的 KV 内存占用。
  

  

3.3 压缩 KV Cache 实现更高效的投机

本节解释了为什么在长上下文、大批量场景下，KV 压缩比轻量级草稿模型更受青睐。主要有两个原因：

• KV Cache 内存占用超过参数内存。与参数内存不同，KV Cache 的大小随批量大小线性增长。如果我们使用 LLaMA-3.1-8B 作为 LLaMA-3.1-70B 的草稿模型，LLaMA-2-7B 作为 LLaMA-2-70B 的草稿模型，草稿模型可能会占用目标模型内存的 38% 到 140%（如图 4a 和 4b 所示），因为它们的 $dim_{kv}/dim_{model}$ 比率更高。因此，在这种情况下，小型草稿模型是不够的，基于压缩 KV 的草稿策略非常有益【34, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, 2024a, arXiv】。这可以从图 3a 中看出，该图说明了对于批量大小为256的 LLaMA-3.1-8B，使用固定 KV 大小的草稿进行自投机时，$T_D/T_T$ 如何随着序列长度的增加而趋近于0。
  

  

• KV 压缩比模型压缩能实现更高的 token 接受率。在服务大批量请求时，高的草稿 token 接受率对于限制昂贵的验证步骤次数至关重要。有趣的是，我们发现 KV Cache 压缩是提高草稿 token 接受率的一种更具成本效益的方式，尤其是在大批量长上下文场景中。图 1c 展示了这一现象：如果一个目标 LLM 使用其自身 KV Cache 的稀疏化版本进行自我推测，那么它能达到的接受率要高于使用完整 KV Cache 的小型草稿模型。

• 总结。总之，一个带有压缩 KV Cache 的草稿模型在长上下文场景中实现了更高加速比的两个重要因素：低草稿成本和高接受率。图 7b 和 7c 通过实验证明了这种草稿策略相比使用小型草稿模型的标准 SD 在实现更高加速比方面的有效性。

A2 方法细节

本节介绍了 MagicDec 为确定正确草稿策略而执行的权衡分析。在第3.3节中，我们已经阐述了在这种情况下采用基于压缩 KV 的草稿策略的原因。然而，要有效利用 KV 压缩，我们需要考虑三个不同因素：（a）草稿模型大小，（b）草稿 KV Cache 大小或草稿 KV 预算，以及（c）KV 压缩算法。所有这三个因素都必须被考虑，以在草稿成本和接受率之间达到完美的平衡。

4.1 基于压缩 KV 草稿的加速通用公式

• 通用公式。首先，我们给出了使用基于压缩 KV 的草稿策略所获得的加速的通用公式。以下分析考虑了稀疏 KV 选择算法；然而，它可以很容易地扩展到其他 KV 压缩方法（例如【16, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】、【24, KIVI: A Tuning-Free Asymmetric 2Bit Quantization for KV Cache, 2024, arXiv】、【32, LoKi: Low-rank Keys for Efficient Sparse Attention, 2024, arXiv】）。稀疏 KV 方法的草稿成本取决于两个主要部分：（1）草稿模型解码成本，和（2）KV 选择成本。对于一个给定的 KV 稀疏化策略（select），固定 KV 预算为 K，选择成本表示为 $T_{select}(B,S,K)$，而 K 个 token 的解码时间为 $T_D(B,K)$。使用此 KV 策略和 KV Cache 预算 K 的草稿总时间为：
  

  

• 最终目标函数。将此作为方程2中的总草稿解码时间，我们的最终目标变为：
  

  

• 影响因素。现在我们详细讨论决定总草稿解码时间 $T_{D,select_K}$ 和最终加速比的三个主要因素。
  

  

4.2 草稿模型大小选择

• 模型大小与性能的权衡。即使使用了压缩的 KV Cache，草稿模型的权重在决定最佳性能方面仍然扮演着角色。当 KV Cache 尺寸较小时，草稿模型参数的加载是草稿成本的主要部分。通常在较低的批量大小时，一个带有压缩 KV Cache 的小型草稿模型由于其较低的草稿与目标成本比，其性能可以超过自投机。当批量大小和序列长度相对较小时，参数加载成本会妨碍草稿性能。此外，对于较小的批量，可以放宽对 token 接受率的要求，以支持一个更高效的草稿模型。然而，当批量大小超过某个阈值时，自投机因其更高的接受率而变得更有效，如图 7c 所示。

4.3 草稿 KV 预算选择

• KV 预算的动态优化。对于一个固定的草稿模型和 KV 压缩算法，最优的草稿 KV Cache 大小会因不同的批量大小和上下文长度而异。因此，在选择最优的 KV 压缩算法之前，我们需要找到候选算法各自的最优 KV 预算。我们在图 5c 中展示了为自投机优化静态 KV 选择算法的 KV 预算的重要性。不同序列长度和批量大小的批次需要不同的最低接受率才能通过投机解码实现任何加速。同样，不同的 KV 预算和不同的草稿模型会有不同的草稿成本-接受率权衡。这张图推荐了能够达到所需最低接受率的合格草稿 KV 预算。这种权衡分析对于服务具有不同序列长度的异构批次特别有用。同一批次中的不同序列可以利用不同的草稿 KV Cache 大小来实现所需的加速。

4.4 KV 选择策略的比较研究

• Top-K 注意力的局限性。最后，MagicDec 必须在不同类型的 KV 选择算法中进行选择，以调节搜索成本 $T_{select}$。尽管 top-k 注意力可以用小得多的 KV Cache 预算实现非常高的接受率，但由于其高昂的 KV 选择成本，它不是一个实用的草稿选项。

• KV 选择算法的分类与权衡。top-k 注意力有许多潜在的替代方案，但确定最优方案并非易事。主要有两种 KV 选择算法：（a）动态 KV 选择算法，如【36, Quest: Query-aware sparsity for efficient long-context llm inference, 2024, arXiv】、【46, PQCache: Product Quantization-based KVCache for Long Context LLM Inference, 2024, arXiv】；（b）静态 KV 选择算法，如【42, Efficient streaming language models with attention sinks, 2024, arXiv】、【43, PyramidInfer: Pyramid KV Cache Compression for High-Throughput LLM Inference, 2024, arXiv】、【21, SnapKV: LLM Knows What You Are Looking for Before Generation, 2024, arXiv】。第一类算法为每个输入查询动态搜索 KV Cache，试图找到 k 个最近邻。虽然这些方法可以实现更高的接受率，但它们会产生大量的搜索成本。相反，静态 KV 选择方法在生成过程中预先收集一个稀疏的 KV Cache 用于注意力近似。这种方法消除了搜索开销，但通常导致接受率较低。

• 静态与动态算法的实证比较。我们使用我们的理论框架和 LLaMA-3.1-8B 模型在各种 Ruler 任务【17, Ruler: What’s the real context size of your long-context language models?, 2024, arXiv】上的自投机解码实证接受率，来评估最先进的 KV 选择策略。我们的分析包括静态（例如 StreamingLLM【42, Efficient streaming language models with attention sinks, 2024, arXiv】、SnapKV【21, SnapKV: LLM Knows What You Are Looking for Before Generation, 2024, arXiv】）和动态（例如 PQCache【46, PQCache: Product Quantization-based KVCache for Long Context LLM Inference, 2024, arXiv】、TopK）的 KV 选择算法，探索不同的 KV 预算和推测长度，以估计最优的理论加速比。

• 任务相关的性能差异。图 5 展示了两种代表性的 KV 稀疏化算法 SnapKV 和 PQCache 之间的权衡，以及它们在三个不同的 Ruler 任务上的各自理论加速比：带密钥的大海捞针3（niah-multikeys-3）、常用词提取（cwe）和问答1（qa-1）。SnapKV 是一种静态算法，其草稿与目标成本比低于 PQCache，因为 PQCache 会产生一个与批量大小相关的 KV 选择成本 $T_{select}$。

• 接受率与搜索成本的博弈。当静态和动态方法的接受率相似时，静态方法往往占主导地位，如在 cwe 和 qa-1 任务中所见。然而，对于 niah-multikeys-3 任务，PQCache 因其更高的接受率而显著受益。当接受率接近1时，PQCache 可以利用更长的推测长度，这显著降低了方程4中的目标函数。然而，随着批量大小的增加，KV 搜索成本再次占据主导，静态算法开始优于动态算法。

A4 实验环境

• 数据集

  • PG-19: 用于主要的端到端加速评估，由【31, Compressive transformers for long-range sequence modelling, 2019, arXiv】提出。
  • Ruler: 用于不同 KV 选择策略的比较，由【17, Ruler: What’s the real context size of your long-context language models?, 2024, arXiv】提出，具体任务包括 niah-multikeys-3、cwe 和 qa-1。

• 模型架构

  • 主要模型: LLaMA-3.1-8B
  • 其他测试模型: LLaMA-3.1-70B, LLaMA-2-7B-32K, Qwen2.5-7B, Qwen2.5-32B, Mistral-7B-v0.3
  • 草稿模型: TinyLlama-1.1B (用于与自投机对比)。

• 硬件配置

  • GPU: 实验在多种 GPU 上进行，包括 8x NVIDIA A100, 8x NVIDIA H100, 以及 NVIDIA L40。
  • 连接: 所有多 GPU 实验均采用 8 路张量并行（8-way Tensor Parallelism）。

• 软件配置

  • 代码实现: 实验基于两个推理框架：
    1. 一个自研的推理后端，基于 GPT-Fast【29, Gpt-fast, 2023, GitHub】构建。
    2. MLC-LLM【37, MLC-LLM, 2023, GitHub】。
  • 依赖与优化: 自研后端使用了多项优化技术以最小化开销并提升加速效果：
    • Flashinfer: 用于加速注意力计算。
    • torch.compile: 用于编译模型。
    • Triton: 使用基于 Triton 的矩阵乘法加速 MLP 层。
    • Pytorch CUDA graphs: 用于减少 CPU 内核启动开销。
    • 张量并行: 为嵌入层实现了张量并行以进一步加速草稿阶段。
  • KV 压缩算法: 实验中评估了多种 KV 压缩算法，包括 StreamingLLM【42, Efficient streaming language models with attention sinks, 2024, arXiv】、SnapKV【21, SnapKV: LLM Knows What You Are Looking for Before Generation, 2024, arXiv】 和 PQCache【46, PQCache: Product Quantization-based KVCache for Long Context LLM Inference, 2024, arXiv】。

A4 实验结果

5.1 端到端加速

• 实验内容：本节旨在验证第3节的理论分析，即投机解码能够为中长序列同时提升吞吐量和延迟。实验使用 LLaMA-3.1-8B 模型，在 PG-19 数据集上，采用基于 StreamingLLM 风格的稀疏 KV 进行自投机，评估了不同批量大小和序列长度下的加速比。
• 实验结果：
  1. SD 对中长上下文有效：如图 6 所示，除了在“大批量+短序列”的组合下，投机解码在所有测试场景中均优于自回归解码，这验证了第 3.2 节的分析。
  2. 大批量下加速比更高：在 8x A100 GPU 上，当序列长度超过 4000 时，投机解码的加速比会随着批量大小的增加而增加，这与理论分析一致。
  3. 硬件影响：如表 1 所示，H100 在相同设置下比 A100 和 L40 取得了更高的加速比，因为 H100 具有更高的 FLOPS-to-memory 带宽比，降低了验证成本。
  4. 模型架构影响：同样在表 1 中，非 GQA 模型 LLaMA-2-7B-32K 在 8k 序列长度上比 GQA 模型 LLaMA-3.1-8B 在 32k 序列长度上获得了更高的加速比，因为非 GQA 模型具有更低的 FLOPS-to-memory 比率。
     

     

     
     

     

5.2 比较不同的 KV 压缩方法

• 实验内容：本节比较了两种静态 KV 压缩方法用于草稿生成：StreamingLLM 和 SnapKV。实验对每种方法的最优推测长度（γ）和 KV 预算进行了扫描。
• 实验结果：
  1. SnapKV 优于 StreamingLLM：如图 7b 和 7c 及表 6 所示，基于 SnapKV 的草稿策略在所有情况下均优于 StreamingLLM。
  2. 接受率是关键因素：根据图 4c 和第 4 节的分析，关键在于接受率。虽然两种方法都没有 KV 搜索开销，但 SnapKV 的接受率要高得多，并且随 KV 预算的增加而迅速提升，这弥补了其草稿成本的增加。相比之下，StreamingLLM 的接受率增长缓慢且上限较低。因此，SnapKV 凭借接受率和草稿成本的综合效应实现了更高的加速比。
  3. 高吞吐量潜力：如表 2 所示，基于 SnapKV 的自投机解码在不同任务上最高可实现 2.51x 的加速，证明了投机解码提升吞吐量的能力。
     

     

5.3 消融研究

• 实验内容：本节对影响投机解码加速比的几个关键因素进行了消融研究。
• 实验结果：
  1. 草稿 KV 预算：如图 6 所示，当批量大小和序列长度很大时，更大的 KV 预算能带来更高的加速比。这是因为此时 LLM 高度受内存限制，验证成本的绝对值远大于固定 KV 预算下的草稿成本。因此，选择具有更高接受率的更大 KV 预算，以增加每步的平均生成长度是更优的选择。
  2. 草稿模型权重：如图 7 所示，不同草稿模型的选择对加速比有显著影响。
     • 在序列长度不足够长、批量不太大的情况下（图 7b），使用小型草稿模型（TinyLlama）结合 KV 压缩的性能优于自投机。因为此时 KV 尚未完全主导推理，自投机中加载模型权重的成本较高。
     • 当序列长度和批量大小都非常大，KV Cache 主导推理时（图 7c），自投机反超了小型草稿模型。此时模型权重加载成本在总延迟中占比较小，而自投机带来的极高接受率使其能够达到更高的加速上限。
  3. 不同模型：为了验证 MagicDec 的通用性，实验还在 Qwen2.5-7B、Qwen2.5-32B 和 Mistral-7B-v0.3 模型上进行了测试。结果显示，投机解码在这些模型上同样有效，分别实现了高达 1.89x、1.51x 和 2.06x 的加速。加速比的趋势与之前的分析和 LLaMA-3.1-8B 的结果一致。（详见附录 A.5）。
     

     

A5 结论

结论总结
优化 LLM 推理的吞吐量和延迟极具挑战性，尤其是在长上下文、大批量场景下。本文的分析揭示，与现有误解相反，投机解码在这种场景下是有益的，并且其效用会随着批量大小的增加而增强。在寻找有效的草稿策略时，我们发现，在相同的内存预算下，KV 压缩比模型压缩更容易实现更高的接受率，这一优势在大批量和长上下文长度的场景下更为显著。基于这些洞见，我们探索了不同的 KV 压缩算法作为草稿策略，并提出了一个感知瓶颈的通用框架，以根据任务、批量大小和序列长度选择合适的草稿策略。

局限性与未来工作
* 局限性：MagicDec 仅关注长上下文 LLM 服务的解码性能，而预填充（prefill）阶段在这些场景中也同样充满挑战。此外，MagicDec 在高端 GPU 上往往能取得更好的加速效果，因为它们具有更高的 FLOPS 与内存带宽比和更大的 HBM 容量。
* 未来工作：
* 可以将 MagicDec 与专注于提升预填充性能的工作（如【2, Mnemosyne: Parallelization strategies for efficiently serving multi-million context length llm inference requests without approximations, 2024, arXiv】、【48, DistServe: Disaggregating Prefill and Decoding for Goodput-Optimized Large Language Model Serving, 2024, arXiv】）相结合，以同时改善预填充和解码性能。
* 可以探索在卸载（offloading）和分布式设置中采用投机解码，以减少通信开销，从而更好地利用普通设备的资源。

A6 附录

A.1 系统实现

• 系统设计。我们的投机解码系统设计如图 8 所示，展示了使用静态 KV 压缩方法的过程。静态压缩的 KV 在预填充（prefill）阶段生成，并用于草稿生成。我们将该投机解码系统实现在了最先进的推理框架 MLC-LLM【37, MLC-LLM, 2023, GitHub】和一个自研的推理后端上。主要结果来自我们的自研后端。我们的后端与 MLC-LLM 的比较见附录 A.3。
  

  

• 自研后端优化。自研推理后端基于 GPT-Fast【29, Gpt-fast, 2023, GitHub】构建，并使用 Flashinfer【12, Flashinfer, GitHub】加速注意力计算。我们使用 torch.compile 来编译模型，并利用基于 Triton 的矩阵乘法来加速 MLP 层。我们使用 Pytorch CUDA graphs 来减少 CPU 内核启动开销。这些优化有助于最小化开销并提高加速比。我们还为嵌入层实现了张量并行，以进一步加速草稿生成。

A.2 在 A100 上的各种批量大小和上下文长度的结果

• 原始数据展示。我们展示了在自研后端上运行投机解码时收集的原始数据点，以支持我们之前的讨论。我们对批量大小和序列长度进行了扫描，并比较了不同模型下不同草稿策略的加速比。所有这些实验都在 8 块 NVIDIA A100 GPU 上以 8 路张量并行方式运行。
  

  

  
  

  

  
  

  

A.3 与 MLC-LLM 结果的比较

• 后端性能对比。我们比较了基于 SnapKV 的自投机在 MLC-LLM 和我们的后端上的结果。由于测量方法不同，我们将它们放在两个表格中，如表 4 和表 5 所示。MLC-LLM 的验证时间包含了一步草稿解码的时间。我们的后端针对投机解码设置进行了高度优化，最小化了草稿和验证的开销，因此带来了更好的加速比。然而，加速比随批量大小增加的趋势是相同的，这与我们在第 3 节的理论分析一致。
  

  

  
  

  

A.4 SnapKV 和 StreamingLLM 的进一步结果

• 原始实验数据。我们展示了原始实验数据。我们比较了基于 StreamingLLM 的自投机、基于 SnapKV 的自投机，以及一个使用 StreamingLLM KV 缓存的小型草稿模型。
  

  

A.5 Qwen 和 Mistral 模型的结果

• 模型通用性验证。
  

  

A.6 TinyLLama1.1B-LLaMA2-7B-32K 结果

• 非 GQA 模型测试。我们还测试了非 GQA 模型 LLaMA-2-7B-32K，包括基于 StreamingLLM 的自投机和使用 StreamingLLM KV 缓存的小型草稿模型。由于非 GQA 模型的 FLOPS 与内存比率较低，它在相同设置下往往能实现比 GQA 模型更高的加速比。