Yggdrasil: Bridging Dynamic Speculation and Static Runtime for Latency-Optimal Tree-Based LLM Decoding

作者/机构: Yue Guan1,2∗, Changming Yu1,2,∗, Shihan Fang1, Weiming Hu1,2, Zaifeng Pan3, Zheng Wang3, Zihan Liu1,2, Yangjie Zhou1,2, Yufei Ding3, Minyi Guo1,2, Jingwen Leng1,2
1Shanghai Jiao Tong University 2Shanghai Qizhi Institute 3University of California, San Diego {bonboru,fichmi,fang-account,weiminghu, altair.liu,yj_zhou,guo-my,leng-jw}@sjtu.edu.cn, {zapan, zhw100, yufeiding}@ucsd.edu

A1 主要贡献

生成式大型语言模型(LLMs)【3, Language models are few-shot learners, 2020, arXiv preprint arXiv:2005.14165】在语言理解【10, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018, arXiv preprint arXiv:1810.04805】、推理【42, Commonsenseqa: A question answering challenge targeting commonsense knowledge, 2018, arXiv preprint arXiv:1811.00937】和多模态任务【11, An image is worth 16x16 words: Transformers for image recognition at scale, 2020, arXiv preprint arXiv:2010.11929】等众多领域表现出色。然而,LLM规模的增长带来了巨大的计算和内存需求【52, A survey on efficient inference for large language models, 2024, arXiv preprint arXiv:2404.14294】,导致高延迟和高昂的部署成本。为了解决这些瓶颈,研究者提出了多种优化策略【53, A Survey on Model Compression for Large Language Models, 2023, arXiv:2308.07633; 28, Vq-llm: High-performance code generation for vector quantization augmented llm inference, 2025, 2025 IEEE International Symposium on High Performance Computer Architecture (HPCA); 19, M-ant: Efficient low-bit group quantization for llms via mathematically adaptive numerical type, 2025, 2025 IEEE International Symposium on High Performance Computer Architecture (HPCA)】。其中,推测解码(speculative decoding)【22, Fast inference from transformers via speculative decoding, 2023, International Conference on Machine Learning】通过在自回归生成过程中利用并行性提供无损加速而脱颖而出。

核心问题与研究目标:
传统的解码方式一次只生成一个token,每个新token都依赖于前一个,这种严格的串行性严重未充分利用现代硬件。推测解码通过生成多个候选token并使用原始模型并行验证它们来打破这一瓶颈。然而,现有推测解码系统由于动态草稿算法与现代基于编译器的运行时的静态假设之间存在根本性不匹配,性能未能达到最优。如图1所示,推测解码在每个解码步骤动态调整树结构和操作符形状以最大化token接受率,但这与深度学习编译器所期望的静态数据流图相冲突,后者依赖固定的计算模式进行图融合、核函数调优和内存规划。此外,现有方法通常在忽略非均匀验证成本的简单假设下优化平均接受长度(AAL),而运行时系统通常将推测解码视为黑盒,未解决其显著的CPU逻辑开销。

为了克服这些限制,本文提出了Yggdrasil,一个协同设计算法和运行时执行的延迟最优推测解码系统。

创新点:
Yggdrasil基于三个核心思想:
1. 延迟感知目标 (Latency-aware objective):一个经过硬件分析的优化目标,通过选择最优的树形状(深度、宽度、验证范围)来最小化真实世界的延迟,而不是像AAL这样的代理指标(§4.1)。
2. 等生长树 (Equal-Growth Tree, EGT):一种新颖的草稿树结构,确保静态的操作符形状以兼容编译时图优化,同时保持适应解码上下文的灵活性(§4.2)。
3. 基于阶段的调度 (Stage-based scheduling):一个针对推测解码的搜索框架,通过消除条件分支和重叠依赖计算来最小化CPU-GPU协调成本(§5)。

Yggdrasil系统构建于开源深度学习编译器TorchInductor【2, Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】之上,无需修改模型架构,支持广泛的LLM。在多种硬件和模型配置下,Yggdrasil相较于包括SpecInfer、Sequoia和vLLM-Spec在内的最先进系统,实现了高达3.98倍的加速。

本文的主要贡献如下:
* 我们设计了一种等生长树草稿算法,该算法优化了一个延迟感知目标,实现了上下文感知的推测解码,同时保持了与图编译的兼容性(§4)。
* 我们引入了一个基于阶段调度的运行时执行框架,通过重新排列和调度推测解码阶段来减少开销(§5)。
* 我们在真实世界的模型和数据集上实现并评估了Yggdrasil,展示了其相对于现有基线的持续优势(§6 & §7)。


图 1: Yggdrasil 概览。


图 2: 推测解码示意图。

A3 背景知识与动机

背景知识

