文章标题:DistServe:解耦预填充(Prefill)和解码(Decoding)以优化大型语言模型服务的良率(Goodput)
作者/机构:Yinmin Zhong1, Shengyu Liu1, Junda Chen3, Jianbo Hu1, Yibo Zhu2, Xuanzhe Liu1, Xin Jin1, Hao Zhang3 (1北京大学计算机学院, 2StepFun, 3加州大学圣迭戈分校)
本文旨在解决现有大型语言模型(LLM)服务系统中预填充(Prefill)和解码(Decoding)两个计算阶段共存(colocation)所引发的性能问题。
核心问题:
现有LLM服务系统将预填充和解码阶段放置在相同的GPU上,并对所有用户和请求的计算进行批处理。这种策略导致了几个关键问题:
1. 预填充-解码干扰:预填充步骤通常比解码步骤耗时得多。当两者批处理在一起时,解码步骤会被预填充拖慢,显著增加每个输出token的时间(TPOT);同时,解码步骤的存在也会增加预填充生成首个token的时间(TTFT)。即使分开调度,它们也会竞争GPU资源,导致排队延迟增加。
2. 资源与并行策略耦合:共存使得预填充和解码阶段必须共享资源分配和并行策略。然而,这两个阶段具有不同的计算特性和延迟要求(预填充通常是计算密集型,解码是内存带宽密集型),需要不同的优化策略。这种耦合导致系统必须为了满足更严格的延迟要求而过度配置资源,从而降低了成本效益。
研究目标:
本文的目标是优化每个LLM查询的成本,同时满足高服务等级目标(SLO)达成率。具体而言,系统应在满足应用对TTFT和TPOT的特定延迟要求下,最大化每GPU的“良率”(Goodput),即在SLO达成目标(如90%)内每秒可服务的最大请求速率。
主要创新点(贡献):
1. 识别并提出解耦方案:识别了现有LLM服务系统中预填充-解码的干扰和资源耦合问题,并提出了将这两个阶段解耦(disaggregate),即分配到不同GPU上执行的方案。
2. 设计新颖的放置算法:设计了一种新颖的放置算法,能够根据应用的TTFT和TPOT要求、集群带宽等因素,自动为预填充和解码实例共同优化资源分配和并行策略,以找到良率最优的部署方案。
3. 构建并全面评估DistServe系统:基于解耦思想构建了一个名为DistServe的良率优化LLM服务系统。通过在真实工作负载下的综合评估,验证了该方法的有效性。评估结果表明,与现有最先进系统相比,DistServe在满足各种延迟约束下,能够服务多达7.4倍的请求或满足严格12.6倍的SLO。
图1:在单张NVIDIA 80GB A100上,使用合成工作负载(输入长度=512,输出长度=64)服务13B参数LLM时的性能。上图:现有系统与仅服务预填充阶段的系统在P90首个token时间(TTFT)延迟上的比较。下图:现有系统与仅服务解码阶段的系统在P90每个输出token时间(TPOT)延迟上的比较。
一个LLM服务遵循客户端-服务器架构:客户端提交一个文本序列作为请求到服务器;服务器在GPU上托管LLM,对请求运行推理,并将生成结果响应(或流式传输)回客户端。由于独特的预填充-解码过程,LLM服务可能对TTFT和TPOT都施加激进的服务等级目标(SLO),这随应用需求而变化。服务系统必须在满足这两个SLO的同时,最小化与昂贵GPU相关的成本。换言之,我们希望服务系统能最大化在遵守SLO达成目标的条件下,每个配置的GPU每秒服务的请求数——即最大化每GPU良率。接下来,我们将详细介绍LLM推理计算(§2.1),并讨论现有的LLM服务优化方法(§2.2)。
现代LLM的计算过程。现代LLM(如【37, GPT-4技术报告, 2023】、【51, Hugo Touvron等人, Llama: Open and efficient foundation language models, 2023】)通过给定一个输入序列来预测下一个token。这个预测过程涉及到为序列中的每个token计算一个隐藏表示。LLM可以接受可变数量的输入token并并行计算它们的隐藏表示,其计算工作负载随着并行处理的token数量超线性增长。无论输入token数量多少,计算过程都需要大量的I/O操作,以将LLM权重和中间状态从GPU的HBM移动到SRAM。这个过程在不同输入大小下是一致的。
预填充与解码阶段的计算差异。预填充步骤处理一个新的序列,通常包含许多token,并并行处理这些token。与预填充不同,每个解码步骤只处理前一步生成的一个新token。这导致了两个阶段之间显著的计算差异。当处理非简短的用户提示时,预填充步骤往往是计算密集型的。例如,对于一个13B的LLM,计算一个512个token序列的预填充会使一个A100接近计算密集状态(见§3.1)。相比之下,尽管每步只处理一个新token,解码阶段产生的I/O量与预填充阶段相似,使其受到GPU内存带宽的限制。
中间状态(KV缓存)的共享。在两个阶段中,每个token位置都会生成中间状态,即KV缓存【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】,这些缓存在后续的解码步骤中需要再次使用。为了避免重新计算,它们被保存在GPU内存中。由于LLM权重和KV缓存在内存中的共享使用,大多数LLM推理引擎选择将预填充和解码阶段放置在同一个GPU上,尽管它们的计算特性截然不同。
在实时在线服务中,多个请求同时到达,必须在SLO内得到服务。对它们的计算进行批处理和并行化是实现低延迟、高吞吐量和高GPU利用率的关键。
批处理(Batching)。当前的推理服务系统【9, Amey Agrawal等人, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, arXiv preprint arXiv:2308.16369, 2023】、【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】、【54, Gyeong-In Yu等人, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, USENIX OSDI】采用一种称为连续批处理的技术。该方法将新请求的预填充与正在进行的请求的解码批处理在一起。这提高了GPU利用率,并最大化了整个系统的吞吐量——即所有用户和请求每秒生成的token数。然而,如§1所述并在§2.3中详述,这种方法导致了TTFT和TPOT之间的权衡。连续批处理的一种高级变体【9】试图通过将长预填充分块并附加解码任务来平衡TTFT和TPOT,但这实质上是用TTFT换取TPOT,无法消除干扰(§2.3)。总之,批处理预填充和解码不可避免地导致在TTFT或TPOT上的妥协。
模型并行(Model parallelism)。在LLM服务中,模型并行通常分为算子内(intra-operator)和算子间(inter-operator)并行【33, Zhuohan Li等人, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, arXiv, 2023】、【46, Mohammad Shoeybi等人, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】、【59, Lianmin Zheng等人, Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning, 2022, USENIX OSDI】。两者都可以用来支持更大的模型,但对服务性能的影响不同。算子内并行将计算密集型算子(如矩阵乘法)划分到多个GPU上,加速了计算但导致大量通信。它减少了执行时间,从而降低了延迟,特别是预填充阶段的TTFT,但这需要GPU之间的高带宽连接(例如NVLINK)。算子间并行将LLM的层组织成多个阶段,每个阶段在一个GPU上运行以形成流水线。由于阶段间的通信,它会适度增加执行时间,但每增加一个GPU,系统的速率容量就能线性扩展。在本文中,我们揭示了模型并行的另一个好处:由于执行时间缩短,预填充和解码阶段的排队延迟都减少了。我们将在§3中进一步探讨这一点。除了模型并行,复制一个模型实例(无论其模型并行配置如何)也能线性扩展系统的速率容量。
并行策略的复杂性。这些并行策略创造了一个复杂的优化空间,需要根据应用的延迟要求进行仔细的权衡。
如现有系统那样,将预填充和解码计算放置在同一位置并进行批处理以最大化整体系统吞吐量,对服务提供商而言是成本有效的。然而,在存在SLO的情况下,由于下述问题,现有方法难以同时维持高质量服务和低成本。
预填充-解码干扰。如图2所示,将单个预填充任务添加到一个解码请求批次中,会显著减慢两个过程,导致TTFT和TPOT显著增加。具体来说,批次中的解码任务必须等待更长的预填充任务完成,从而延长了TPOT;这种减速随着预填充的加长而加剧,如图2(b)所示。将解码任务添加到预填充中也会增加完成预填充任务的时间,特别是在GPU已经满负荷时(图2蓝色曲线)。
分块预填充(Chunked-prefill)的局限性。一种试图减轻这种干扰的方法叫做分块预填充与搭载(chunked-prefill with piggyback)【3, Deepspeed model implementations for inference (mii), 2023】、【9, Amey Agrawal等人, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, arXiv preprint arXiv:2308.16369, 2023】。它提议将长预填充分割成块,并将一个预填充块与少量解码任务批处理在一起(即搭载)。这项技术减轻了长预填充任务对解码任务的减速,但并未消除它。此外,它给预填充任务带来了额外的开销,这个开销不易通过调整块大小来缓解。首先,如果块大小设置得远低于能使GPU饱和的拐点,那么预填充任务的执行时间会变长,因为它在同一批次中与解码任务竞争,无法独占GPU资源。其次,如果我们增加块大小以接近饱和GPU,搭载的机会将减少,因为留给解码token的槽位有限。同时,分块预填充导致预填充任务的内存访问显著增多。这是因为所有先前块的KV缓存必须重复地从HBM加载到SRAM来计算每个后续块。具体来说,如果一个预填充任务被分成N个等大小的块,总共需要加载 $N + (N-1) + ... + 1 = O(N^2)$ 个KV缓存块,而无分块情况下是 $O(N)$。随着上下文长度的增加,这个开销会更大。
图2:在服务一个13B LLM时,批处理执行时间随批大小增加的变化情况。比较了纯解码批处理与增加一个预填充任务的批处理。
无效的调度策略。不批处理预填充和解码任务,而是顺序调度它们,并不能减轻干扰。解码任务可能会因为等待GPU上正在进行的预填充任务而经历更长的排队延迟。此外,专门用于解码的批次常常导致GPU利用率不足。优先处理任一阶段的任务都会对另一阶段的延迟产生负面影响,使得优先级调度变得无效。
资源与并行策略的耦合。在相同的GPU上共存预填充和解码阶段,不可避免地导致它们共享资源和并行设置。然而,每个阶段都有其独特的计算特性和延迟要求,需要更多异构的资源分配。例如,预填充阶段倾向于计算密集型,并且可以从更多的算子内并行中受益,以减少执行时间,满足对TTFT的严格SLO。相比之下,解码阶段的最佳并行配置取决于运行的批处理大小。在现有系统中,由于耦合,资源分配和并行计划被调整以满足TTFT和TPOT中要求更高的一方,这可能对另一方并非理想。这通常导致为了满足两个SLO而过度配置资源。
机遇:解耦预填充与解码。为解决这些问题,我们提议解耦预填充和解码阶段。我们使用术语实例(instance)来表示管理模型权重的一个完整副本的资源单元。当应用模型并行时,一个实例可以对应多个GPU。请注意,当我们把两个阶段解耦到不同的GPU上时,每个阶段都管理自己的一份模型权重副本,从而产生预填充实例和解码实例。一个预填充实例在收到请求后,只为该请求执行预填充计算以生成第一个输出token。然后,它将中间结果(主要是KV缓存)发送给一个解码实例,该实例负责后续的解码步骤。因为解码计算的GPU利用率通常较低,我们可以为每个解码实例分配多个预填充实例。这允许批处理更多的解码任务,以实现更高的GPU利用率。
解耦的优势。解耦预填充和解码自然解决了两个阶段之间的干扰,并使每个阶段能够专注于其优化目标——TTFT或TPOT。每种类型的实例可以采用不同的资源和并行策略来满足各种延迟要求。通过调整提供给两种实例的GPU数量和并行方式,我们可以最大化整个系统的每设备良率,避免过度配置,最终转化为在遵守服务质量的前提下降低每个查询的成本。接下来,我们将研究如何为每个阶段找到最佳的资源分配和并行计划。
解耦将两个阶段分开,允许对每个阶段的特性进行独立分析,为算法设计提供了宝贵的见解。它也扩展了设计空间:现在每个阶段都需要根据其延迟要求独立地进行扩展和调度。
在本节中,我们分析了解耦后预填充(§3.1)和解码实例(§3.2)的计算模式。我们的目标是确定关键参数,并为每个阶段的批处理和并行化推导出指导原则。然后,我们重点介绍了一些实际部署的考虑因素(§3.3)。本节为每GPU良率优化奠定了基础。
解耦后,预填充阶段通过并行处理用户提示的所有token来生成第一个token。假设给定的到达率,我们的目标是使用最少的资源来满足服务对TTFT的延迟要求。
批处理策略。预填充步骤通常是计算密集型的。图3(a)显示了预填充阶段的吞吐量如何随输入长度和批处理大小变化。对于一个13B参数的LLM,处理一个512个token的单一序列可以完全占用一个A100 GPU。一旦GPU变得计算受限,向批次中添加更多请求不再能提高GPU效率。相反,它会按比例延长批次的总处理时间,无意中延迟了所有包含的请求。因此,对于预填充实例,有必要预先对特定的LLM和GPU进行性能分析,以确定一个关键的输入长度阈值,记为Lm,超过该阈值,预填充阶段就变得计算受限。只有当被调度请求的输入长度低于Lm时,才应考虑批处理更多请求。在实践中,用户提示的平均长度通常超过数百个token【8, Sharegpt teams, https://sharegpt.com/, 2023】。预填充实例的批处理大小通常保持较小。
图3:在服务一个13B参数的LLM时,两个阶段在不同批处理大小和输入长度下的吞吐量。
并行计划。为了研究纯预填充实例的并行偏好,我们在两台A100 GPU上使用算子间或算子内并行策略服务一个66B的LLM。为简化问题,我们假设请求的输入长度统一为512个token,并采用泊松到达过程。我们在图4(a)中比较了不同到达率下的平均TTFT:算子内并行在较低的到达率下更有效,而随着速率的增加,算子间并行则更具优势。解耦使得预填充阶段的功能类似于一个M/D/1队列,因此我们可以使用排队论来验证这一观察。
排队论模型。我们首先使用单设备无并行的情况来建立符号体系:每个请求的执行时间,记为D,由于预填充长度统一而保持不变。由于一个请求就能使GPU饱和,我们通过先到先服务(FCFS)调度请求,不进行批处理。假设泊松到达率为R,且满足利用率条件 $RD < 1$,平均TTFT($Avg\_TTFT$)可以用M/D/1队列模型【47, John F Shortle等人, Fundamentals of queueing theory, 2018】的封闭形式来建模:
其中第一项代表执行时间,第二项对应排队延迟。基于公式1,我们在下面引入并行性。
算子间并行模型。对于2路算子间并行,我们假设请求级延迟变为 $D_s$,最慢的阶段需要 $D_m$ 完成。我们有 $D \approx D_s \approx 2 \times D_m$,因为层间激活通信可忽略不计【33, Zhuohan Li等人, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, arXiv, 2023】、【59, Lianmin Zheng等人, Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning, 2022, USENIX OSDI】。2路算子间并行的平均TTFT推导如下:
图4:在使用不同并行策略于两台A100 GPU上服务一个66B参数LLM时的平均TTFT。
算子内并行模型。对于算子内并行,我们引入一个加速系数 $K$,其中 $1 < K < 2$,反映了由算子内并行的高通信开销导致的不完美加速。执行时间为 $D_s = D/K$,2路算子内并行的平均TTFT为:
并行策略比较。比较公式2和公式3:在较低速率下,执行时间(第一项)是主要因素,算子内并行因其减少了执行时间而更有效。随着速率增加,排队延迟(第二项)变得更加显著,算子间并行变得更有优势,这与图4(a)的观察一致。
影响并行偏好的其他因素。预填充阶段对并行性的偏好也受到TTFT SLO和加速系数K的影响。从图4(a)可以看出:更严格的SLO会使算子内并行更具优势,因为它能减少执行时间。K的值取决于输入长度、模型架构、通信带宽和放置策略等因素【46, Mohammad Shoeybi等人, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】、【59, Lianmin Zheng等人, Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning, 2022, USENIX OSDI】。如图4(b)所示,K值的降低显著减弱了算子内并行的效果。§4将开发算法来优化资源和并行配置,同时考虑这些因素。
与预填充实例不同,解码实例遵循一种独特的计算模式:它从预填充实例接收KV缓存和第一个输出token,然后逐个生成后续的token。对于解码实例,我们的优化目标是使用最少的计算资源来满足应用的TPOT要求。
批处理策略。由于单个解码任务是严重受带宽限制的,批处理是避免GPU利用率低(从而提高每GPU良率)的关键,如图3(b)所示。在预填充和解码阶段共存的现有系统中,增加解码批处理大小很困难,因为它与满足延迟目标相冲突,特别是在高请求率的场景下。这是因为共享GPU导致预填充和解码任务之间产生竞争,从而在TTFT和TPOT之间形成权衡。例如,更高的到达率会产生更多的预填充任务,如果优先处理预填充任务,就需要更多的GPU时间来满足TTFT要求,这反过来又会对TPOT产生不利影响。
图5:在批大小=128、输入长度=256的条件下,服务13B LLM时,不同并行度下解码阶段的延迟和吞吐量。
解耦后的批处理优势。相反,解耦提供了一个解决方案,它允许将多个预填充实例分配给单个解码实例。这种方法可以在不牺牲TPOT的情况下,在专用的GPU上为解码阶段累积更大的批处理大小。
并行计划。解耦后,解码的批处理大小可能受到GPU内存容量的限制,因为它需要为所有活动请求维护KV缓存。通过模型并行扩展解码实例或利用LLM KV缓存的先进内存管理技术(如Paged-Attention【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】和GQA【10, Joshua Ainslie等人, Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023】),可以进一步扩展解码批处理大小,使其接近计算密集型。随着解码批处理大小持续增加并接近计算密集型,解码计算开始类似于预填充阶段。基于这一观察,我们在图5中研究了在大批量条件下,不同并行度下的延迟和吞吐量变化:算子内并行能降低延迟,但收益递减,这是由通信和分区后利用率降低引起的。算子间并行几乎可以线性地扩展吞吐量。因此,当TPOT SLO严格时,算子内并行对于降低TPOT以满足延迟目标至关重要。在此之上,算子间并行更适合线性提升吞吐量。
副本(Replication)的作用。值得注意的是,当模型可以装入单个GPU内存时,除了模型并行,副本也是预填充和解码实例的一个有竞争力的选项,可以线性扩展系统的速率容量。它也可能通过将R替换为R/N(假设请求被平均分配到N个副本)来减少排队延迟——如公式1所示——但代价是在GPU内存中维护额外的模型权重副本。
我们已经为每个阶段选择批处理和并行策略制定了基本原则。在本节中,我们讨论并解决在实际部署解耦的预填充和解码阶段时遇到的一些挑战。
可变的预填充长度。§3假设所有请求的提示长度是统一的。在实际部署中,根据LLM应用的不同,请求的长度是不均匀的。这种不均匀性可能导致应用了算子间并行的预填充实例出现流水线气泡(pipeline bubbles)【28, Yanping Huang等人, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019】、【36, Deepak Narayanan等人, Pipedream: Generalized pipeline parallelism for dnn training, 2019, SOSP】,因为不同长度请求的流水线阶段执行时间会变化。这导致结果与使用M/D/1队列模型得出的结论有轻微偏差。为了解决这个问题,§4开发了基于工作负载搜索并行策略的算法,并借助调度来最小化气泡(§4.3)。
通信开销。从预填充实例向解码实例传输KV缓存会产生显著的开销。例如,一个512-token请求在OPT-66B上的KV缓存大小约为1.13GB。假设平均到达率为10 rps,我们需要每秒传输11.3GB数据——或等效于90Gbps的带宽才能使开销变得不可见。虽然许多现代LLM的GPU集群配备了InfiniBand(例如800Gbps),但在跨节点带宽有限的情况下,DistServe依赖于普遍可用的节点内NVLINK,其中A100 GPU之间的峰值带宽为600GB/s,同样使得传输开销可以忽略不计(见§6.3)。然而,这一要求对预填充和解码实例的放置施加了额外约束,我们在下一节中会考虑到这一点。
综合考量。通过本节的分析,我们确定了工作负载模式、放置约束、SLO要求、并行策略和资源分配是设计解耦服务系统时需要考虑的关键参数,它们形成了一个复杂的考量网络。如何自动地在这个搜索空间中导航,找到能实现最佳每GPU良率的配置是具有挑战性的,我们将在接下来解决这个问题。
我们构建了DistServe来解决上述挑战。给定模型、工作负载特性、延迟要求和SLO达成目标,DistServe将确定 (a) 预填充和解码实例的并行策略,(b) 每种实例类型的部署数量,以及 (c) 如何将它们放置到物理集群上。我们将这个解决方案称为放置(placement)。我们的目标是找到一个能最大化每GPU良率的放置方案。
如§3.3所述,一个关键的设计考虑是管理解耦的预填充和解码阶段之间的通信,考虑到不同的集群设置。在本节中,我们首先提出两种放置算法:一种用于具有高速跨节点网络的集群(§4.1),另一种用于缺乏此类基础设施的环境(§4.2);后者引入了额外的约束。然后,我们开发了在线调度优化,以适应现实世界工作负载的细微差别(§4.3)。
算法1描述。在配备了Infiniband的高节点亲和性集群上,跨节点传输KV缓存的开销可以忽略不计,DistServe可以在任意两个节点间部署预填充和解码实例,不受限制。我们为这种情况提出了一种两级放置算法:我们首先分别为预填充和解码实例优化并行配置,以获得阶段级别的最佳每GPU良率;然后,我们使用副本来匹配整体的流量速率。
挑战与解决方案。然而,为单一实例类型(如预填充实例)找到最优的并行配置仍然具有挑战性,因为缺乏一个简单的解析公式来计算SLO达成率(即满足TTFT要求的请求百分比),特别是当工作负载具有多样化的输入输出长度和不规则的到达模式时。通过真实测试平台进行性能分析来评估SLO是耗时的。因此,我们借助构建一个模拟器来估计SLO达成率,假设事先知道工作负载的到达过程以及输入输出长度的分布。虽然短期内的间隔无法预测,但较长时间尺度(如小时或天)的工作负载模式通常是可预测的【33, Zhuohan Li等人, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, arXiv, 2023】、【55, Hong Zhang等人, Shepherd: Serving dnns in the wild, 2023】。DistServe从历史请求跟踪中拟合一个分布,并从该分布中重新采样新的跟踪作为模拟器的输入工作负载来计算SLO达成率。接下来,DistServe简单地枚举放置方案,并通过二分搜索和模拟试验找到满足SLO达成率目标的最大速率。
算法1流程。算法1概述了这个过程。我们枚举所有可行的并行配置,受限于集群容量,用于预填充和解码实例。然后,对于一个特定的预填充阶段配置,我们使用simu_prefill通过二分搜索模拟并找到其最大良率(类似地,使用simu_decode处理解码阶段)。在确定了预填充和解码实例的最优并行配置后,我们根据它们的良率复制它们,以达到用户要求的整体流量速率。
输入: LLM G, 每个实例的节点数限制 N, 每个节点的GPU数 M, GPU内存容量 C, 工作负载 W, 流量速率 R.
输出: 最佳放置方案 best_plm.
configp, configd ← ∅,∅
for intra_op ∈ {1, 2,..., M} do
for inter_op ∈ {1, 2, ..., N × M / intra_op} do
if G.size / (inter_op × intra_op) < C then
config ← (inter_op, intra_op)
Ĝ ← parallel(G, config)
// 模拟预填充
config.goodput ← simu_prefill(Ĝ, W)
if configp.goodput / configp.num_gpus < config.goodput / config.num_gpus then
configp ← config
// 模拟解码
config.goodput ← simu_decode(Ĝ, W)
if configd.goodput / configd.num_gpus < config.goodput / config.num_gpus then
configd ← config
n, m ← ⌈ R / configp.goodput ⌉, ⌈ R / configd.goodput ⌉
best_plm ← (n, configp, m, configd)
return best_plm
算法复杂度。算法1的复杂度为$O(NM^2)$,其中N是每个实例的节点限制,M代表现代集群中每个节点的典型GPU数量(例如8)。搜索空间是可控的,在我们最大的设置中,求解时间不到1.3分钟,如§6.5所示。
模拟器构建。算法1依赖于一个模拟器,在给定工作负载和并行计划的情况下,估算不同SLO和SLO达成目标下的良率。为了构建一个准确的模拟器,我们分别分析了预填充和解码阶段的FLOPs和内存访问次数,并使用一个延迟模型来近似推理执行时间。详细信息见附录A。由于DNN工作负载的高度可预测性【23, Arpan Gujarati等人, Serving DNNs like clockwork: Performance predictability from the bottom up, 2020, OSDI】、【33, Zhuohan Li等人, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, arXiv, 2023】,模拟器与真实的性能分析结果非常吻合,这在§6.4中得到了验证。
带宽受限场景。到目前为止,我们已经开发了算法1,假设我们可以在集群的任意两个节点之间(或在同一个节点上)放置预填充和解码实例,并且KV缓存传输利用高带宽网络。在许多实际集群中,一个节点内的GPU可以访问高带宽的NVLINK,而分布在不同节点间的GPU带宽有限。我们接下来开发一个算法来解决这个约束。
基本解决方案。一个直接的解决方案是始终将预填充和解码实例放置在同一个节点上,利用节点内普遍可用的NVLINK。对于大型模型,例如具有175B参数(350GB)的模型,我们可能甚至无法在一个8-GPU节点(80G × 8 = 640G < 350 × 2GB)中托管一对预填充和解码实例。我们将此作为额外的放置约束,并与模型并行性共同优化,如算法2所示。
核心洞察。关键洞察在于,KV缓存传输只发生在预填充和解码实例的相应层之间。利用算子间并行,我们将层分组为阶段,并将每个实例划分为段,称为实例段(instance segments),每个段维护一个特定的算子间阶段。通过将相同阶段的预填充和解码段放置在单个节点内,我们强制中间状态的传输只通过NVLINK进行。在一个节点内,我们为同一实例的段设置相同的并行和资源分配。鉴于每个节点的GPU数量通常有限(通常为8),我们可以枚举一个节点内的所有可能配置,并使用模拟器来确定产生最佳良率的配置。
算法2流程。如算法2所述,我们首先枚举算子间并行度,以获得所有可能的实例段。对于每个段,我们通过调用get_intra_node_configs获取所有可能的节点内并行配置。然后我们使用模拟来找到最优的一个,并复制它以满足目标流量速率。
输入: LLM G, 每个实例的节点数限制 N, 每个节点的GPU数 M, GPU内存容量 C, 工作负载 W, 流量速率 R.
输出: 最佳放置方案 best_plm.
config* ← ∅
for inter_op ∈ {1, 2,..., N} do
P ← get_intra_node_configs(G, M, C, inter_op)
for Pp ∈ P do
for Pd ∈ P do
if Pp.num_gpus + Pd.num_gpus ≤ M then
config ← (inter_op, Pp, Pd)
Ĝp, Ĝd ← parallel(G, config)
config.goodput ← simulate(Ĝp, Ĝd, W)
if config*.goodput / config*.num_gpus < config.goodput / config.num_gpus then
config* ← config
n ← ⌈ R / config*.goodput ⌉
best_plm ← (n, config*)
return best_plm
系统架构。DistServe的运行时架构如图6所示。DistServe采用简单的FCFS(先到先服务)调度策略。所有传入的请求到达一个中心控制器,然后被分派到队列最短的预填充实例进行预填充处理,随后分派到负载最轻的解码实例进行解码步骤。这个设置虽然简单,但通过几项关键增强进行了优化,以适应现实世界工作负载的细微差别。
图6:DistServe运行时系统架构
减少流水线气泡。为了减轻由非均匀提示长度引起的流水线气泡(§3.3),我们调度请求的方式是平衡流水线中所有批次的执行时间。这是通过注意到对于预填充和解码实例,批次中新token的数量是批次实际执行时间的可靠指标来实现的。对于预填充实例,我们对目标模型和GPU进行性能分析,找出使GPU饱和所需的最短提示长度Lm。我们通过批处理多个短于Lm的请求或单独调度长于Lm的请求,来调度总序列长度接近Lm的预填充批次。对于解码实例,我们将Lm设置为最大批次大小。
应对突发流量。工作负载中的突发性可能导致大量KV缓存从预填充实例传输到解码实例,有使解码实例内存过载的风险。为了规避这一点,DistServe采用“拉取”方法进行KV缓存传输,而不是“推送”方法——解码实例根据需要从预填充实例获取KV缓存,使用预填充实例的GPU内存作为排队缓冲区。这样,预填充实例可以通过在处理完提示后简单地将KV缓存保留在GPU内存中,继续处理其他预填充任务。因此,每种类型的实例都按自己的节奏运行,无需复杂的协调。
重新规划。DistServe中的资源和并行计划是为特定的工作负载模式优化的,如果工作负载模式随时间变化,它可能会变得次优。DistServe实现了周期性的重新规划。一个工作负载分析器监控关键参数,如请求的平均输入输出长度、平均到达率等。如果检测到显著的模式变化,DistServe将根据最近的历史数据触发重新运行放置算法。这个过程是迅速的——所提出的算法在几秒钟内运行完毕(§6.5),并且重新加载LLM权重可以在几分钟内完成——远短于现实世界工作负载变化的每小时尺度。
抢占和容错。DistServe没有实现像抢占【26, Mingcong Han等人, Microsecond-scale preemption for concurrent GPU-accelerated DNN inferences, 2022, OSDI】和容错【58, Kai Zhao等人, Ft-cnn: Algorithm-based fault tolerance for convolutional neural networks, 2021, IEEE Transactions on Parallel and Distributed Systems】这样的高级运行时策略,这些策略与解耦是互补的。尽管如此,我们讨论了它们如何融入DistServe。在DistServe中,FCFS策略可能导致“护航效应”,即较长的请求在预填充阶段阻塞较短的请求。结合现有文献【53, Bingyang Wu等人, Fast distributed inference serving for large language models, arXiv preprint arXiv:2305.05920, 2023】中建议的抢占策略,可以提高效率,并且在我们的系统架构内是可行的。虽然不是当前DistServe的主要焦点,但容错是需要考虑的一个关键方面。在传统的基于共存和副本的系统中,一个实例的故障通常不会干扰其他副本实例。然而,在DistServe中,预填充和解码实例之间的依赖性引入了故障传播的风险。例如,映射到多个预填充实例的单个解码实例的故障可能会瘫痪整个服务和集群。我们将这两者都作为未来的工作。
DistServe是一个用于LLM的端到端分布式服务系统,包括一个放置算法模块、一个RESTful API前端、一个编排层和一个并行执行引擎。算法模块、前端和编排层用6.5K行Python代码实现。并行执行引擎用8.1K行C++/CUDA代码实现。
系统组件。放置算法模块实现了§4中提到的算法和模拟器,为特定的模型和集群设置提供放置决策。前端支持与OpenAI API兼容的接口,客户端可以指定采样参数,如最大输出长度和温度。编排层管理预填充和解码实例,负责请求分发、KV缓存传输和结果交付。它利用NCCL【6, Nvidia collective communications library (nccl), 2023】进行跨节点GPU通信,并使用异步的CudaMemcpy进行节点内通信,这避免了在传输过程中阻塞GPU计算。每个实例由一个并行执行引擎驱动,该引擎使用Ray【35, Philipp Moritz等人, Ray: A distributed framework for emerging AI applications, 2018, USENIX OSDI】的actor来实现GPU工作者,这些工作者以分布式方式执行LLM推理并管理KV缓存。它集成了许多近期的LLM优化,如连续批处理【54, Gyeong-In Yu等人, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, USENIX OSDI】、FlashAttention【20, Tri Dao等人, Flashattention: Fast and memoryefficient exact attention with io-awareness, 2022】、PagedAttention【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】,并支持流行的开源LLM,如OPT【56, Susan Zhang等人, Opt: Open pre-trained transformer language models, 2022】和LLaMA【51, Hugo Touvron等人, Llama: Open and efficient foundation language models, 2023】。
集群测试平台
- 硬件配置:
- 4个节点,共32个GPU。
- 每个节点配备8个NVIDIA SXM A100-80GB GPU,通过NVLINK连接。
- 跨节点带宽为25Gbps,属于低节点亲和性环境。因此,实验主要采用针对该环境的放置算法(§4.2)。
模型与工作负载
- 模型:
- 选择了学术界和工业界广泛使用的OPT模型系列,包括13B、66B和175B三种规模。
- 使用FP16精度进行所有实验。
- 选用使用经典多头注意力(MHA)的OPT模型,以对KV缓存传输开销施加足够压力。
- 数据集与应用场景(如表1和图7所示):
- 聊天机器人(Chatbot):使用ShareGPT数据集【8】,该数据集包含用户与ChatGPT的对话。针对不同大小的OPT模型设置了不同的TTFT和TPOT SLO。
- 代码补全(Code completion):使用HumanEval数据集【14】,包含164个编程问题。由于是实时编码助手,设置了严格的TTFT和TPOT SLO。
- 摘要生成(Summarization):使用LongBench数据集【13】的摘要任务。该数据集输入长度远长于其他两个,因此设置了宽松的TTFT和严格的TPOT SLO。
- 请求生成:由于数据集不含时间戳,使用泊松分布生成不同请求率的请求到达时间。
- 基线系统:
- vLLM【32】:一个代表性的LLM服务系统,支持连续批处理和PagedAttention,但将预填充和解码计算共存。
- DeepSpeed-MII【3】:支持分块预填充(chunked-prefill),通过分解长提示来缓解预填充-解码干扰,但无法完全消除。
- 软件配置:DistServe由Python和C++/CUDA实现,依赖于NCCL和Ray等库。
表1:评估中使用的工作负载及延迟要求。
图7:(a) ShareGPT, (b) HumanEval, 和 (c) LongBench 数据集的输入输出长度分布。
评估指标
- SLO达成率(SLO attainment):主要评估指标。
- 在给定的SLO达成率目标(例如90%)下,关注两个方面:
1. 最大每GPU良率(goodput):系统能处理的最大请求率。
2. 系统能承受的最小SLO:衡量系统对延迟要求的鲁棒性。
聊天机器人(Chatbot)应用(图8):
- 实验内容:在ShareGPT数据集上,针对OPT-13B、66B、175B模型,比较DistServe与vLLM和DeepSpeed-MII的性能。
- 实验结果:
- 良率:在90% SLO达成率下,DistServe支持的请求率比vLLM高2.0至4.6倍,比DeepSpeed-MII高1.6至7.4倍。
- SLO鲁棒性:DistServe能承受比vLLM严格1.8至3.2倍、比DeepSpeed-MII严格1.7至1.8倍的SLO。
- 分析结论:
- DistServe通过解耦消除了预填充-解码干扰。其放置算法能找到非平凡的并行策略(例如,对于175B模型,预填充实例为3路流水线并行+3路张量并行,解码实例为3路流水线并行+4路张量并行),有效平衡负载并以最低成本满足SLO。
- vLLM因共存导致解码阶段严重减速,TPOT过高,无法满足聊天应用的严格TPOT要求。
- DeepSpeed-MII的分块预填充虽然有所缓解,但其本身比完整预填充慢,导致难以满足TTFT要求。
图8:在ShareGPT数据集上使用OPT模型的聊天机器人应用。
代码补全(Code completion)应用(图9a):
- 实验内容:在HumanEval数据集上,使用OPT-66B模型进行比较。
- 实验结果:DistServe支持的请求率比vLLM高5.7倍,比DeepSpeed-MII高1.6倍;能承受的SLO比两者都严格1.4倍。
- 分析结论:代码补全任务对TTFT要求极高。DistServe通过解耦消除了解码任务的干扰,并利用搜索算法自动增加预填充实例的算子内并行度,从而降低了预填充延迟,满足了更多请求的TTFT要求。
摘要生成(Summarization)应用(图9b):
- 实验内容:在LongBench数据集上,使用OPT-66B模型进行比较。
- 实验结果:DistServe支持的请求率比vLLM高4.3倍,比DeepSpeed-MII高1.8倍;能承受的SLO比vLLM严格12.6倍,比DeepSpeed-MII严格2.6倍。
- 分析结论:摘要任务输入长,预填充计算压力大,但TTFT要求宽松,TPOT变得至关重要。vLLM由于共存,在长预填充任务下解码阶段严重减速,无法满足TPOT要求。DistServe则无此问题。
图9:分别在HumanEval和LongBench数据集上使用OPT-66B进行代码补全和摘要任务。
图10:左图:在ShareGPT数据集上使用DistServe服务OPT-175B的延迟分解。右图:三种OPT模型的KV缓存传输时间的累积分布函数(CDF)。
表2:模拟器与真实系统在不同速率下报告的SLO达成率比较。
图11:消融实验。
图12:算法运行时间
不同场景下的适用性。本文专注于优化良率的场景,并提出了DistServe。然而,LLM服务场景多样,优化目标和资源限制各异,不存在“一刀切”的解决方案。
- 吞吐量优化场景:在对延迟不敏感的离线应用中,用户对响应时间要求较低。此时,服务系统可以更关注最大化整体吞吐量而非良率。在这种情况下,诸如带“搭载”功能的分块预填充技术【3, Deepspeed model implementations for inference (mii), 2023】、【9, Amey Agrawal等人, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, arXiv preprint arXiv:2308.16369, 2023】可能更受青睐,因为它们能将每个批次填满到计算密集型阈值,从而在每次迭代中保持更高的GPU利用率。
- 资源受限场景:小型企业和个人研究者通常缺乏在大型集群上部署LLM的资源。在资源受限的环境(如只有少数甚至单个GPU)中,DistServe的设计空间被极大地限制,难以通过调整并行策略和资源分配来有效提升服务性能。在这种情况下,更简单的架构选择,如非解耦系统【3, Deepspeed model implementations for inference (mii), 2023】、【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】,可能会降低部署复杂性并优化运营效率。
- 长上下文场景:越来越多的GPT模型支持极长的上下文,例如Claude3【11】、Gemini-1.5【22】和Large World Model (LWM)【34】,它们都拥有1M的上下文窗口。在这种场景下,随着提示长度的增加,KV缓存的大小线性增长,传输开销也会增加。然而,预填充的计算量是二次方增长的,因此传输与预填充任务的相对时长会减少。同时,更长的上下文进一步加剧了预填充和解码任务之间计算需求的不匹配,导致它们之间的干扰增加。因此,DistServe提出的解耦方法在长上下文服务中仍然具有前景。
推理服务。近期有大量关于推理服务的工作,从通用生产级系统如TorchServe【7】和NVIDIA Triton【19】,到专门为Transformer-based LLM优化的系统【9, 18, 21, 33, 50, 53, 54, 60】。其中,Orca【54, Gyeong-In Yu等人, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, USENIX OSDI】引入了连续批处理以增加吞吐量。vLLM【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】提出了paged-attention用于细粒度的KV缓存管理。SARATHI【9, Amey Agrawal等人, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, arXiv preprint arXiv:2308.16369, 2023】提出了一种分块预填充方法。FastServe【53, Bingyang Wu等人, Fast distributed inference serving for large language models, arXiv preprint arXiv:2305.05920, 2023】实现了迭代级抢占调度以减轻长任务造成的排队延迟。然而,它们都采用了预填充和解码处理的共存方法,从而导致了严重的干扰。也有一些同期的工作如Splitwise【38】、TetriInfer【27】和DéjàVu【49】采用了类似的解耦思想来优化LLM推理,进一步证实了这种方法的有效性。不同的是,DistServe更强调良率优化场景,并更深入地研究了网络带宽方面的问题。
良率优化系统。优化良率是深度学习应用中的一个热门话题。Pollux【39】通过动态调整作业资源来提高DL集群的调度性能。Sia【29】引入了一种异构感知调度方法。Clockwork【23】和Shepherd【55】为传统小模型提供了延迟感知的调度和抢占。AlpaServe【33, Zhuohan Li等人, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, arXiv, 2023】专注于LLM,采用模型并行来统计复用GPU执行。然而,它只针对非自回归生成。DistServe是首个为自回归LLM推理优化良率的工作。
资源解耦。资源解耦系统【17, 25, 43】将硬件资源从传统的单体服务器基础设施中解耦出来,分离成资源池进行独立管理。DistServe共享了这一理念,通过解耦其系统组件,允许独立的资源扩展和管理。
用于训练的模型并行。DistServe与大量关于训练中模型并行的工作【28, 36, 40, 46, 59】是正交的。如§3.3所述,推理服务工作负载具有训练场景中没有的独特特性。这些系统与DistServe的交集在于它们实现多维度模型并行的方法。DistServe可以将新的并行优化集成到其放置搜索算法中。
本文提出了DistServe,一种新的LLM服务架构,它解耦了预填充和解码的计算。DistServe最大化了每GPU的良率——即在遵守SLO达成目标的条件下,每个配置的GPU每秒可服务的最大请求速率,从而在保证满足SLO的前提下,将每个LLM查询的成本降低了高达7.4倍。我们的发现证实,随着延迟成为LLM服务越来越重要的指标,预填充和解码的解耦是保证性能提升和服务质量的关键策略。
为了准确模拟不同放置策略的良率,我们使用一个分析模型来预测LLM推理中预填充和解码阶段的执行时间。在现代LLM服务系统【18, NVIDIA Corporation, Fastertransformer, 2019】、【32, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】、【53, Bingyang Wu等人, Fast distributed inference serving for large language models, arXiv preprint arXiv:2305.05920, 2023】中,像Softmax和LayerNorm这样的内存密集型操作通常与矩阵乘法核心融合以提高效率。因此,GEMM(通用矩阵乘法)主导了整体延迟,我们的分析主要集中在它们上面。
模型架构相关符号:
- $h$: 隐藏层大小
- $n$: 注意力头数量
- $s$: 每个头的大小 ($h = n \cdot s$)
- $m$: FFN中间层大小
注意:如果使用张量并行,h, n, 和 m 应除以张量并行的大小。
批处理特征相关符号:
- $B$: 批大小
- $l_0, l_1, \dots, l_{B-1}$: 批内每个请求的输入长度
- $t$: 批内的总token数 ($t = \sum_{i=0}^{B-1} l_i$)
- $t_2$: 输入长度的平方和 ($t_2 = \sum_{i=0}^{B-1} l_i^2$),用于FlashAttention【20, Tri Dao等人, Flashattention: Fast and memoryefficient exact attention with io-awareness, 2022】。
非注意力操作。由于注意力操作使用特殊优化的核心,我们首先讨论预填充阶段的其他四个矩阵乘法:
这些操作的算术强度(AI)为$O(t)$。在NVIDIA A100-80GB GPU上,当AI超过156时,操作是计算密集型的。由于在实际情况中t通常可以达到几百,所有这些操作都是计算密集型的。因此,我们可以根据总FLOPs来建模这些操作的延迟:
注意力操作。接下来,我们讨论使用FlashAttention【20】优化的预填充注意力操作。由于注意力只在同一请求的token之间操作,当前实现为同一批次中的每个请求启动注意力核心。对于一个注意力头和一个有$l$个token的请求,注意力核心需要执行总共 $2sl + 3sl \cdot (l/b) \approx 3sl \cdot (l/b)$ 次内存读写,以及 $2sl^2 + sl(l/b) \approx 2sl^2$ FLOPs。所以AI是 $2b/3 = 10.677$ (当b=16时) 或 $21.333$ (当b=32时),表明它在A100 GPU上是内存密集型操作。因此,整个注意力层的延迟(包括所有请求和所有头)可以建模为:
整体延迟。总体而言,预填充阶段的延迟可以建模为:
我们使用$C_3$来量化其他开销,如Python运行时、系统噪音等。然后我们使用性能分析和插值来确定$C_1$, $C_2$, 和$C_3$的值。
非注意力操作。同样,我们首先关注解码阶段的以下GEMM:
这些操作的AI为$O(B)$。B受限于GPU内存大小和严格的延迟要求,因此在现有的服务场景中,这些操作是内存密集型的。总内存读写量为 $8Bh + 4h^2 + 2hm + 2Bm$,由于h和m通常远大于B,我们可以将延迟建模为:
注意力操作。对于解码注意力操作,对于一个注意力头和一个有$l$个已生成token的请求,它需要执行$3sl$次内存读写,以及$2sl$ FLOPs。它是内存密集型的,所以我们可以将解码注意力的延迟建模为:
整体延迟。总的来说,解码阶段的延迟是:
这里我们没有引入开销项(如预填充阶段的$C_3$),因为$4h^2 + 2hm$已经是一个常数,开销可以并入$C_4$。同样,我们使用性能分析和插值来确定$C_4$和$C_5$的值。
表3显示了在§6.2的端到端实验中,DistServe为预填充和解码实例选择的张量并行(TP)和流水线并行(PP)配置。
表3:DistServe在端到端实验中选择的并行策略。
图13和图14显示了DistServe与基线系统在§6.2相同设置下的端到端性能,但SLO达成率目标改为99%。我们可以看到,在更严格的SLO达成率目标下,与vLLM相比,DistServe仍然可以支持3倍至8倍的更高速率和1.24倍至6.67倍更严格的SLO。与DeepSpeed-MII相比,DistServe可以实现1.32倍至8倍的更高速率和1.20倍至1.58倍更严格的SLO。
图13:在ShareGPT数据集上使用OPT模型的聊天机器人应用。
图14:分别在HumanEval和LongBench数据集上使用OPT-66B进行代码补全和摘要任务。