分布式大语言模型服务引擎vLLM论文解读

1 minute read

Published:

论文地址:Efficient Memory Management for Large Language Model Serving with PagedAttention

摘要

大语言模型(LLMs)的高吞吐量服务需要一次对足够多的请求进行批处理。然而,现有系统面临困境,因为每个请求的键值缓存(KV缓存)内存占用巨大,且会动态增长和收缩。若管理效率低下,内存会因碎片化和冗余重复而被大量浪费,从而限制批处理规模。为解决这一问题,我们提出了分页注意力(PagedAttention),这是一种受操作系统中经典虚拟内存和分页技术启发的注意力算法。在此基础上,我们构建了vLLM,这是一个大语言模型服务系统,它实现了:(1)将KV缓存内存浪费降至几乎为零;(2)在请求内部和请求之间灵活共享KV缓存,以进一步降低内存使用。我们的评估表明,与诸如FasterTransformer和Orca等最先进的系统相比,vLLM在相同延迟水平下,能将流行大语言模型的吞吐量提高2至4倍。在处理更长序列、更大模型以及更复杂的解码算法时,这种提升更为显著。vLLM的源代码可在https://github.com/vllm-project/vllm 公开获取。

图1。左图:在英伟达A100上为一个130亿参数的大语言模型提供服务时的内存布局。模型参数(灰色部分)在整个服务过程中一直保留在GPU内存中。键值缓存(KV缓存)的内存(红色部分)会根据每个服务请求进行分配和释放。少量内存(黄色部分)会临时用于激活操作。右图:vLLM 平滑了现有系统[31, 60]中键值缓存内存的快速增长曲线,显著提升了服务吞吐量。

1. 引言

像GPT [5, 37]和PaLM [9]这样的大语言模型(LLMs)的出现,催生了诸如编程助手[6, 18]和通用聊天机器人[19, 35]等新应用,这些应用正开始深刻影响我们的工作和日常生活。许多云服务公司[34, 44]竞相以托管服务的形式提供这些应用。然而,运行这些应用成本高昂,需要大量诸如GPU之类的硬件加速器。据近期估算,处理一个大语言模型请求的成本可能比传统关键字查询高出10倍[43]。鉴于这些高昂成本,提高大语言模型服务系统的吞吐量,进而降低每个请求的成本,变得愈发重要。

大语言模型的核心是自回归Transformer模型[53]。该模型基于输入(提示词)以及到目前为止它所生成的输出词元序列,一次生成一个词(词元)。对于每个请求,这一高成本的过程会不断重复,直到模型输出一个终止词元。这种顺序生成过程使得工作负载受限于内存,无法充分利用GPU的计算能力,从而限制了服务吞吐量。

通过将多个请求进行批处理,有可能提高吞吐量。然而,要批量处理多个请求,就需要高效管理每个请求的内存空间。例如,图1(左)展示了在拥有40GB内存的英伟达A100 GPU上运行一个130亿参数大语言模型时的内存分布情况。大约65%的内存分配给模型权重,这些权重在服务过程中保持不变。近30%的内存用于存储请求的动态状态。对于Transformer模型,这些状态由与注意力机制相关的键值张量组成,通常称为键值缓存(KV缓存)[41],它代表了先前词元的上下文,以便按顺序生成新的输出词元。其余一小部分内存用于其他数据,包括激活值——即在评估大语言模型时临时创建的张量。由于模型权重是固定的,且激活值仅占用GPU内存的一小部分,因此键值缓存的管理方式对于确定最大批处理规模至关重要。如图1(右)所示,如果管理效率低下,键值缓存内存会显著限制批处理规模,进而影响大语言模型的吞吐量。

在本文中,我们注意到现有的大语言模型服务系统[31, 60]在高效管理键值缓存内存方面存在不足。这主要是因为它们像大多数深度学习框架[33, 39]要求张量存储在连续内存中一样,将请求的键值缓存存储在连续的内存空间中。然而,与传统深度学习工作负载中的张量不同,键值缓存具有独特的特性:随着模型生成新的词元,它会随时间动态增长和收缩,并且其生存期和长度事先无法得知。这些特性使得现有系统的方法在两方面效率极低: 首先,现有系统[31, 60]存在内部和外部内存碎片化问题。为了在连续空间中存储请求的键值缓存,它们会预先分配一块与请求最大长度(例如2048个词元)相等的连续内存块。这可能会导致严重的内部碎片化,因为请求的实际长度可能比其最大长度短得多(例如图11)。此外,即使事先知道实际长度,预先分配仍然效率低下:由于在请求的生存期内整个内存块都被保留,其他较短的请求无法利用该内存块中当前未使用的任何部分。此外,外部内存碎片化也可能很严重,因为每个请求的预先分配大小可能不同。实际上,我们在图2中的分析结果表明,在现有系统中,键值缓存内存中只有20.4% - 38.2%用于存储实际的词元状态。 其次,现有系统无法利用内存共享的机会。大语言模型服务通常使用先进的解码算法,如并行采样和束搜索,每个请求会生成多个输出。在这些场景中,请求由多个序列组成,这些序列可以部分共享它们的键值缓存。然而,在现有系统中无法实现内存共享,因为这些序列的键值缓存存储在不同的连续空间中。

为了解决上述限制,我们提出了分页注意力(PagedAttention)算法,这是一种受操作系统(OS)解决内存碎片化和共享问题的方案——虚拟内存分页启发的注意力算法。分页注意力将请求的键值缓存划分为多个块,每个块可以包含固定数量词元的注意力键值。在分页注意力中,键值缓存的块不一定要存储在连续空间中。因此,我们可以像操作系统的虚拟内存那样,以更灵活的方式管理键值缓存:可以将块视为页面,词元视为字节,请求视为进程。这种设计通过使用相对较小的块并按需分配来减轻内部碎片化。此外,由于所有块大小相同,它消除了外部碎片化。最后,它能够在块的粒度上,在与同一请求相关的不同序列之间,甚至在不同请求之间实现内存共享。