起草与验证 (Drafting and Verification)。在推测解码过程中,如图2所示,起草阶段涉及为未来步骤生成候选的token序列。这些序列随后进入验证阶段,由oracle模型评估其正确性。相应的模型被称为起草模型(drafter)和验证/目标模型(verifier/target model)。通过实现并行的token验证,推测解码减少了解码迭代的次数。

平均接受长度 (Average Accepted Length)。推测解码中的一个关键性能指标是平均接受长度(AAL),它衡量每次解码迭代平均接受的token数量。更高的AAL表示更高的效率,因为有更多的token被附加到输出中。

改进来源 (Source of Improvements)。推测解码有效性的核心洞见在于其利用了未被充分使用的计算资源。许多现代计算设备,如H100 GPU卡【33, Nvidia h100 tensor core gpu, 2023, https://www.nvidia.com/en-us/ data-center/h100/】,在解码推理期间是内存受限的,其计算能力未被充分利用。这种未充分利用为推测解码最大化资源效率创造了机会。此外,作为大多数LLM骨干的Transformer【43, Attention is all you need, 2017, Advances in Neural Information Processing Systems】解码器架构天生支持并行性,使其非常适合推测解码。具体来说,Transformer内部的自注意力机制使用序列级掩码来编码输入序列中的因果依赖关系【43, Attention is all you need, 2017, Advances in Neural Information Processing Systems】。如图2所示,这种设计可以在一次并行的推理执行中验证推测序列。推测性起草和高效的序列级并行验证相结合,增强了解码过程。


图 3: 起草结构比较。


表 1: 先前工作比较。

基于树的起草 (Tree-Based Drafting)。虽然推测解码在验证过程中利用了冗余的计算资源,但其生成性的起草阶段仍然未被充分利用,如图2-(b)所示。推测解码中的一种先进方法是使用基于树的起草,其中预测以树状结构分层组织。这种结构能够实现并行验证,并扩展了候选token的探索空间。虽然基于树的起草显著扩大了推测探索空间并提高了AAL指标,但其有效性在很大程度上受树的结构和构建策略的影响,这会影响系统性能和计算开销。

  • SpecInfer 【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】提出使用独立的起草模型生成多个草稿序列,并将它们归约为一个树结构。虽然这种方法扩大了探索空间,但它引入了重复的推测,导致起草开销增加。SpecInfer的另一种方法是通过在每个起草步骤中对预测分布的top-K个token进行采样来构建一个K元树。这种方法覆盖了广泛的推测可能性,并保证了高AAL。然而,由于其指数级的复杂性,对于更深的树来说变得不切实际,这限制了它们的草稿长度。
  • DISCO 【41, Spectr: Fast speculative decoding via optimal transport, 2024, Advances in Neural Information Processing Systems】引入了一种完全动态的树结构,其中树的深度和宽度是根据序列上下文动态确定的。这种动态方法在固定的验证预算下实现了高AAL,但由于其依赖动态控制流和可变的操作形状,会产生显著的运行时开销,这在§3中有进一步讨论。
  • Sequoia 【8, Sequoia: Scalable, robust, and hardware-aware speculative decoding, 2024, arXiv preprint arXiv:2402.12374】采取了不同的方法,提出了一种适应数据集的树结构,以平衡AAL和运行时开销。通过为给定的验证预算和数据集分布分析出最优的树结构,Sequoia优化了性能。然而,它对所有输入token应用相同的树结构,忽略了不同上下文中起草探索空间的可变性。

动机

尽管推测解码在加速LLM推理方面具有巨大潜力,但现有系统由于算法设计和运行时支持之间的根本差距而未能实现最佳性能。具体来说,我们强调以下两个瓶颈。

动态草稿破坏了静态编译。如§2所述,上下文感知的树状起草在不增加验证token数量的情况下提高了AAL。然而,这种动态性与当今编译器优化中根深蒂固的静态图假设直接冲突,如图4所示。例如,使用CUDA Graphs捕获Llama-2-7B可带来2.32倍的加速,但动态起草引入的控制流分支无法被固化到单个图中,从而阻碍了这一收益。操作符级别的自动调优【17, Fractal: Joint multi-level sparse pattern tuning of accuracy and performance for dnn pruning, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】也同样受到影响。为固定形状编译的内核可额外带来1.23倍的加速,但这仅在形状保持不变时才有效。简而言之,正是提高AAL的灵活性,使我们无法获得使LLM推理快速的编译时优化。


图 4: 不同运行时的基准测试。


图 5: 延迟和加速特性。

