作者/机构: 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
生成式大型语言模型(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: 推测解码示意图。
起草与验证 (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指标,但其有效性在很大程度上受树的结构和构建策略的影响,这会影响系统性能和计算开销。
尽管推测解码在加速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和低运行时延迟,暴露了对调和动态起草与编译效率的技术的迫切需求 。
本节介绍一种新颖的树状起草方法,以解决动态推测解码的挑战。我们首先为推测解码提出了一个延迟感知的优化目标,该目标最大化了实际的系统墙上时间加速比。然后,我们介绍了一种上下文感知的动态树状起草算法,该算法对编译时优化友好。
现有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数量。通过这个目标,我们得到了一个能反映实际动态树结构、对系统延迟更好的估计。
核心挑战与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】的注意力掩码,为并行验证准备好生成的树结构。
CPU开销问题。通过EGT,我们可以得到一个运行时友好的起草算法,该算法为GPU执行启动静态计算图。然而,这并非故事的全部。推测解码的复杂执行流和动态控制逻辑导致了沉重的CPU开销,如图9-(a)所示。这意味着一些GPU操作是由CPU评估结果动态启动的,这导致了GPU墙上时间的“气泡”。例如,根据验证结果,我们在CPU上应用接受推测管理,并动态决定是否为最后一个尾部token进行起草。这使得GPU在应用接受推测阶段时处于空闲状态。为此,我们讨论了两种源于推测解码动态和多阶段特性的特定运行时优化。我们首先通过推测打破某些阶段之间的依赖关系,以提前其执行,从而获得更好的重叠效率。然后,我们利用性能分析统计数据,在当前EGT设置下搜索性能最佳的执行计划。
图 9: 基于阶段的推测解码调度运行时。
打破控制流依赖。在朴素的推测解码流水线中,由于控制流依赖,会出现明显的空闲时段,如图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内核。我们提出在整个即将到来的序列上进行起草,在前一次迭代的奖励起草之后立即发出头部起草。释放的依赖关系使我们能够将头部起草与接受阶段重叠,进一步压缩执行时间线。
权衡与搜索。虽然提前执行解锁了重叠,但它也通过处理额外的token增加了模型执行开销。这种权衡取决于EGT参数、每个阶段的相对成本以及目标硬件。这导致了一个巨大且高度依赖环境的设计空间。我们通过一个离线的、基于性能分析的搜索框架来解决这个问题,该框架在图9-(b)的块区域内探索阶段重叠策略。
搜索最佳执行计划。在图9-(c)所示的原始依赖图上,我们有提前执行打破了依赖关系,以及可以自由调度的并行阶段。因此,我们可以利用每个阶段的性能分析结果来搜索最佳执行计划,该结果估计了每个候选调度方案的端到端延迟。对于提前执行的阶段,将一个阶段前移会引入更多的起草计算,但会创造额外的重叠机会,使得最佳点不那么显而易见。我们保留了原始执行阶段和提前执行阶段的性能分析指标,并使用简单的网格搜索来找到最佳执行计划。由于依赖图定义明确,搜索空间很小,可以在编译时离线完成。我们在图9-(c)中可视化了几个代表性的计划。
系统架构。图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工作负载的多功能且有影响力的解决方案。
硬件配置:
软件配置:
模型架构:
数据集:
评估指标: 主要使用每token延迟 (TPOT, per-token latency) 来评估解码效率。
图 10: 相较于SpecInfer[31]的端到端每token生成延迟加速结果。
树结构AAL比较:
上下文感知的动态树:
图 11: Wikitest数据集上不同树结构的AAL和加速比比较。
图 12: Llama-2-7B模型与Llama-68m草稿模型的优化分解。
EGT结构:
加速比目标:
温度影响:
图 13: Llama-2-7B和Llama-68m的EGT参数敏感性分析。
图 14: 以加速比和AAL为优化目标的比较。
图 15: 使用Llama-7B和Llama-68M在不同采样温度下的延迟。
推测解码变体和服务系统。除了§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在单请求场景中推动了延迟的极限,但将该框架扩展到延迟-吞吐量联合优化空间,其中推测解码和批处理调度被一同解决,仍然是一个开放的研究方向。
本文我们提出了Yggdrasil,一个新颖的延迟最优推测解码框架,它能够在保持编译时优化的同时实现上下文感知的推测。此外,它通过基于性能分析驱动的运行时洞察,自适应地选择最优的草稿树,从而动态地最小化执行延迟。为了进一步提高效率,Yggdrasil采用了一种基于阶段的执行调度策略,简化了推测解码的工作流程。实验评估表明,Yggdrasil实现了显著的性能提升,相较于最先进的系统,提供了卓越的加速效果。