在这项工作中,我们基于分页注意力构建了vLLM,这是一个高吞吐量的分布式大语言模型服务引擎,它实现了键值缓存内存几乎零浪费。vLLM采用了与分页注意力协同设计的块级内存管理和抢占式请求调度。vLLM支持各种不同规模的流行大语言模型,如GPT [5]、OPT [62]和LLaMA [52],包括那些超出单个GPU内存容量的模型。我们在各种模型和工作负载上的评估表明,与最先进的系统[31, 60]相比,vLLM将大语言模型服务的吞吐量提高了2 - 4倍,且完全不影响模型准确性。在处理更长序列、更大模型以及更复杂的解码算法时,这种提升更为显著(§4.3)。

综上所述,我们做出了以下贡献:

  • 我们明确了大语言模型服务中内存分配的挑战,并量化了它们对服务性能的影响。
  • 我们提出了分页注意力算法,这是一种在非连续分页内存中存储的键值缓存上运行的注意力算法,它受操作系统中的虚拟内存和分页机制启发。
  • 我们设计并实现了vLLM,这是一个基于分页注意力构建的分布式大语言模型服务引擎。
  • 我们在各种场景下对vLLM进行了评估,并证明它显著优于诸如FasterTransformer [31]和Orca [60]等先前的最先进解决方案。

图3. 现有系统中的键值(KV)缓存内存管理。存在三种类型的内存浪费——预留浪费、内部碎片浪费和外部碎片浪费,这些浪费导致其他请求无法装入内存。每个内存中的词元代表其键值缓存。请注意,相同的词元在不同位置时可能具有不同的键值缓存。

2. 背景

在本节中,我们将描述典型大语言模型的生成与服务流程,以及大语言模型服务中使用的迭代级调度。

2.1 基于 Transformer 的大语言模型

2.2 大语言模型服务与自回归生成

2.3 大语言模型的批处理技术

通过对多个请求进行批处理,可以提高大语言模型服务中的计算利用率。因为请求共享相同的模型权重,移动权重的开销会分摊到一批请求中,并且当批处理大小足够大时,计算开销会占主导地位。然而,对大语言模型服务的请求进行批处理并非易事,原因有二:

首先,请求可能在不同时间到达。简单的批处理策略要么让较早的请求等待较晚的请求,要么延迟传入的请求,直到较早的请求完成,这会导致显著的排队延迟。 其次,请求的输入和输出长度可能差异很大(图 11)。直接的批处理技术会对请求的输入和输出进行填充,使其长度相等,这会浪费 GPU 计算资源和内存。 为了解决这个问题,已经提出了细粒度的批处理机制,如蜂窝批处理 [16] 和迭代级调度 [60]。与在请求级别工作的传统方法不同,这些技术在迭代级别上运行。每次迭代后,已完成的请求从批处理中移除,新的请求被添加进来。因此,一个新请求在等待一次迭代后就可以被处理,而无需等待整个批处理完成。此外,通过特殊的 GPU 内核,这些技术无需对输入和输出进行填充。通过减少排队延迟和填充带来的低效率,细粒度批处理机制显著提高了大语言模型服务的吞吐量。

3. 大语言模型服务中的内存挑战

尽管细粒度批处理减少了计算资源的浪费,并使请求能够以更灵活的方式进行批处理,但能够一起批处理的请求数量仍然受到GPU内存容量的限制,特别是分配用于存储键值(KV)缓存的空间。换句话说,服务系统的吞吐量受限于内存。要克服这种内存限制,需要应对内存管理中的以下挑战:

  • 庞大的KV缓存:KV缓存的大小会随着请求数量迅速增长。例如,对于拥有130亿参数的OPT模型[62],单个词元的KV缓存需要800KB的空间,计算方式为2(键值向量)×5120(隐藏状态大小)×40(层数)×2(每个FP16的字节数)。由于OPT可以生成长达2048个词元的序列,存储一个请求的KV缓存所需的内存可能高达1.6GB。目前常见的GPU内存容量为几十GB。即使将所有可用内存都分配给KV缓存,也只能容纳几十个请求。此外,如图2所示,低效的内存管理会进一步降低批处理大小。另外,鉴于当前的发展趋势,GPU的计算速度增长比内存容量增长更快[17]。例如,从英伟达A100到H100,浮点运算次数(FLOPS)增加了两倍多,但GPU内存最大仍保持在80GB。因此,我们认为内存将成为越来越严重的瓶颈。
  • 复杂的解码算法:大语言模型服务为用户提供了一系列解码算法可供选择,每种算法对内存管理的复杂性都有不同影响。例如,当用户从单个输入提示词请求多个随机样本时(这是程序建议中的常见用例[18]),提示词部分的KV缓存(在我们的实验中占总KV缓存内存的12%,见§6.3)可以共享,以尽量减少内存使用。另一方面,由于样本结果不同以及它们对上下文和位置的依赖,自回归生成阶段的KV缓存不应共享。KV缓存的共享程度取决于所采用的具体解码算法。在诸如束搜索[49]等更复杂的算法中,不同请求的搜索束可以共享更大比例(内存节省高达55%,见§6.3)的KV缓存,并且共享模式会随着解码过程的推进而变化。
  • 未知输入和输出长度的调度:对大语言模型服务的请求在输入和输出长度上表现出可变性。这就要求内存管理系统能够适应各种提示词长度。此外,随着请求在解码过程中输出长度的增加,其KV缓存所需的内存也会扩大,可能会耗尽用于新请求或现有提示词正在进行的生成所需的可用内存。系统需要做出调度决策,例如从GPU内存中删除或交换出某些请求的KV缓存。

    3.1 现有系统中的内存管理

    由于当前大多数深度学习框架[33, 39]中的操作符要求张量存储在连续内存中,以前的大语言模型服务系统[31, 60]也将一个请求的KV缓存作为跨不同位置的连续张量进行存储。由于大语言模型的输出长度不可预测,它们会根据请求的最大可能序列长度为请求静态分配一块内存,而不考虑请求的实际输入或最终输出长度。

