Splitwise: Efficient Generative LLM Inference Using Phase Splitting
文章标题: Splitwise: 利用阶段切分实现高效的生成式LLM推理
作者/机构: Pratyush Patel (University of Washington), Esha Choukse, Chaojie Zhang, Aashaka Shah, Íñigo Goiri, Saeed Maleki, Ricardo Bianchini (Microsoft)
A1 主要贡献
核心问题: 生成式大型语言模型(LLM)的推理任务通常在昂贵且耗电的GPU上运行,而全球范围内的GPU产能正面临紧缩。这些推理任务具有两个计算特性截然不同的阶段:计算密集的“提示词计算(prompt computation)”阶段和内存密集的“令牌生成(token generation)”阶段。即使采用最先进的批处理(batching)和调度技术,令牌生成阶段也无法充分利用最新GPU的计算资源,导致资源浪费和成本增加。
研究目标: 本文旨在通过分离并独立管理LLM推理的两个阶段来提升系统整体的GPU利用率和效率。研究目标是设计出在吞吐量、成本和功耗方面都得到优化的LLM推理集群,以应对当前LLM服务部署中的挑战。
创新点 (Splitwise):
本文提出了Splitwise,一种模型部署和调度技术,其核心思想是将LLM推理请求的两个阶段(提示词计算和令牌生成)拆分到不同的机器上执行。这一方法带来了以下创新和优势:
1. 阶段性资源管理: 通过将两个阶段分离,Splitwise允许为每个阶段使用最适合其特性的硬件,并进行独立的资源管理。例如,计算密集的提示词计算阶段可以使用最新的、计算能力强的GPU,而内存密集的令牌生成阶段则可以运行在计算能力较弱但成本和功耗更低的硬件上。
2. 高效的状态迁移: 为了实现阶段分离,提示词计算阶段生成的上下文缓存(KV-cache)需要被高效地传输到令牌生成机器。Splitwise利用当今GPU集群中常见的高速后端互连(如InfiniBand)上的优化网络库,以低延迟传输请求状态,从而避免了性能损失。
3. 灵活的集群设计: 基于Splitwise,本文探索了同构和异构的LLM推理集群设计,旨在优化吞吐量、成本和功耗。研究表明,可以针对不同的目标(如为用户提供更高的性价比Perf/\$,或为云服务提供商提供更高的每瓦性能Perf/W)来配置集群。
4. **显著的性能提升**: 评估结果显示,与现有设计相比,采用Splitwise的集群可以在成本降低20%的情况下实现高达1.4倍的吞吐量提升;或者在相同的功耗和成本预算下,提供2.35倍的吞吐量。
**本文的主要贡献总结如下**:
1. 对LLM推理中提示词和令牌生成阶段在NVIDIA A100和H100 GPU上的执行和利用率模式进行了广泛的特性分析,并使用了生产环境中的真实追踪数据。
2. 提出了Splitwise技术,通过将提示词计算和令牌生成阶段拆分到不同机器上,优化了可用硬件的利用率。
3. 利用Splitwise对同构和异构集群部署进行了设计探索,以优化整体成本、请求吞吐量和配置功率。
4. 使用生产环境追踪数据对基于Splitwise设计的系统进行了评估。
## A3 背景知识与关键洞察
### 背景知识
**大型语言模型(LLM)**: 现代LLM基于Transformer架构,使用注意力机制和多层感知机层来理解输入并生成输出。本文关注的生成式LLM通常是仅解码器(decoder-only)或编码器-解码器(encoder-decoder)模型。
**生成式LLM的推理阶段**: 生成式LLM的推理过程如图1所示。当收到提示词查询后,所有输入令牌在一次迭代中并行计算,生成第一个输出令牌,这被称为**提示词处理阶段(prompt processing phase)**。此阶段注意力层生成的上下文被保存在键值缓存(KV-cache)中,供后续所有令牌生成迭代使用。第一个令牌生成后,后续令牌的生成仅使用最新生成的令牌和KV-cache作为模型前向传播的输入,这使得**令牌生成阶段(token generation phase)**相比计算密集的提示词阶段,更加受到内存带宽和容量的限制。
TABLE I: NVIDIA A100 vs. H100 规格。
TABLE IV: 在Llama-70B上无批处理时,A100与H100的P50请求指标对比。
A2 方法细节
Splitwise的总体设计: 基于上述洞察,本文提出Splitwise技术,将LLM推理的提示词阶段和生成阶段拆分到不同的机器上。如图10所示,系统维护两个独立的机器池:一个用于提示词处理,一个用于令牌处理。此外,还有一个混合池,可根据工作负载动态伸缩。所有机器都预装了所需的模型。当新请求到达时,调度器会为其分配一对机器(一个提示词机器和一个令牌机器)。提示词机器负责处理输入提示词并生成第一个令牌和KV缓存,然后将KV缓存发送给令牌机器,由后者继续生成剩余的令牌直到响应完成。在请求率较低时,Splitwise旨在优化延迟;在请求率较高时,则旨在避免因机器池分片而导致的性能或吞吞吐量下降。
分层调度: Splitwise采用如图10所示的两级分层调度。集群级调度器(CLS)①负责机器池管理和路由传入的推理请求。机器级调度器(MLS)②则维护每个机器上的待处理队列并管理请求的批处理。
集群级调度(CLS)
机器池管理: CLS维护提示词池、令牌池和混合池③。Splitwise根据预期的请求负载和输入/输出令牌分布,初始分配机器到提示词池或令牌池。为了减少分片并在高负载下满足SLO,提示词池或令牌池中的机器可以被动态地移入和移出混合池。混合池中的机器保留其原始身份(提示词或令牌机器),一旦其队列中没有相反类型的任务,就会返回其原始池。如果负载分布与初始假设有较大偏差,Splitwise会采用粗粒度的机器再利用策略,在提示词池和令牌池之间移动机器。
请求路由: CLS使用“加入最短队列”(Join the Shortest Queue, JSQ)调度算法【39, Analysis of Join-the-Shortest-Queue Routing for Web Server Farms, Performance Evaluation 2007】、【85, Analysis of JSQ Policy on Soft Real-time Scheduling in Cluster, HPCAsia 2000】为每个请求分配一个提示词机器和一个令牌机器。队列长度由待处理的令牌数量定义。为了重叠KV缓存传输与提示词计算以减少开销,系统会同时分配提示词和令牌机器。当路由请求时,如果待处理队列超过某个阈值,CLS会首先在混合池中寻找目标机器。如果混合池也满了,它会进一步在相反的池中寻找(例如,让一个令牌机器运行提示词任务),并将该机器移入混合池。混合池中的机器采用混合批处理方式运行。一旦混合请求的队列被清空,CLS会将该机器转换回其原始池。
机器级调度(MLS)
调度器职责: MLS在每台机器上运行,负责跟踪GPU内存利用率,维护待处理队列④,为每次迭代决定批处理内容,并向CLS报告相关状态。
提示词机器: MLS采用先到先服务(FCFS)来调度提示词。由于实验结果(图6a)显示,当提示词令牌总数超过2048后吞吐量会下降,MLS将多个提示词的批处理限制在总共2048个令牌以内。
令牌机器: MLS对令牌同样采用FCFS调度,并尽可能多地进行批处理。由于令牌生成的吞吐量会随着批处理大小的增加而持续扩展直到内存耗尽(图6b),MLS会跟踪内存使用情况,并在机器接近内存耗尽时开始对令牌进行排队。
混合机器: 为了满足TTFT的SLO,MLS必须优先运行提示词,并立即调度待处理队列中的任何新提示词。如果机器正在运行令牌阶段且没有额外容量来运行提示词阶段,MLS将抢占令牌。为了避免令牌阶段因抢占而饿死,系统会增加令牌的优先级随其年龄增长,并限制每个请求可被抢占的次数。
KV缓存传输
传输开销: 在Splitwise中,KV缓存需要在提示词机器和令牌机器之间传输⑤,这个传输延迟是Splitwise引入的主要开销。如图11a所示,一种简单的方式是在提示词阶段完全结束后,串行地传输KV缓存。这个传输延迟会直接影响TBT和端到端延迟,特别是对于大型提示词,即使使用高速InfiniBand,开销也可能相当可观。
优化传输: Splitwise通过将KV缓存传输与提示词阶段的计算重叠来优化这一过程。如图11b所示,当提示词机器计算完LLM的每一层后,该层对应的KV缓存也随之生成。在每层计算结束时,系统会触发一个异步传输,将该层的KV缓存发送出去,同时提示词计算继续进行到下一层。这种逐层传输的方式减少了传输开销,并使得令牌机器可以更早开始令牌生成阶段,提示词机器也可以更早释放KV缓存内存。
自适应传输策略: 逐层KV缓存传输需要在每层进行细粒度的同步,这可能会对TTFT产生性能干扰,特别是对于较小的提示词。然而,小提示词的KV缓存总体积较小,本身就不需要通过逐层传输来隐藏延迟。因此,Splitwise会根据批处理中的令牌数选择最佳的传输技术:对较小的提示词使用串行KV缓存传输,对较大的提示词使用逐层传输。
使用Splitwise进行资源配置
机器类型: 本文提出了四种主要的基于Splitwise的系统变体,命名方式为“提示词机器类型-令牌机器类型”:
* Splitwise-AA: 提示词和令牌池都使用DGX-A100机器。
* Splitwise-HH: 提示词和令牌池都使用DGX-H100机器。
* Splitwise-HA: 提示词池使用DGX-H100,令牌池使用DGX-A100。
* Splitwise-HHcap: 提示词和令牌池都使用DGX-H100,但令牌机器的功率被限制在额定功率的70%(每个GPU限制50%)。
TABLE V: 评估的Splitwise设计,均以DGX-A100为基准进行归一化。
机器数量确定: 为了确定集群中提示词和令牌机器的合适数量,本文使用一个事件驱动的集群模拟器来搜索设计空间。该过程需要输入:(1) 目标集群设计,(2) LLM特定的性能模型,(3) 服务的目标提示词和令牌大小分布的简短追踪,(4) SLO,(5) 约束条件(如吞吐量),(6) 优化目标(如最小化成本)。
搜索空间与优化: 图12展示了一个为编码工作负载在Splitwise-HH设计下寻找提示词和令牌机器数量的二维搜索空间示例。模拟器输出各种延迟指标的百分位,然后筛选出满足所有SLO的集群配置,并根据目标函数(最小化成本)找到最优点。例如,图中用⋆标记的配置(27个提示词机器和3个令牌机器)是在满足70 RPS吞吐量下成本最低的方案。优化目标可以有三个:吞吐量、成本和功耗,它们对云服务提供商(CSP)和用户的吸引力各有不同。
A4 实验
实验环境
硬件配置: 实验在Microsoft Azure上的两台DGX-A100和两台DGX-H100虚拟机上进行,每台机器包含8个GPU,通过InfiniBand连接。H100机器的互连带宽是A100的两倍(400 Gbps vs 200 Gbps)。
软件配置: Splitwise的KV缓存传输机制在vLLM【51, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】之上实现,使用了优化的GPU驱动通信库MSCCL++【11, MSCCL++: A GPU-driven communication stack for scalable AI applications. GitHub】。代码已开源【1, Add Splitwise Implementation to vLLM. GitHub】。此外,实现了先进的混合连续批处理机制。
数据集: 使用了来自Microsoft Azure两个生产LLM推理服务(编码和对话)的追踪数据【4, Azure Public Dataset: Azure LLM Inference Trace 2023. GitHub】。
模型架构: 实验使用了BLOOM-176B和Llama-70B模型。
模拟器: 为了在更大规模上评估Splitwise,构建了一个开源的事件驱动模拟器SplitwiseSim【20, SplitwiseSim: LLM Serving Cluster Simulator. GitHub】。该模拟器使用通过硬件剖析构建的分段线性性能模型(MAPE < 3%)和通过基准测试建立的通信模型,并已通过与硬件实验的交叉验证确保了准确性。
服务等级目标(SLOs): 如表VI所示,SLO被定义为与在无竞争的DGX-A100上运行的请求相比的减速因子。评估要求P50、P90和P99的TTFT、TBT和E2E延迟共九个SLO都必须满足。
TABLE VI: SLO表示为与在无竞争的DGX-A100上运行的请求相比的减速。
基线: 将Splitwise设计与两个基线进行比较:Baseline-A100(仅由DGX-A100组成)和Baseline-H100(仅由DGX-H100组成)。两个基线都使用混合连续批处理。
实验结果
KV缓存传输延迟:
* 实验内容: 测量了随着提示词大小增加,naive(串行)和optimized(逐层)KV缓存传输的可见延迟。
* 实验结果: 如图14所示,与提示词计算时间相比,优化后的传输开销非常小(< 7%)。串行传输的延迟随提示词大小线性增加,而优化传输则将大部分延迟隐藏起来,在A100上非重叠传输时间约为8ms,在H100上约为5ms。如图15所示,对于端到端延迟,Splitwise引入的开销仅为0.8%,对用户几乎无感知。
等功耗下吞吐量优化的集群:
* 实验内容: 在固定的总功耗预算下(以40台DGX-H100的功耗为基准),配置不同设计的集群,并评估其在不同负载下的延迟和吞吐量。
* 实验结果: 如图16所示,对于编码和对话追踪,所有Splitwise设计都比Baseline-H100表现更好。随着负载增加,Baseline-H100因混合批处理导致TBT过高而首先无法满足SLO。Splitwise-HA在低TTFT和高吞吐量之间取得了良好平衡。图17显示,Splitwise机器实现了更好的批处理效率,大部分时间运行在较大的批次大小下。图18a总结了对话追踪的结果:与Baseline-A100相比,Splitwise-AA在相同功耗和成本下提供了2.15倍的吞吐量;Splitwise-HA在成本降低10%和相同功耗下提供了1.18倍的吞吐量。
其他集群优化:
* 实验内容: 评估了等成本下吞吐量优化、等吞吐量下功耗优化、等吞吐量下成本优化的集群设计。
* 实验结果:
* 等成本吞吐量优化 (图18b): Splitwise-AA以相同成本提供了比Baseline-H100高1.4倍的吞吐量,这对于不关心功耗和空间的用户非常有吸引力。
* 等吞吐量功耗优化 (图19a): Splitwise-HHcap能以比Baseline-H100低25%的功耗实现相同的吞吐量,这对CSP是明显的好处。
* 等吞吐量成本优化 (图19b): Splitwise-AA能以比Baseline-H100低25%的成本实现相同的吞吐量。
工作负载变化的鲁棒性:
* 实验内容: 测试Splitwise对工作负载变化的鲁棒性,包括将为编码服务优化的集群用于对话追踪,以及将为BLOOM-176B优化的集群用于Llama-70B。
* 实验结果: 如图20所示,Splitwise的同构设计(AA和HH)能很好地适应新工作负载,性能无影响。异构设计(HA和HHcap)吞吐量有约7%的下降,但仍远优于基线。当更换模型时,所有Splitwise设计都表现出优于基线的性能。
* 结论: Splitwise通过其智能调度能够适应工作负载需求的变化,对LLM模型、请求负载和令牌分布的变化具有鲁棒性。
批处理作业的集群设计:
* 实验内容: 分析在不考虑严格延迟SLO的批处理作业场景下各设计的性能。
* 实验结果: 在高负载下,Splitwise-AA和Baseline-A100具有最高的每成本吞吐量(0.89 RPS/$)。此时,由于所有机器都在混合池中进行混合批处理,Splitwise的行为趋同于基线。
A7 补充细节
讨论
对新模型的可扩展性: 所有现代基于Transformer的生成式LLM(包括混合专家模型MoE)都具有提示词处理和令牌生成这两个不同阶段。Splitwise利用了这两个阶段的特性,因此只要工作负载的自回归性质需要这两个阶段,该方法就适用。
替代计算硬件: 本文使用了NVIDIA H100和A100 GPU,但Splitwise的方法也适用于其他符合提示词和令牌阶段计算需求的硬件(包括CPU、FPGA、ASIC)。提示词阶段需要高计算能力和内存带宽,而令牌阶段需要中等计算能力、高内存容量和带宽。因此,像AMD MI-250 GPU或配备HBM的Intel SapphireRapids CPU可能成为有效的令牌机器。
提示词和令牌机器间的互连: 本文假设机器间使用InfiniBand连接。尽管H100和A100之间的InfiniBand连接在技术上可行,但在实践中可能不常见。替代方案可以是使用CPU连接的InfiniBand或以太网上的RoCE。由于优化的KV缓存传输减少了关键路径上的延迟,即使带宽降低10倍,该方法可能仍然有益。
异构提示词/令牌机器: 尽管Splitwise对不同的模型和输入追踪具有鲁棒性,但在数据中心部署不同类型的GPU(如Splitwise-HA)可能会给CSP带来其自身的挑战。
来回对话: 未来的聊天API可能会在GPU上缓存对话上下文,以避免重计算。这可能会改变提示词阶段的内存使用模式,并可能需要将KV缓存传回提示词机器,为下一轮对话做准备。
相关工作
异构调度和数据流系统: 先前的工作研究了各种交互式服务的异构调度,但通常将整个工作负载运行在同一台机器上。分布式数据流系统提供了通用的可编程性,但Splitwise是针对生成式LLM推理的专门两阶段设计,并利用了阶段感知的资源管理和高效批处理。
模型服务系统: 近期有许多工作在优化LLM推理服务的批处理、调度和内存使用。这些方法通常在同一台机器上处理提示词和令牌阶段。结合Splitwise,它们可以通过阶段拆分来进一步提高吞吐量和延迟。
其他推理系统: 在视频和机器学习服务中,调度器需要处理具有数据依赖性的模型链。推荐系统推理在模型内部和模型之间都表现出计算/内存的异构性,已有工作利用这种异构性在CPU和加速器之间选择性地调度请求或在异构硬件上划分计算/内存。Splitwise同样利用了LLM推理请求内部的异构性,但由于LLM工作负载特性和需求的不同,采用了不同的优化方法。
A5 结论
本文深入剖析了LLM推理中提示词计算和令牌生成阶段在系统利用率模式上的差异。基于这些洞察,我们设计了Splitwise,将这两个阶段分离到不同的机器上,从而实现了阶段性的资源管理。利用Splitwise,我们探索了为吞吐量、成本和功耗优化的集群设计,并证明了即使工作负载发生变化,这些设计也能保持良好性能。在满足性能SLO的前提下,与现有设计相比,Splitwise集群能够以低15%的功耗和相同的成本实现1.76倍的吞吐量提升,或者在成本和功耗相同的情况下实现2.35倍的吞吐量提升。
A6 附录
摘要
我们开源了评估Splitwise所需的关键组件;这些组件也可以被重新用于评估未来的LLM推理服务系统。我们的工件包括:
* 来自微软Azure两个LLM推理服务的生产环境追踪数据。
* 在vLLM中实现的Splitwise KV缓存传输机制的原型。
* SplitwiseSim,一个用于评估LLM推理集群中模型服务的离散事件模拟器。
由于硬件可用性有限,工件的功能性仅针对追踪数据和SplitwiseSim进行了测试。
工件清单(元信息)
- 数据集: 生产环境追踪数据,作为工件的一部分提供。
- 运行时环境: Linux / Ubuntu。
- 硬件: vLLM原型需要两台通过GPU InfiniBand连接的机器(例如NVIDIA DGX-A100, NVIDIA DGX-H100)。SplitwiseSim需要一台x86-64 CPU机器。
- 是否公开可用: 是。
- 代码许可证: MIT。
- 数据许可证: CC-BY。
- 存档DOI: 10.5281/zenodo.11003049。
描述
如何访问: 整个工件作为一个存档在Zenodo上提供:https://doi.org/10.5281/zenodo.11003049。各个组件也可在线获取:
* 生产环境追踪数据可从Azure公共数据集GitHub仓库下载【4】。
* KV缓存传输原型可从vLLM GitHub仓库下载,目前以拉取请求的形式提供【1】。
* SplitwiseSim及相关的实验和绘图脚本可从一个独立的GitHub仓库下载【20】。
硬件依赖: KV缓存传输原型需要两台通过InfiniBand连接的GPU机器。SplitwiseSim需要标准的x86-64 CPU机器。
软件依赖: KV缓存传输原型基于vLLM【51】和MSCCL++【11】构建。SplitwiseSim依赖于一小组公开可用的Python包,可通过附带的requirements.txt
文件安装。
数据集: 来自微软Azure的编码和对话追踪数据作为工件发布的一部分在线提供【4】。
安装和实验流程
请参考工件内的README文件以获取安装和使用说明。
💬 评论讨论
欢迎在这里分享您的想法和见解!