高AAL不等于高终端到端加速。大多数先前的工作都最大化AAL,并默认假设接受更多token能线性地转化为时钟时间的节省。这种等价关系仅在验证成本保持低廉时成立,如图5-(a)所示。一旦验证token的数量增加,每一步的验证延迟就会升级。在这种情况下,我们仍然可以观察到AAL的加速,但实际的每token加速曲线会变平甚至最终反转(图5-(b))。孤立地优化AAL会产生递减甚至负面的回报。任何现实的目标都必须在当前解码上下文中权衡接受增益与验证开销。


图 6: AAL、每步延迟和每token延迟的比较。

总结。这些瓶颈导致了次优的推测系统设计,如图6所示。基于树的方法如Sequoia将AAL推高,但保留了较高的每步延迟。面向编译的系统如vLLM【27, Optimizing speculative decoding for serving large language models using goodput, 2024, arXiv preprint arXiv:2406.14066】和TRT-LLM【34, Tensorrt-llm, 2023, https://github.com/NVIDIA/TensorRT-LLM】通过静态内核削减了步骤延迟,但无法利用动态推测,限制了AAL。这些都导致了次优的每token延迟,并且没有现有的框架能同时提供高AAL和低运行时延迟,暴露了对调和动态起草与编译效率的技术的迫切需求 。

A2 方法细节

4. 延迟最优的树状推测

本节介绍一种新颖的树状起草方法,以解决动态推测解码的挑战。我们首先为推测解码提出了一个延迟感知的优化目标,该目标最大化了实际的系统墙上时间加速比。然后,我们介绍了一种上下文感知的动态树状起草算法,该算法对编译时优化友好。

4.1 延迟感知的优化目标

现有AAL优化目标的局限性。如§3所述,现有系统通常最大化AAL作为系统加速的近似值。

$$Speedup_{\text{naive}} = \frac{T_{\text{generative}}}{T_{\text{speculative}}} \simeq \frac{Num_{\text{iteration}} * T_{\text{target}}}{T_{\text{target}}} = AAL$$

其中,$Num_{iteration}$ 是在朴素生成范式中执行的前向推理步骤数,以达到与推测解码相同的输出长度,这个值等于AAL。然而,这种近似忽略了(1)起草迭代的开销和(2)模型的延迟特性。

提出更精确的加速比度量。因此,我们建议采用一个更精确的加速比度量,该度量考虑了实际的执行延迟。我们从普通的推测解码设置开始解释,如下所示。

$$Speedup = \frac{T_{\text{generative}}}{T_{\text{speculative}}} = \frac{AAL * T_{\text{verifier}}(1)}{Num_{\text{draft}} * T_{\text{drafter}} + T_{\text{verifier}}(Num_{\text{draft}} + 1)}$$

$T_{verifier}(W)$ 是验证模型在并行验证W个token时的验证时间,如图5-(a)所示。而$Num_{draft}$是起草模型产生推测所需的推理迭代次数,该值大于或等于AAL-1。这里的-1项源于验证模型的奖励token。然而,这个项很难估计,因为它依赖于硬件和模型大小。先前的工作通过严格限制验证数量小于一个阈值并将其视为一个常数来简化它。这种简化忽略了用额外的验证时间生成比阈值更多token的机会,这在当前推测质量高时是有益的。

树状推测下的复杂性。随着树状推测的引入,问题变得更加复杂,因为AAL和$T_{draft}(W_{draft})$依赖于树的结构,具体如下。

$$Speedup = \frac{AAL(W_{\text{draft}}, D_{\text{draft}}, W_{\text{verify}}) * T_{\text{verifyer}}(1)}{\sum_{D_{\text{draft}}} T_{\text{drafer}}(W_{\text{draft}}) + T_{\text{verifier}}(W_{\text{verify}})}$$

其中,$W_{draft}$、$D_{draft}$、$W_{verify}$是与树相关的设置,分别代表树的宽度和深度。具体来说,树的深度决定了要启动的起草迭代次数。树的宽度决定了要并行起草的token数量。通过这个目标,我们得到了一个能反映实际动态树结构、对系统延迟更好的估计。

4.2 等生长树 (Equal-Growth Tree)

核心挑战与EGT算法。挑战在于找到最优的树结构,该结构通过适应上下文来最大化加速目标,同时对编译时优化友好。在上一节中,公式3揭示了加速不仅取决于AAL,还取决于紧密耦合的三元组$\langle W_{draft}, D_{draft}, W_{verify} \rangle$。精确搜索这个高维空间是不可行的,所以我们将问题分解为三个贪婪但协调的子决策,并引入等生长树(EGT)算法将它们联系在一起。EGT预测$D_{draft}$并启动所有草稿图,以消除每个token的控制流停顿。然后,我们选择在预测的树深度下最大化加速目标的$W_{draft}$。最后,我们修剪草稿树,得到一个大小为$W_{verify}$的子树,该子树最大化了加速目标。这些步骤共同产生了一个上下文自适应且运行时友好的树,其在AAL和吞吐量上都优于先前的动态序列方法【30, Accelerating speculative decoding using dynamic speculation length, 2024, arXiv preprint arXiv:2405.04304; 27, Optimizing speculative decoding for serving large language models using goodput, 2024, arXiv preprint arXiv:2406.14066】。