图3展示了两个请求:请求A的最大可能序列长度为2048,请求B的最大可能序列长度为512。现有系统中的内存块预分配方案主要存在三种内存浪费来源:为未来词元预留的槽位、由于为潜在的最大序列长度过度分配而导致的内部碎片,以及像伙伴分配器这样的内存分配器产生的外部碎片。外部碎片在服务请求之前就已知永远不会用于生成的词元。内部碎片同样未被使用,但只有在请求完成采样后才会意识到。它们都是纯粹的内存浪费。虽然预留的内存最终会被使用,但在整个请求期间保留该空间,特别是当预留空间很大时,会占用原本可用于处理其他请求的空间。我们在图2中展示了实验中内存浪费的平均百分比,结果表明以前系统中的实际有效内存可能低至20.4%。 尽管有人提出压缩[54]作为解决碎片问题的潜在方案,但在对性能敏感的大语言模型服务系统中,由于庞大的KV缓存,执行压缩并不现实。即使进行压缩,现有内存管理系统中为每个请求预先分配的内存块空间也会阻碍特定于解码算法的内存共享。

4. 方法

在这项工作中,我们开发了一种新的注意力算法——分页注意力(PagedAttention),并构建了一个大语言模型服务引擎vLLM,以应对第3节中概述的挑战。vLLM的架构如图4所示。vLLM采用集中式调度器来协调分布式GPU工作节点的执行。KV缓存管理器通过分页注意力,以分页方式有效地管理KV缓存。具体来说,KV缓存管理器通过集中式调度器发送的指令,管理GPU工作节点上的物理KV缓存内存。 接下来,我们将在4.1节描述分页注意力算法。在此基础上,分别在4.2节展示KV缓存管理器的设计,以及在4.3节说明它如何辅助分页注意力。然后,我们将阐述这种设计如何为各种解码方法促进有效的内存管理(4.4节),以及如何处理可变长度的输入和输出序列(4.5节)。最后,我们将展示vLLM的系统设计在分布式环境中是如何工作的(4.6节)。

4.1 分页注意力

4.2 KV缓存管理器

vLLM内存管理器背后的核心思想类似于操作系统中的虚拟内存[25]。操作系统将内存划分为固定大小的页面,并将用户程序的逻辑页面映射到物理页面。连续的逻辑页面可以对应非连续的物理内存页面,这使得用户程序可以像访问连续内存一样访问内存。此外,物理内存空间无需提前完全预留,使操作系统能够根据需要动态分配物理页面。vLLM利用虚拟内存背后的思想来管理大语言模型服务中的KV缓存。在分页注意力的支持下,我们将KV缓存组织成固定大小的KV块,类似于虚拟内存中的页面。

一个请求的KV缓存表示为一系列逻辑KV块,随着新词元及其KV缓存的生成,从左到右填充。最后一个KV块中未填充的位置预留给未来生成的词元。在GPU工作节点上,块引擎分配一块连续的GPU动态随机存取存储器(DRAM),并将其划分为物理KV块(在CPU随机存取存储器(RAM)上进行交换时也会执行相同操作;见4.5节)。KV块管理器还维护块表——每个请求的逻辑KV块与物理KV块之间的映射。每个块表项记录一个逻辑块对应的物理块以及已填充位置的数量。将逻辑KV块和物理KV块分离,使vLLM能够动态增长KV缓存内存,而无需提前为所有位置预留内存,这消除了现有系统中的大部分内存浪费,如图2所示。

4.3 使用分页注意力和vLLM进行解码

接下来,我们通过图6中的示例,展示vLLM在单个输入序列的解码过程中如何执行分页注意力并管理内存:

  1. 初始阶段:与操作系统的虚拟内存类似,vLLM最初不需要为最大可能生成的序列长度预留内存。相反,它只预留必要的KV块,以容纳提示词计算期间生成的KV缓存。在这种情况下,提示词有7个词元,因此vLLM将前2个逻辑KV块(0和1)分别映射到2个物理KV块(7和1)。在预填充步骤中,vLLM使用传统的自注意力算法(例如[13])生成提示词和第一个输出词元的KV缓存。然后,vLLM将前4个词元的KV缓存存储在逻辑块0中,接下来3个词元的KV缓存存储在逻辑块1中。剩余的槽位预留给后续的自回归生成阶段。
  2. 第一次自回归解码步骤:vLLM在物理块7和1上使用分页注意力算法生成新词元。由于最后一个逻辑块中仍有一个槽位可用,新生成的KV缓存存储在那里,并且块表中的“已填充”记录会更新。
  3. 第二次解码步骤:由于最后一个逻辑块已满,vLLM将新生成的KV缓存存储在一个新的逻辑块中;vLLM为其分配一个新的物理块(物理块3),并将此映射存储在块表中。 总体而言,对于每次解码迭代,vLLM首先选择一组候选序列进行批处理(详见4.5节),并为新需要的逻辑块分配物理块。然后,vLLM将当前迭代的所有输入词元(即提示词阶段请求的所有词元以及生成阶段请求的最新词元)连接成一个序列,并将其输入到大语言模型中。在大语言模型的计算过程中,vLLM使用分页注意力内核访问以逻辑KV块形式存储的先前KV缓存,并将新生成的KV缓存保存到物理KV块中。在一个KV块中存储多个词元(块大小(>1)),使得分页注意力内核能够并行处理更多位置的KV缓存,从而提高硬件利用率并减少延迟。然而,较大的块大小也会增加内存碎片化。我们在7.2节研究了块大小的影响。

同样,随着更多词元及其KV缓存的生成,vLLM动态地为逻辑块分配新的物理块。由于所有块都是从左到右填充,并且只有在所有先前块都已满时才会分配新的物理块,vLLM将一个请求的所有内存浪费限制在一个块内,因此它可以有效地利用所有内存,如图2所示。这使得更多请求能够装入内存进行批处理,从而提高吞吐量。一旦一个请求完成生成,其KV块可以被释放,以存储其他请求的KV缓存。在图7中,我们展示了vLLM为两个序列管理内存的示例。两个序列的逻辑块映射到GPU工作节点中块引擎预留空间内的不同物理块。两个序列相邻的逻辑块在物理GPU内存中不需要连续,并且物理块的空间可以被两个序列有效利用。

4.4 在其他解码场景中的应用

4.3节展示了分页注意力和vLLM如何处理基本的解码算法,如贪心解码和采样,这些算法将一个用户提示词作为输入并生成单个输出序列。在许多成功的大语言模型应用[18, 34]中,大语言模型服务必须提供更复杂的解码场景,这些场景具有复杂的访问模式和更多的内存共享机会。在本节中,我们将展示vLLM在这些场景中的普遍适用性。

  • 并行采样:在基于大语言模型的编程助手[6, 18]中,大语言模型为单个输入提示词生成多个采样输出;用户可以从各种候选输出中选择一个最喜欢的输出。到目前为止,我们隐含地假设一个请求生成单个序列。在本文的其余部分,我们假设更一般的情况,即一个请求生成多个序列。在并行采样中,一个请求包括多个共享相同输入提示词的样本,这使得提示词的KV缓存也可以共享。通过其分页注意力和分页内存管理,vLLM可以轻松实现这种共享并节省内存。

图8展示了两个输出的并行解码示例。由于两个输出共享相同的提示词,在提示词阶段,我们只为提示词状态的一份副本预留空间;两个序列提示词的逻辑块都映射到相同的物理块:两个序列的逻辑块0和1分别映射到物理块7和1。由于单个物理块可以映射到多个逻辑块,我们为每个物理块引入一个引用计数。在这种情况下,物理块7和1的引用计数均为2。在生成阶段,两个输出采样不同的输出词元,需要为KV缓存单独存储。vLLM在块粒度上为需要被多个序列修改的物理块实现了写时复制机制,类似于操作系统虚拟内存中的写时复制技术(例如,在创建进程时)。具体来说,在图8中,当样本A1需要写入其最后一个逻辑块(逻辑块1)时,vLLM识别到相应物理块(物理块1)的引用计数大于1;它分配一个新的物理块(物理块3),指示块引擎从物理块1复制信息,并将引用计数减为1。接下来,当样本A2写入物理块1时,引用计数已减为1;因此A2直接将其新生成的KV缓存写入物理块1。 总之,vLLM使得存储提示词KV缓存的大部分空间可以在多个输出样本之间共享,除了最后一个逻辑块,它由写时复制机制管理。通过在多个样本之间共享物理块,内存使用可以大大减少,特别是对于长输入提示词。

  • 束搜索:在诸如机器翻译[59]等大语言模型任务中,用户期望大语言模型输出最合适的前(k)个翻译。束搜索[49]被广泛用于从大语言模型中解码出最可能的输出序列,因为它减轻了完全遍历样本空间的计算复杂性。该算法依赖于束宽参数(k),它决定了在每一步保留的顶级候选数量。在解码过程中,束搜索通过考虑所有可能的词元来扩展束中的每个候选序列,使用大语言模型计算它们各自的概率,并从(k \cdotV)个候选中保留最可能的前(k)个序列,其中(V)是词汇表大小。

与并行解码不同,束搜索不仅可以共享初始提示词块,还可以在不同候选之间共享其他块,并且共享模式会随着解码过程的推进而动态变化,类似于操作系统中由复合分叉创建的进程树。图9展示了vLLM如何为(k = 4)的束搜索示例管理KV块。在虚线所示的迭代之前,每个候选序列都使用了4个完整的逻辑块。所有束候选共享第一个块0(即提示词)。候选3从第二个块开始与其他候选不同。候选0 - 2共享前3个块,并在第四个块处发散。在后续迭代中,前4个最可能的候选都来自候选1和2。由于原始候选0和3不再是顶级候选,它们的逻辑块被释放,相应物理块的引用计数减少。vLLM释放所有引用计数达到0的物理块(块2、4、5、8)。然后,vLLM分配新的物理块(块9 - 12)来存储新候选的新KV缓存。现在,所有候选共享块0、1、3;候选0和1共享块6,候选2和3进一步共享块7。 以前的大语言模型服务系统需要在束候选之间频繁复制KV缓存。例如,在图9所示的情况下,虚线之后,候选3需要复制候选2的很大一部分KV缓存才能继续生成。vLLM的物理块共享显著减少了这种频繁的内存复制开销。在vLLM中,不同束候选的大多数块都可以共享。与并行解码一样,仅当新生成的词元在旧的共享块内时,才应用写时复制机制。这只涉及复制一个块的数据。