图 7: 等生长树的起草与剪枝。

草稿深度预测。我们首先预测$D_{draft}$,即草稿模型的调用次数。我们插入了一个轻量级的多头深度预测器,它包含一个两层MLP编码器和多个深度预测头。该预测器使用目标模型的最后一个token的嵌入作为输入,并输出预期的接受长度。我们为每个数据集和起草/验证模型对离线训练此预测器,训练数据通过在领域内验证语料库上进行性能分析一次性收集。与逐次迭代的方法不同,我们同时实例化所有$D_{draft}$个草稿图,从而消除了CPU分支误预测并最大化了GPU的重叠。

草稿宽度选择。然后,我们选择$W_{draft}$以在预测的树深度下最大化加速目标。这是一个贪婪的选择,通过选择最优宽度来平衡AAL和验证时间之间的权衡。虽然EGT在每个草稿步骤中精确地增长$W_{draft}$个叶子节点,但我们可以将它们附加到部分树的任何位置:我们选择那些路径预期AAL增益最大的位置,使用生成概率作为接受可能性的替代【44, Opt-tree: Speculative decoding with adaptive draft tree structure, 2024, arXiv preprint arXiv:2406.17276】。

验证宽度剪枝。起草后,我们解决一个最大价值子树问题,以满足验证数量$W_{verify}$的限制。由于公式3中的其他项在此时已确定,问题简化为在起草的推测树中找到一个能最大化整体加速目标的子树。这通过一个动态规划算法解决,该算法以自底向上的方式遍历树,计算每个节点及其子节点的加速值。这种后处理是必要的,仅因为在线的等生长是贪婪预测的,因为如果在生长过程中知道oracle的最优解,剪枝将是不必要的。

并行验证准备。最后,我们通过生成遵循树的因果依赖关系【36, Fasttree: Optimizing attention kernel and runtime for tree-structured LLM inference, 2025, Eighth Conference on Machine Learning and Systems】的注意力掩码,为并行验证准备好生成的树结构。

5. 基于阶段的调度运行时

CPU开销问题。通过EGT,我们可以得到一个运行时友好的起草算法,该算法为GPU执行启动静态计算图。然而,这并非故事的全部。推测解码的复杂执行流和动态控制逻辑导致了沉重的CPU开销,如图9-(a)所示。这意味着一些GPU操作是由CPU评估结果动态启动的,这导致了GPU墙上时间的“气泡”。例如,根据验证结果,我们在CPU上应用接受推测管理,并动态决定是否为最后一个尾部token进行起草。这使得GPU在应用接受推测阶段时处于空闲状态。为此,我们讨论了两种源于推测解码动态和多阶段特性的特定运行时优化。我们首先通过推测打破某些阶段之间的依赖关系,以提前其执行,从而获得更好的重叠效率。然后,我们利用性能分析统计数据,在当前EGT设置下搜索性能最佳的执行计划。


图 9: 基于阶段的推测解码调度运行时。

5.1 提前阶段执行 (Ahead-of-Time Stage Execution)

打破控制流依赖。在朴素的推测解码流水线中,由于控制流依赖,会出现明显的空闲时段,如图9-(a)所示。我们的EGT方法缓解了起草阶段的大部分气泡,但验证后的验证步骤仍然存在沉重的CPU开销。因此,我们将CPU管理与GPU推理重叠,以减少墙上时间。挑战在于推测解码包含似乎排除了这种重叠的数据依赖关系——例如,我们必须知道已接受的token集合,然后才能确定下一个要起草的头token,从而将执行划分为离散的阶段(图9-(c))。我们的关键观察是,通过使用实际所需token的超集来执行下游阶段,可以推测性地打破其中几个依赖关系。具体来说,我们提前执行两个阶段,使它们可以提前进行。

提前尾部起草 (Ahead-of-Time Tail Draft)。我们不再有条件地只起草最后的尾部token,而是推测性地起草整个候选序列。一旦任何叶子节点被接受,我们只需重用相应的结果,从而消除了条件分支。因为接受和尾部起草都是轻量级的,它们可以并行运行。这将尾部起草的CPU逻辑从关键路径中移除。

提前头部起草 (Ahead-of-Time Head Draft)。头部起草通常使用已确认的根节点解码单个token,在执行关键路径上引入一个小的GPU内核。我们提出在整个即将到来的序列上进行起草,在前一次迭代的奖励起草之后立即发出头部起草。释放的依赖关系使我们能够将头部起草与接受阶段重叠,进一步压缩执行时间线。