共享前缀:通常,大语言模型(LLM)的用户会提供对任务的(较长)描述,其中包括指令以及示例输入和输出,这也被称为系统提示[36]。该描述与实际任务输入连接起来,形成请求的提示。大语言模型会基于完整的提示生成输出。图10展示了一个示例。此外,通过提示工程,共享前缀还可以进一步优化,以提高下游任务的准确性[26, 27]。

对于这类应用,许多用户提示都有一个共享前缀,因此大语言模型服务提供商可以提前存储该前缀的键值(KV)缓存,以减少在前缀上的冗余计算。在vLLM中,大语言模型服务提供商可以通过为一组预定义的共享前缀预留一组物理块来轻松实现这一点,就像操作系统处理跨进程的共享库一样。带有共享前缀的用户输入提示可以简单地将其逻辑块映射到已缓存的物理块(最后一个块标记为写时复制)。提示阶段的计算只需在用户的任务输入上执行。

混合解码方法:前面讨论的解码方法呈现出不同的内存共享和访问模式。尽管如此,vLLM能够促进对具有不同解码偏好的请求进行同时处理,而现有系统无法高效做到这一点。这是因为vLLM通过一个将逻辑块转换为物理块的通用映射层,隐藏了不同序列之间复杂的内存共享关系。大语言模型及其执行内核只需看到每个序列的物理块ID列表,无需处理跨序列的共享模式。与现有系统相比,这种方法拓宽了具有不同采样要求的请求的批处理机会,最终提高了系统的整体吞吐量。

4.5 调度与抢占

当请求流量超过系统容量时,vLLM必须对一部分请求进行优先级排序。在vLLM中,我们对所有请求采用先来先服务(FCFS)的调度策略,以确保公平性并防止饥饿现象。当vLLM需要抢占请求时,它会确保先到达的请求先得到处理,最新的请求先被抢占。

大语言模型服务面临一个独特的挑战:大语言模型的输入提示长度可能差异很大,并且生成的输出长度事先无法得知,这取决于输入提示和模型本身。随着请求数量及其输出的增加,vLLM可能会耗尽GPU的物理块来存储新生成的KV缓存。在这种情况下,vLLM需要回答两个经典问题:(1)应该逐出哪些块?(2)如果再次需要,如何恢复被逐出的块?通常,逐出策略会使用启发式方法来预测哪个块在未来最晚被访问,然后逐出该块。在我们的情况中,由于我们知道一个序列的所有块是一起被访问的,所以我们实施一种要么全有要么全无的逐出策略,即要么逐出一个序列的所有块,要么一个都不逐出。此外,一个请求中的多个序列(例如,一次束搜索请求中的束候选序列)会作为一个序列组进行协同调度。由于这些序列之间可能存在内存共享,一个序列组中的序列总是会一起被抢占或重新调度。为了回答如何恢复被逐出块的第二个问题,我们考虑两种技术:

  • 交换:这是大多数虚拟内存实现中使用的经典技术,即将被逐出的页面复制到磁盘上的交换空间。在我们的场景中,我们将被逐出的块复制到CPU内存。如图4所示,除了GPU块分配器,vLLM还包括一个CPU块分配器,用于管理交换到CPU随机存取存储器(RAM)的物理块。当vLLM为新生成的词元耗尽空闲物理块时,它会选择一组序列进行逐出,并将它们的KV缓存转移到CPU。一旦它抢占了一个序列并逐出其块,vLLM会停止接受新请求,直到所有被抢占的序列完成处理。一旦一个请求完成,其块将从内存中释放,被抢占序列的块会被重新调入,以继续该序列的处理。请注意,通过这种设计,交换到CPU RAM的块数量永远不会超过GPU RAM中的总物理块数量,因此CPU RAM上的交换空间受分配给KV缓存的GPU内存限制。
  • 重新计算:在这种情况下,当被抢占的序列重新调度时,我们只需重新计算KV缓存。请注意,重新计算的延迟可能会显著低于原始延迟,因为解码时生成的词元可以与原始用户提示连接起来作为一个新提示,这样它们在所有位置的KV缓存都可以在一次提示阶段迭代中生成。

交换和重新计算的性能取决于CPU RAM和GPU内存之间的带宽以及GPU的计算能力。我们将在7.3节研究交换和重新计算的速度。

4.6 分布式执行

许多大语言模型的参数规模超过了单个GPU的容量[5, 9]。因此,有必要将它们分布在多个GPU上,并以模型并行的方式执行[28, 63]。这就需要一个能够处理分布式内存的内存管理器。vLLM通过支持在Transformer上广泛使用的Megatron - LM风格的张量模型并行策略,在分布式环境中表现出色[47]。这种策略遵循单程序多数据(SPMD)执行调度,其中线性层被划分以执行块级矩阵乘法,并且GPU通过全规约操作不断同步中间结果。具体来说,注意力算子在注意力头维度上进行拆分,每个SPMD进程负责多头注意力中的一部分注意力头。

我们观察到,即使采用模型并行执行,每个模型分片仍然处理相同的一组输入词元,因此需要相同位置的KV缓存。因此,如图4所示,vLLM在集中式调度器中设置了一个单一的KV缓存管理器。不同的GPU工作节点共享该管理器,以及逻辑块到物理块的映射。这种通用映射允许GPU工作节点使用调度器为每个输入请求提供的物理块来执行模型。尽管每个GPU工作节点具有相同的物理块ID,但每个工作节点只为其对应的注意力头存储一部分KV缓存。

在每一步中,调度器首先为批处理中的每个请求准备包含输入词元ID的消息以及每个请求的块表。接下来,调度器将此控制消息广播给GPU工作节点。然后,GPU工作节点使用输入词元ID开始执行模型。在注意力层中,GPU工作节点根据控制消息中的块表读取KV缓存。在执行过程中,GPU工作节点无需调度器的协调,通过全规约通信原语同步中间结果,如[47]所述。最后,GPU工作节点将此迭代采样的词元发送回调度器。总之,GPU工作节点无需在内存管理上进行同步,因为它们只需要在每次解码迭代开始时接收所有内存管理信息以及步骤输入。

5. 实现

vLLM是一个端到端的服务系统,拥有FastAPI [15]前端和基于GPU的推理引擎。前端扩展了OpenAI API [34]接口,允许用户为每个请求定制采样参数,例如最大序列长度和束宽 (k) 。vLLM引擎由8500行Python代码和2000行C++/CUDA代码编写而成。我们用Python开发了包括调度器和块管理器在内的控制相关组件,同时为诸如分页注意力(PagedAttention)等关键操作开发了定制的CUDA内核。对于模型执行器,我们使用PyTorch [39]和Transformers [58]实现了如GPT [5]、OPT [62]和LLaMA [52]等流行的大语言模型。我们使用NCCL [32]在分布式GPU工作节点间进行张量通信。

5.1 内核级优化

由于分页注意力引入了现有系统无法高效支持的内存访问模式,我们开发了多个GPU内核来对其进行优化。

  • 融合重塑与块写入:在每个Transformer层中,新的键值(KV)缓存被分割成块,重塑为针对块读取优化的内存布局,然后保存到块表指定的位置。为了最小化内核启动开销,我们将这些操作融合为单个内核。
  • 融合块读取与注意力:我们对FasterTransformer [31]中的注意力内核进行调整,使其根据块表读取KV缓存并即时执行注意力操作。为确保合并内存访问,我们分配一个GPU线程束来读取每个块。此外,我们增加了对请求批内可变序列长度的支持。
  • 融合块复制:写时复制机制发出的块复制操作可能作用于不连续的块。如果我们使用cudaMemcpyAsync API,这可能导致大量小数据移动的调用。为减轻这种开销,我们实现了一个内核,将不同块的复制操作批处理为单个内核启动。

    5.2 支持多种解码算法

    vLLM使用三种关键方法实现了多种解码算法:分叉(fork)、追加(append)和释放(free)。分叉方法从现有序列创建一个新序列。追加方法将一个新词元追加到序列中。最后,释放方法删除序列。例如,在并行采样中,vLLM使用分叉方法从单个输入序列创建多个输出序列。然后,它在每次迭代中使用追加方法向这些序列添加新词元,并使用释放方法删除满足停止条件的序列。vLLM在束搜索和前缀共享中也应用了相同的策略。我们相信,未来通过组合这些方法也可以支持新的解码算法。

    6. 评估

    在本节中,我们评估vLLM在各种工作负载下的性能。

6.1 实验设置

  • 模型与服务器配置:我们使用参数规模为130亿、660亿和1750亿的OPT [62]模型,以及参数规模为130亿的LLaMA [52]模型进行评估。如大语言模型排行榜[38]所示,130亿和660亿参数规模在大语言模型中较为常见,而1750亿参数规模与著名的GPT - 3 [5]模型相同。在所有实验中,我们使用谷歌云平台上配备英伟达A100 GPU的A2实例。详细的模型规模和服务器配置如表1所示。
  • 工作负载:我们基于ShareGPT [51]和Alpaca [50]数据集合成工作负载,这些数据集包含真实大语言模型服务的输入和输出文本。ShareGPT数据集是用户与ChatGPT [35]共享对话的集合。Alpaca数据集是由GPT - 3.5通过自指令(self - instruct)[57]生成的指令数据集。我们对数据集进行词元化,并使用其输入和输出长度来合成客户端请求。如图11所示,ShareGPT数据集的输入提示平均比Alpaca数据集长8.4倍,输出平均长5.8倍,且方差更高。由于这些数据集不包含时间戳,我们使用泊松分布以不同的请求速率生成请求到达时间。
  • 基线1:FasterTransformer:FasterTransformer [31]是一个针对延迟高度优化的分布式推理引擎。由于FasterTransformer没有自己的调度器,我们实现了一个自定义调度器,其动态批处理机制类似于Triton [30]等现有服务系统。具体来说,在每次实验中,我们根据GPU内存容量尽可能设置一个较大的最大批大小 (B) 。调度器最多接收 (B) 个最早到达的请求,并将这批请求发送给FasterTransformer进行处理。
  • 基线2:Orca:Orca [60]是一个针对吞吐量优化的先进大语言模型服务系统。由于Orca未公开可用,我们实现了自己版本的Orca。我们假设Orca使用伙伴分配算法来确定存储KV缓存的内存地址。基于为请求输出过度预留空间的程度,我们实现了三个版本的Orca:
    • Orca(Oracle):我们假设系统知晓请求实际将生成的输出长度。这展示了Orca的性能上限,在实际中难以实现。
    • Orca(Pow2):我们假设系统为输出过度预留的空间最多为2倍。例如,如果真实输出长度为25,它会为输出预留32个位置。
    • Orca(Max):我们假设系统总是预留直至模型最大序列长度的空间,即2048个词元。
  • 关键指标:我们关注服务吞吐量。具体而言,使用不同请求速率的工作负载,我们测量系统的归一化延迟,即每个请求的端到端延迟的平均值除以其输出长度,这与Orca [60]中的做法相同。一个高吞吐量的服务系统在高请求速率下应保持较低的归一化延迟。在大多数实验中,我们使用1小时的追踪数据评估系统。作为例外,由于成本限制,对于OPT - 175B模型,我们使用15分钟的追踪数据。