5.2 基于性能分析的执行计划搜索

权衡与搜索。虽然提前执行解锁了重叠,但它也通过处理额外的token增加了模型执行开销。这种权衡取决于EGT参数、每个阶段的相对成本以及目标硬件。这导致了一个巨大且高度依赖环境的设计空间。我们通过一个离线的、基于性能分析的搜索框架来解决这个问题,该框架在图9-(b)的块区域内探索阶段重叠策略。

搜索最佳执行计划。在图9-(c)所示的原始依赖图上,我们有提前执行打破了依赖关系,以及可以自由调度的并行阶段。因此,我们可以利用每个阶段的性能分析结果来搜索最佳执行计划,该结果估计了每个候选调度方案的端到端延迟。对于提前执行的阶段,将一个阶段前移会引入更多的起草计算,但会创造额外的重叠机会,使得最佳点不那么显而易见。我们保留了原始执行阶段和提前执行阶段的性能分析指标,并使用简单的网格搜索来找到最佳执行计划。由于依赖图定义明确,搜索空间很小,可以在编译时离线完成。我们在图9-(c)中可视化了几个代表性的计划。

6. 实现

系统架构。图8展示了Yggdrasil的架构,该架构经过精心设计,以优化现代推理系统的推测解码。该系统分为两个紧密集成的子系统:(1)编译时优化模块和(2)运行时执行模块。在编译时,用户提供一个草稿模型、一个验证模型和一个小的校准数据集。从这一点开始,工作流以完全模型无关的方式运行,无需对原始网络进行源代码级别的修改。Yggdrasil利用最先进的ML编译器来降低两个模型的级别,并采用基于性能分析的成本模型来生成针对目标硬件优化的执行计划。此外,它还训练一个轻量级的深度预测器来指导推测解码,确保高效和准确的执行。


图 8: 系统实现。

运行时核心抽象。在运行时,核心抽象是TokenTree类,它封装了EGT的语义。TokenTree与草稿和验证模型无缝交互,以生成推测性token并验证其正确性。它提供了健壮的接口,用于与草稿和验证模型交换token序列及其相应的KV缓存张量。KV缓存存储以及其他关键元数据,如树结构数组、注意力掩码和叶子节点位置,都作为TokenTree的属性得到高效管理。

实现框架。Yggdrasil系统是在PyTorch【2, Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】框架上实现的,利用TorchInductor编译器后端来实现高性能。重要的是,所提出的动态推测树生成算法和流水线运行时被设计为可以推广到广泛的推理系统【34, Tensorrt-llm, 2023, https://github.com/NVIDIA/TensorRT-LLM; 20, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】,使Yggdrasil成为现代AI工作负载的多功能且有影响力的解决方案。

A4 实验环境

  • 硬件配置:

    • GPU: NVIDIA A100 (80G), A40。
    • CPU: Intel(R) Xeon(R) E5-2620 v3 @ 2.40GHz。
  • 软件配置:

    • CUDA: CUDA-11.7。
    • 框架/编译器: PyTorch,使用TorchInductor作为编译器后端。
  • 模型架构:

    • 目标模型: Llama-2-7B, Llama-2-13B。
    • 草稿模型: Llama-68M, Llama-160M 【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】。
  • 数据集:

    • 名称与用途: C4 【39, Searching ´ for efficient transformers for language modeling, 2021, Advances in neural information processing systems】, Wikipedia 【45, Plagiarism — Wikipedia, the free encyclopedia, 2004, [Online; accessed 22-July-2004]】, CNNDaily 【18, Teaching machines to read and comprehend, 2015, Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 1】。用于评估语言建模任务的性能。
  • 评估指标: 主要使用每token延迟 (TPOT, per-token latency) 来评估解码效率。

  • 基线系统:
    • 端到端系统: vLLM 【20, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】, FlexFlow 【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】。
    • 推测解码系统/算法: Sequoia 【8, Sequoia: Scalable, robust, and hardware-aware speculative decoding, 2024, arXiv preprint arXiv:2402.12374】, SpecInfer 【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】。

A4 实验结果

LLM解码延迟基准测试

  • 实验内容: 评估Yggdrasil在LLM解码任务上的端到端延迟性能,重点关注每token延迟。
  • 实验结果: 如图10所示,Yggdrasil在所有基准测试中始终表现最佳,在A100 GPU上平均解码延迟提升了3.98倍,在A40 GPU上提升了2.76倍。
  • 分析结论: 这证明了其协同设计方法的有效性,超越了以算法为中心的Sequoia和以系统为中心的SpecInfer。SpecInfer因GPU执行和服务运行时的巨大开销表现最差。Sequoia和vLLM-Spec利用TorchInductor优化模型执行,分别比SpecInfer平均提速2.45倍和2.66倍。Yggdrasil通过结合上下文感知推测和CUDA-graph优化,将加速比进一步提升至3.37倍。性能增益在A100 GPU上更为显著,因为它存在更严重的硬件未充分利用问题,使Yggdrasil能够最大化其优势。