6.2 基本采样

我们在三个模型和两个数据集上,通过基本采样(每个请求一个样本)评估vLLM的性能。图12的第一行展示了在ShareGPT数据集上的结果。这些曲线表明,随着请求速率的增加,延迟最初以缓慢的速度上升,然后突然激增。这是因为当请求速率超过服务系统的容量时,队列长度会持续无限增长,请求的延迟也会随之增加。

在ShareGPT数据集上,与Orca(Oracle)相比,vLLM能够维持高出1.7倍至2.7倍的请求速率,与Orca(Max)相比,高出2.7倍至8倍,同时保持相近的延迟。这是因为vLLM的分页注意力(PagedAttention)能够高效管理内存使用,因此比Orca能够处理更多的请求批处理。例如,如图13a所示,对于OPT - 13B模型,vLLM同时处理的请求数量比Orca(Oracle)多2.2倍,比Orca(Max)多4.3倍。与FasterTransformer相比,vLLM能够维持高出22倍的请求速率,因为FasterTransformer没有采用细粒度调度机制,并且像Orca(Max)一样,内存管理效率低下。

图12的第二行和图13b展示了在Alpaca数据集上的结果,其趋势与ShareGPT数据集相似。一个例外是图12(f),在该图中vLLM相对于Orca(Oracle)和Orca(Pow2)的优势不太明显。这是因为OPT - 175B的模型和服务器配置(表1)提供了大量的GPU内存空间来存储键值缓存(KV缓存),而Alpaca数据集的序列较短。在这种设置下,尽管Orca(Oracle)和Orca(Pow2)内存管理效率低下,但它们也能够对大量请求进行批处理。因此,这些系统的性能受限于计算能力而非内存。

6.3 并行采样和束搜索

我们使用两种流行的采样方法——并行采样和束搜索,评估分页注意力中内存共享的有效性。在并行采样中,一个请求中的所有并行序列可以共享提示的KV缓存。如图14第一行所示,随着要采样的序列数量增加,与Orca基线相比,vLLM带来了更大的性能提升。同样,图14第二行展示了不同束宽下束搜索的结果。由于束搜索允许更多的共享,vLLM展现出更大的性能优势。在OPT - 13B和Alpaca数据集上,vLLM相对于Orca(Oracle)的性能提升,从基本采样时的1.3倍,在束宽为6的束搜索中提升到了2.3倍。

图15绘制了内存节省量,通过共享节省的块数除以不共享时的总块数来计算。我们发现在并行采样中内存节省6.1% - 9.8%,在束搜索中节省37.6% - 55.2%。在ShareGPT数据集的相同实验中,我们看到并行采样内存节省16.2% - 30.5%,束搜索节省44.3% - 66.3%。

6.4 共享前缀

我们探究vLLM在不同输入提示共享前缀情况下的有效性,如图10所示。对于模型,我们使用支持多语言的LLaMA - 13B [52]。对于工作负载,我们使用WMT16 [4]英德翻译数据集,并合成两个前缀,其中包含一条指令和几个翻译示例。第一个前缀包含单个示例(即单样本),而另一个前缀包含5个示例(即少样本)。如图16(a)所示,当共享单样本前缀时,vLLM的吞吐量比Orca(Oracle)高1.67倍。此外,当共享更多示例时(图16(b)),vLLM的吞吐量比Orca(Oracle)高3.58倍。

6.5 聊天机器人

聊天机器人[8, 19, 35]是大语言模型最重要的应用之一。为实现聊天机器人,我们将聊天历史和最后一条用户查询连接成一个提示,让模型生成回复。我们使用ShareGPT数据集合成聊天历史和用户查询。由于OPT - 13B模型的上下文长度有限,我们将提示截断为最后1024个词元,并让模型最多生成1024个词元。我们不在不同对话轮次之间存储KV缓存,因为这样做会在对话轮次之间占用其他请求的空间。

图17表明,与三个Orca基线相比,vLLM能够维持高出2倍的请求速率。由于ShareGPT数据集包含许多长对话,大多数请求的输入提示有1024个词元。由于伙伴分配算法,无论Orca基线如何预测输出长度,它们都会为请求输出预留1024个词元的空间。因此,三个Orca基线表现相似。相比之下,vLLM能够有效处理长提示,因为分页注意力解决了内存碎片化和预留问题。

7. 消融研究

在本节中,我们研究vLLM的各个方面,并通过消融实验评估我们所做的设计选择。

7.1 内核微基准测试

分页注意力中的动态块映射会影响涉及存储的KV缓存的GPU操作性能,即块的读/写和注意力计算。与现有系统相比,我们的GPU内核(§5)涉及额外的开销,如访问块表、执行额外分支以及处理可变序列长度。如图18a所示,与高度优化的FasterTransformer实现相比,这导致注意力内核延迟增加20 - 26%。我们认为这种开销较小,因为它仅影响注意力算子,而不影响模型中的其他算子,如线性层。尽管存在开销,但分页注意力使vLLM在端到端性能上显著优于FasterTransformer(§6)。

7.2 块大小的影响

块大小的选择对vLLM的性能有重大影响。如果块大小过小,vLLM可能无法充分利用GPU并行读取和处理KV缓存的能力。如果块大小过大,内部碎片化会增加,共享的概率会降低。

在图18b中,我们使用ShareGPT和Alpaca的跟踪数据,在固定请求速率下通过基本采样评估不同块大小的vLLM性能。在ShareGPT跟踪数据中,块大小从16到128时性能最佳。在Alpaca跟踪数据中,虽然块大小为16和32时效果良好,但更大的块大小会显著降低性能,因为序列变得比块大小更短。在实践中,我们发现块大小为16足以有效利用GPU,并且在大多数工作负载中足够小以避免显著的内部碎片化。因此,vLLM将默认块大小设置为16。

7.3 比较重新计算和交换

vLLM同时支持重新计算和交换作为恢复机制。为了解这两种方法之间的权衡,我们评估它们的端到端性能,并对其开销进行微基准测试,如图19所示。

我们的结果表明,在块大小较小时,交换会产生过多开销。这是因为小块大小通常会导致CPU和GPU之间大量的小数据传输,从而限制了有效PCIe带宽。相比之下,重新计算的开销在不同块大小下保持不变,因为重新计算不使用KV块。因此,当块大小较小时,重新计算更高效;当块大小较大时,交换更高效,尽管重新计算的开销从未超过交换延迟的20%。对于16到64的中等块大小,这两种方法在端到端性能上表现相当。

8. 讨论

  • 将虚拟内存和分页技术应用于其他GPU工作负载:虚拟内存和分页的思想在大语言模型服务中管理键值缓存是有效的,因为该工作负载需要动态内存分配(由于输出长度事先未知),并且其性能受GPU内存容量的限制。然而,这并非适用于所有的GPU工作负载。例如,在深度神经网络训练中,张量的形状通常是静态的,因此内存分配可以提前进行优化。再如,在为非大语言模型的深度神经网络提供服务时,提高内存效率可能并不会带来任何性能提升,因为其性能主要受计算能力的限制。在这些场景中,引入vLLM的技术可能会由于内存间接寻址和非连续块内存带来的额外开销而降低性能。不过,如果vLLM的技术能应用于与大语言模型服务具有相似特性的其他工作负载,我们会感到十分欣喜。
  • 在应用虚拟内存和分页时针对大语言模型的优化:vLLM通过利用特定应用的语义,对虚拟内存和分页的思想进行了重新诠释和扩展。其中一个例子是vLLM的全有或全无的换出策略,该策略利用了处理一个请求需要将其所有相应的词元状态存储在GPU内存中的这一事实。另一个例子是用于恢复被逐出块的重新计算方法,这在操作系统中是不可行的。此外,vLLM通过将用于内存访问操作的GPU内核与注意力计算等其他操作的内核进行融合,减轻了分页中内存间接寻址的开销。

9. 相关工作

  • 通用模型服务系统:近年来,模型服务一直是一个活跃的研究领域,人们提出了许多系统来解决深度学习模型部署中的不同问题。Clipper[11]、TensorFlow Serving[33]、Nexus[45]、InferLine[10]和Clockwork[20]是一些早期的通用模型服务系统。它们研究了为单个或多个模型提供服务时的批处理、缓存、部署和调度问题。最近,DVABatch[12]引入了多入口多出口批处理。REEF[21]和Shepherd[61]提出了用于服务的抢占机制。AlpaServe[28]利用模型并行性进行统计复用。然而,这些通用系统没有考虑到大语言模型推理的自回归特性和词元状态,从而错过了优化的机会。
  • 针对Transformer的专用服务系统:由于Transformer架构的重要性,人们开发了许多针对它的专用服务系统。这些系统利用GPU内核优化[1, 29, 31, 56]、先进的批处理机制[14, 60]、模型并行性[1, 41, 60]和参数共享[64]来实现高效服务。其中,Orca[60]与我们的方法最为相关。
  • 与Orca的比较:Orca[60]中的迭代级调度和vLLM中的分页注意力是互补的技术:虽然这两个系统都旨在提高GPU利用率,进而提高大语言模型服务的吞吐量,但Orca是通过对请求进行调度和交错处理,以便更多请求可以并行处理来实现这一目标,而vLLM则是通过提高内存利用率,使更多请求的工作集能够装入内存来实现。通过减少内存碎片化并实现共享,vLLM能够在一个批次中并行运行更多请求,与Orca相比实现了2 - 4倍的加速。实际上,像Orca中那样对请求进行细粒度的调度和交错处理,使得内存管理更具挑战性,这也让vLLM中提出的技术显得更为关键。
  • 内存优化:加速器计算能力和内存容量之间不断扩大的差距,使得内存成为训练和推理的瓶颈。交换[23, 42, 55]、重新计算[7, 24]以及它们的组合[40]已被用于降低训练的峰值内存。值得注意的是,FlexGen[46]研究了如何在GPU内存有限的情况下,为大语言模型推理交换权重和词元状态,但它并非针对在线服务场景。OLLA[48]通过优化张量的生命周期和存储位置来减少碎片化,但它没有进行细粒度的块级管理或在线服务。FlashAttention[13]应用平铺和内核优化来降低注意力计算的峰值内存并减少I/O成本。本文在在线服务的背景下引入了块级内存管理的新思想。

10. 结论

本文提出了分页注意力(PagedAttention),这是一种新的注意力算法,它允许注意力的键和值存储在非连续的分页内存中,并介绍了vLLM,这是一个由分页注意力实现高效内存管理的高吞吐量大语言模型服务系统。受操作系统的启发,我们展示了如何调整已有的技术,如虚拟内存和写时复制,以有效地管理大语言模型服务中的键值缓存并处理各种解码算法。我们的实验表明,与最先进的系统相比,vLLM的吞吐量提高了2 - 4倍。

致谢

我们要感谢刘晓萱、陈智峰、黄彦平、匿名的SOSP审稿人以及我们的指导人周立东,感谢他们富有见地的反馈。这项研究部分得到了Andreessen Horowitz、Anyscale、Astronomer、谷歌、IBM、英特尔、Lacework、微软、穆罕默德·本·扎耶德人工智能大学、三星SDS、优步和VMware的资助。