图 10: 相较于SpecInfer[31]的端到端每token生成延迟加速结果。

树结构比较

  • 实验内容: 评估本文提出的EGT树结构和硬件感知优化目标的有效性。
  • 树结构AAL比较:

    • 实验结果: 如图11-(a)所示,序列推测解码的AAL有限且很快饱和。Sequoia通过静态适应数据集和验证数量,实现了良好的AAL,但对所有输入使用相同的树结构,忽略了上下文。
    • 分析结论: 相比之下,EGT动态调整其宽度,在最优验证设置下始终优于Sequoia的静态树,凸显了上下文感知树结构选择的重要性。
  • 上下文感知的动态树:

    • 实验结果: 如图11-(b)所示,使用公式3的理论加速比来评估树结构有效性。序列推测解码的加速比最低。像N-way和Sequoia这样的静态树随着验证数量的增加而性能下降,这是由于资源饱和导致验证延迟增加。
    • 分析结论: 不同宽度的EGT在不同的验证设置下实现了最佳加速比。通过为每个请求动态选择最佳宽度和验证数量,Yggdrasil利用上下文和硬件特性,始终优于所有基准,实现了卓越的性能。


图 11: Wikitest数据集上不同树结构的AAL和加速比比较。

优化分解

  • 实验内容: 分析Yggdrasil系统中各项优化技术的有效性,如图12所示。
  • 实验结果与分析:
    • O1 (延迟最优树推测): 作为基线,其改进体现在AAL指标上。
    • O2 (图编译): 使用TorchInductor进行图编译,贡献了平均2.775倍的加速,是四项优化中最重要的部分。这强调了运行时支持的必要性,以及EGT设计在同时实现高AAL和编译优化中的关键作用。
    • O3 (验证宽度剪枝): 通过动态调整验证数量,比O2平均提速1.07倍。改进主要来自自适应验证策略,该策略将验证数量与推理时间曲线的饱和区域对齐,减少了开销。
    • O4 (基于图的调度优化): 比之前优化平均提升1.21倍。这表明在现有执行运行时内利用推测解码特定优化的重要性。
    • O5 (草稿深度预测器): 通过显式考虑上下文难度,为系统带来了1.10倍的加速。基线是使用固定的树结构(深度16,宽度8)。


图 12: Llama-2-7B模型与Llama-68m草稿模型的优化分解。

敏感性分析

  • EGT结构:

    • 实验内容: 对EGT参数进行敏感性分析,如图13所示。
    • 分析结论: 结果显示,草稿宽度、草稿深度和验证数量之间的相互作用显著影响每token延迟。例如,静态分析中,$D_{draft}=8$, $W_{draft}=8$, $W_{verify}=64$的组合性能最佳。
  • 加速比目标:

    • 实验内容: 在C4数据集上进行消融实验,比较优化加速比目标与优化AAL的优势,如图14所示。
    • 分析结论: 与直接优化AAL相比,将公式3作为优化目标在所有草稿-验证模型设置中额外带来了8%的性能提升。这是因为加速比目标能更好地捕捉运行时特性,包括动态的草稿和验证执行时间。
  • 温度影响:

    • 实验内容: 研究采样温度值对Sequoia和Yggdrasil的影响,如图15所示。
    • 分析结论: 当温度为0时,Sequoia和Yggdrasil都获得了更好的性能,因为低温下草稿模型能更好地与目标模型对齐,从而获得更高的AAL。在不同温度下,Yggdrasil始终优于Sequoia,平均加速比为1.49倍。


图 13: Llama-2-7B和Llama-68m的EGT参数敏感性分析。


图 14: 以加速比和AAL为优化目标的比较。


图 15: 使用Llama-7B和Llama-68M在不同采样温度下的延迟。

A7 补充细节

相关工作

推测解码变体和服务系统。除了§3中讨论的工作外,还有大量研究提出了推测解码的变体。首先,模型透明的研究在不修改目标模型结构的情况下提高了草稿的准确性或效率。Lookahead decoding【14, Break the sequential dependency of llm inference using lookahead decoding, 2024, arXiv preprint arXiv:2402.02057】、ANPD【35, Lossless acceleration of large language model via adaptive n-gram parallel decoding, 2024, arXiv preprint arXiv:2404.08698】和NEST【23, Nearest neighbor speculative decoding for llm generation and attribution, 2024, arXiv preprint arXiv:2405.19325】应用更高效的检索或统计方法来生成草稿token,而不是使用草稿模型。一些工作【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3; 41, Spectr: Fast speculative decoding via optimal transport, 2024, Advances in Neural Information Processing Systems】设计了更复杂的采样程序以获得更好的草稿准确性。这些工作与Yggdrasil是正交的。相比之下,模型侵入式方法修改目标模型,以对齐目标和草稿模型【51, Distillspec: Improving speculative decoding via knowledge distillation, 2023, arXiv preprint arXiv:2310.08461】,或使用目标模型的一部分【26, Kangaroo: Lossless self-speculative decoding via double early exiting, 2024, arXiv preprint arXiv:2404.18911; 12, Layer skip: Enabling early exit inference and self-speculative decoding, 2024, arXiv preprint arXiv:2404.16710】或额外的子模块【4, Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024, arXiv preprint arXiv:2401.10774; 25, Eagle: Speculative sampling requires rethinking feature uncertainty, 2024, arXiv preprint arXiv:2401.15077】来产生推测,以减少草稿开销。这些工作需要对原始模型进行显式修改,超出了本文的支持范围。此外,除了在§7中进行基准测试的服务系统【31, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3; 8, Sequoia: Scalable, robust, and hardware-aware speculative decoding, 2024, arXiv preprint arXiv:2402.12374; 20, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】外,SmartSpec【27, Optimizing speculative decoding for serving large language models using goodput, 2024, arXiv preprint arXiv:2406.14066】和MagicDec【6, Magicdec: Breaking the latency-throughput tradeoff for long context generation with speculative decoding, 2024, arXiv preprint arXiv:2408.11049】研究了将推测解码与连续批处理相结合,并在在线服务场景中实现服务水平目标(SLO)和吞吐量提升之间的平衡。然而,它们没有利用推测解码本身的执行效率,而这正是本文的主要关注点。

ML系统中的动态支持。许多工作研究了在机器学习系统中处理动态工作负载。像TensorFlow Eager【1, {TensorFlow}: a system for {Large-Scale} machine learning, 2016, 12th USENIX symposium on operating systems design and implementation (OSDI 16)】和PyTorch Eager【37, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】这样的框架支持在运行时动态构建图【16, Amanda: Unified instrumentation framework for deep neural networks, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1】,但硬件利用率不佳。ML编译器,包括PyTorch 2.0【2, Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】、DietCode【49, Dietcode: Automatic optimization for dynamic tensor programs, 2022, Proceedings of Machine Learning and Systems】、Nimble【21, Nimble: Lightweight and parallel gpu task scheduling for deep learning, 2020, Advances in Neural Information Processing Systems】和BladeDISC【50, Bladedisc: Optimizing dynamic shape machine learning workloads via compiler approach, 2023, Proceedings of the ACM on Management of Data】,通过自适应图转换和代码生成来优化具有可变输入形状的模型。对于具有数据依赖执行路径的模型,Cocktailer【47, Cocktailer: Analyzing and optimizing dynamic control flow in deep learning, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】、Grape【48, Grape: Generalizing robot policy via preference alignment, 2024, arXiv preprint arXiv:2411.19309】、Cortex【13, Cortex: Production infrastructure for machine learning at scale, 2022, https://github.com/ cortexlabs/cortex】、BrainStorm【24, Think outside the code: Brainstorming boosts large language models in code generation, 2023, arXiv preprint arXiv:2305.10679】和DyCL【7, Dycl: Dynamic neural network compilation via program rewriting and graph optimization, 2023, Proceedings of the 32nd ACM SIGSOFT International Symposium on Software Testing and Analysis】优化动态控制流以减少运行时开销。此外,在动态服务场景中,DyNet【32, Dynet: The dynamic neural network toolkit, 2017, arXiv preprint arXiv:1701.03980】、BatchMaker【15, Low latency rnn inference with cellular batching, 2018, Proceedings of the Thirteenth EuroSys Conference】、DVABatch【9, {DVABatch}: Diversity-aware {Multi-Entry}{Multi-Exit} batching for efficient processing of {DNN} services on {GPUs}, 2022, 2022 USENIX Annual Technical Conference (USENIX ATC 22)】、Brainstorm【24, Think outside the code: Brainstorming boosts large language models in code generation, 2023, arXiv preprint arXiv:2305.10679】、ORCA【46, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】和Llumnix【40, Llumnix: Dynamic scheduling for large language model serving, 2024, arXiv preprint arXiv:2406.03243】通过动态调整批次大小或优化调度来提高GPU利用率和吞吐量。尽管有上述进展,这些工作在处理推测解码的额外开销时仍然面临限制,其中不可预测的批次大小和控制流增加了显著的复杂性。

局限性

核心假设与适用场景。Yggdrasil的核心结果是在一个延迟最优的环境下得出的,其中单个交互式请求独占所有可用的GPU内存和计算资源。这种配置反映了边缘或本地部署的场景,其中专用加速器服务于单个用户或延迟关键的流水线。

对批处理服务的不适用性。虽然上述假设使得能够清晰地分析内存带宽的权衡并激发了EGT起草方法,但它不适用于面向吞吐量的批处理服务。生产推理堆栈通常将数十到数百个提示进行批处理,以最大化加速器利用率。在这种情况下,将推测解码与批处理调度程序协调起来对整体效率至关重要。

未来研究方向。设计一个统一的策略,共同决定何时进行推测以及如何打包请求是至关重要的。简而言之,虽然Yggdrasil在单请求场景中推动了延迟的极限,但将该框架扩展到延迟-吞吐量联合优化空间,其中推测解码和批处理调度被一同解决,仍然是一个开放的研究方向。

A5 结论

本文我们提出了Yggdrasil,一个新颖的延迟最优推测解码框架,它能够在保持编译时优化的同时实现上下文感知的推测。此外,它通过基于性能分析驱动的运行时洞察,自适应地选择最优的草稿树,从而动态地最小化执行延迟。为了进一步提高效率,Yggdrasil采用了一种基于阶段的执行调度策略,简化了推测解码的工作流程。实验评估表明,Yggdrasil实现了显著的性能提升,相较于最先进的系统,提供了卓越的加速效果。

A6 附录

NeurIPS 论文核对清单

  1. 声明: 摘要和引言中的主要声明是否准确反映了论文的贡献和范围? [是] 我们在引言部分明确强调了贡献并进行了详细讨论,确保与摘要和引言中的声明一致。
  2. 局限性: 论文是否讨论了作者工作的局限性? [是] 我们在附录中包含了一个讨论延迟导向优化局限性的部分。
  3. 理论假设与证明: 对于每个理论结果,论文是否提供了完整的假设集和完整(且正确)的证明? [不适用] 论文不包含理论结果。
  4. 实验结果可复现性: 论文是否充分披露了复现主要实验结果所需的所有信息,以影响论文的主要主张和/或结论(无论是否提供代码和数据)? [是] 我们在评估设置部分讨论了复现结果的实验设置。
  5. 数据和代码的开放获取: 论文是否提供数据和代码的开放获取,并附有足够的说明以忠实地复现主要实验结果,如补充材料中所述? [否] 一旦被接受,我们将开源我们的代码。
  6. 实验设置/细节: 论文是否详细说明了所有理解结果所必需的训练和测试细节(例如,数据划分、超参数、选择方法、优化器类型等)? [是] 我们在论文的实验设置部分讨论了训练细节。
  7. 实验统计显著性: 论文是否报告了适当且正确定定义的误差棒或关于实验统计显著性的其他适当信息? [否] 评估结果是在评估集上对不同上下文相关样本进行静态平均得出的。
  8. 实验计算资源: 对于每个实验,论文是否提供了复现实验所需的足够计算机资源信息(计算单元类型、内存、执行时间)? [是] 我们在实验设置部分指明了细节。
  9. 道德准则: 论文中进行的研究是否在各方面都符合NeurIPS道德准则 https://neurips.cc/public/EthicsGuidelines[是] 该研究在各方面均符合NeurIPS道德准则。
  10. 更广泛的影响: 论文是否讨论了所进行工作的潜在积极和消极社会影响? [不适用] 所进行的工作没有社会影响。
  11. 保障措施: 论文是否描述了为负责任地发布具有高滥用风险的数据或模型(例如,预训练语言模型、图像生成器或抓取的数据集)而设置的保障措施? [不适用] 论文不构成此类风险。
  12. 现有资产的许可证: 论文中使用的资产(例如,代码、数据、模型)的创建者或原始所有者是否得到适当的致谢,并且许可证和使用条款是否被明确提及并得到适当遵守? [是] 所有资产均被适当引用。
  13. 新资产: 论文中引入的新资产是否得到良好记录,并且文档是否与资产一同提供? [不适用] 论文不发布新资产。
  14. 众包与人类受试者研究: 对于众包实验和涉及人类受试者的研究,论文是否包括给予参与者的完整说明文本和截图(如果适用),以及有关报酬的详细信息(如果有)? [不适用] 论文不涉及众包或人类受试者研究。
  15. 机构审查委员会(IRB)批准: 论文是否描述了研究参与者可能承担的潜在风险,是否向受试者披露了这些风险,以及是否获得了机构审查委员会(IRB)的批准(或根据您所在国家或机构的要求获得的同等批准/审查)? [不适用] 论文不涉及众包或人类受试者研究。
  16. LLM使用声明: 如果LLM是本研究核心方法中一个重要的、原创的或非标准组成部分,论文是否描述了其使用情况? [不适用] 本研究的核心方法开发不涉及将LLM作为任何重要的、原创的或非标准的组成部分。