【论文阅读】Efficient Memory Management for Large Language Model Serving with PagedAttention_efficientmemory management for large language mode

Efficient Memory Management for Large Language Model Serving with PagedAttention

• 出处: ACM-2023 [2309.06180] Efficient Memory Management for Large Language Model Serving with PagedAttention (arxiv.org)

• 源码: 项目地址

Abstract

高吞吐量的LLM服务需同时处理多个请求。但是现有系统非常困难，因为KV cache非常巨大并且是动态伸缩的，因为显存管理不善，导致碎片和重复，造成显存的巨大浪费，从而限制了batch的大小和吞吐量。为了解决这个问题，本文借鉴操作系统的分页内存管理方法，提出了PagedAttention。基于这个方法，实现了vLLM，它能够实现：1) 接近零的KV cache浪费；2) 同一请求内和不同请求间KV cache的灵活共享。实验证明本方法的吞吐量是SOTA系统的2-4倍。

1 Introduction

像GPT和PaLM这样的大型语言模型(llm)的出现使编程助理和通用聊天机器人等新应用成为可能，它们开始深刻地影响我们的工作和日常生活。许多云计算公司[34,44]正在竞相提供这些应用程序作为托管服务。然而，运行这些应用程序是非常昂贵的，需要大量的硬件加速器，如gpu。

llm的核心是自回归的Transformer模型。该模型基于输入(提示)和迄今为止生成的输出标记的前一个序列，一次生成一个单词(token)。对于每个请求，重复这个昂贵的过程，直到模型输出一个终止令牌。这种顺序生成过程使工作负载受到内存限制，使gpu的计算能力得不到充分利用，并限制了服务吞吐量。

通过将多个请求批处理在一起，可以提高吞吐量。但是，为了批量处理许多请求，应该有效地管理每个请求的内存空间。例如，图1(左)说明了在具有40GB RAM的NVIDIA A100 GPU上13B参数LLM的内存分布。大约65%的内存分配给模型参数，这些权重在服务期间保持静态。接近30%的内存用于存储请求的动态状态。对于transformer，这些状态由与注意力机制相关的key值和value值组成，通常被称为KV cache，它表示来自早期令牌的上下文，以按顺序生成新的输出令牌。剩下的一小部分用于其他数据，包括激活——在评估LLM时产生的短暂张量。由于模型权重是恒定的，并且激活仅占用GPU内存的一小部分，因此管理KV缓存的方式对于确定最大批大小至关重要。当管理效率低下时，KV高速缓存会显著限制批处理大小，从而限制LLM的吞吐量，如图1(右)所示。

在本文中，我们观察到现有的LLM服务系统无法有效地管理KV缓存。这主要是因为它们将请求的KV缓存存储在连续内存空间中，因为大多数深度学习框架要求将张量存储在连续内存中。然而，与传统深度学习工作负载中的张量不同，KV cache 具有独特的特征:随着模型生成新的token，它会随着时间的推移动态增长和缩小，并且它的生命周期和长度是未知的。这些特点使现有系统的方法在两个方面显着效率低下：

首先，现有系统存在内部和外部内存碎片。为了在连续空间中存储请求的KV cache，它们预先分配了一个具有请求最大长度的连续内存块(例如，2048个token)。这可能导致严重的内部碎片，因为请求的实际长度可能比它的最大长度短得多(例如，图11)。此外，即使预先知道实际长度，预分配仍然是低效的:由于在请求的生命周期内保留了整个块，其他较短的请求不能利用当前未使用的块的任何部分。此外，外部内存碎片也很重要，因为每个请求的预分配大小可能不同。事实上，我们在图2中的分析结果显示，在现有系统中，只有20.4% - 38.2%的KV cache内存用于存储实际 token 状态。

其次，现有系统无法利用内存共享的机会。LLM服务通常使用高级解码算法，例如parallel sampling 和 beam search，每个请求生成多个输出。在这些场景中，请求由多个序列组成，这些序列可以部分共享它们的KV cache。然而，在现有系统中，内存共享是不可能的，因为序列的KV cache存储在单独的连续空间中。

为了解决上述限制，我们提出了PagedAttention，这是一种注意力算法，灵感来自于操作系统(OS)对内存碎片和共享的解决方案:带分页的虚拟内存。PagedAttention将请求的KV缓存划分为块，每个块可以包含固定数量的token的key和value。在PagedAttention中，KV cache的块不一定存储在连续的空间中。因此，我们可以像在OS的虚拟内存中那样更灵活地管理KV cache:可以将块视为页，token视为字节，request视为进程。这种设计通过使用相对较小的块并按需分配来减轻内部碎片。此外，它消除了外部碎片，因为所有块都具有相同的大小。最后，它支持以块粒度、跨与相同请求关联的不同序列甚至跨不同请求共享内存。

在这项工作中，我们在PagedAttention的基础上构建了一个高吞吐量的分布式LLM服务引擎vLLM，它在KV高速缓存中实现了接近零的浪费。vLLM使用与PagedAttention共同设计的块级内存管理和抢占式请求调度。vLLM支持GPT[5]、OPT[62]、LLaMA[52]等流行的llm，支持不同大小的llm，包括超过单个GPU内存容量的llm。我们对各种模型和工作负载的评估表明，与最先进的系统相比，vLLM将LLM服务吞吐量提高了2-4倍，而完全不影响模型的准确性。对于更长的序列、更大的模型和更复杂的解码算法，改进更加明显。

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

• 我们确定了在服务llm时内存分配方面的挑战，并量化了它们对服务性能的影响。
• 受操作系统中虚拟内存和分页的启发，提出了一种基于非连续分页内存中KV cache 的注意力算法PagedAttention。
• 我们设计并实现了vLLM，一个基于PagedAttention的分布式LLM服务引擎。
• 我们在各种情况下评估了vLLM，并证明它大大优于以前的核心解决方案，如FasterTransformer[31]和Orca[60]。

2. Background

在本节中，我们描述了典型LLM的生成和服务过程，以及LLM服务中使用的iteration 调度。

2.1 Transformer-Based Large Language Models

语言建模的任务是对标记列表 $(x_1,\dots ,x_n)$ 的概率进行建模。由于语言具有自然的顺序排序，因此通常将整个序列的联合概率分解为条件概率的乘积(也称为自回归分解[3]):
$$P(x)=P(x_1)\cdot P(x_2|x_1)\cdot P(x_n|x_1,\cdots ,x_{n-1})$$
Transformer[53]已经成为在大范围内对上述概率进行建模的事实上的标准架构。基于transformer的语言模型最重要的组件是它的自注意力层。对于输入隐藏层状态 $(x_1,\cdots ,x_n)\in {\mathbb{R}}^{n\times d}$ ，自注意力层首先对每个位置向量进行线性变换，得到Query, key 和 value 向量:
$$q_i=W_q{x}_i,k_i=W_k{x}_i,v_i=W_v{x}_i$$
然后，自注意层通过将某一位置的 Query 向量与其之前的所有 key 向量相乘来计算关注分数 $a_{ij}$，并计算输出 $o_i$ 作为 value 向量的加权平均值:
$$a_{ij}=\frac{exp(q_i^T{k}_j/\sqrt{d})}{{\sum }_{t=1}^{i}exp(q_i^T{k}_t/\sqrt{d})},o_i=\sum _{j=1}^i{a}_{ij}{v}_j.$$
除Eq. 4中的计算外，Transformer模型中的所有其他组成部分，包括嵌入层、前馈层、层归一化[2]、剩余连接[22]、输出logit计算以及Eq. 2中的Query, key, value 值转换，都是按照位置独立应用的，形式为: $y_i=f(x_i)$。

2.2 LLM Service & Autoregressive Generation

经过训练后，llm通常被部署为生成服务(例如，生成API[34]或聊天机器人[19,35])。对LLM服务的请求提供了一个输入提示token序列$(x_1,\cdots ,x_n)$，LLM服务根据式(1)生成一个输出令牌列表$(x_{n+1},\cdots ,xn+T)$。我们将提示符列表和输出列表的连接称为序列。

LLM只能逐个采样并生成新的令牌，并且每个新token的生成过程取决于该序列中所有先前的tonken，特别是它们的key和value。在这个顺序生成过程中，通常缓存现有token的key和value向量，以生成未来的token，称为KV cache。注意，一个token的KV cache依赖于它之前的所有token。这意味着同一token在序列中出现在不同位置的KV cache 将是不同的。

给定一个请求prompt，LLM服务中的生成计算可以分解为两个阶段:

**The prompt phase: ** 将用户的整个prompt $(x_1,\cdots ,x_n)$ 作为输入计算第一个token的概率$P(x_{n+1}|x_1,\cdots ,x_n)$ ，在此过程中，还生成了key向量$k_1,\cdots ,k_n$ 和 value 向量$v_1,\cdots ,v_n$，由于token $(x_1,\cdots ,x_n)$ 都是已知的，提示阶段的计算可以使用矩阵乘法运算并行化。因此，这一阶段可以有效地利用gpu固有的并行性。

**The autoregressive generation phase: ** 依次生成剩余的新token，在第 t 个迭代，模型需要token $x_{n+t}$ 作为输入并使用key向量 $k_1,\cdots ,k_{n+t}$ 和value 向量 $v_1,\cdots ,v_{n+t}$ 计算概率 $P(x_{n+t+1}|x_1,\cdots ,x_{n+t})$ ,注意，$1\to n+t-1$ 的key和value向量在之前的迭代已经缓存了，这个迭代值生成新的token的key 和 value 。当序列达到最大长度(由用户或者llm限制)或发出序列结束($<eos>$)token时，自回归阶段完成。由于数据依赖性，不同迭代的计算不能并行化，通常采用矩阵-向量乘法，效率较低。因此，这一阶段严重未充分利用GPU计算并成为内存限制，带来单个请求的大部分延迟。

2.3 Batching Techniques for LLMs

通过批量处理多个请求，可以提高llm服务的计算利用率。由于请求共享相同的模型权重，因此移动权重的开销在批处理请求中平摊，并且当批处理大小足够大时，可能会被计算开销所抵消(相比变得很小)。但是，由于两个原因，将请求批处理到LLM服务是非常重要的。首先，请求可能在不同的时间到达。直接的批处理策略要么让较早的请求等待较晚的请求，要么将之后的请求延迟到较早的请求完成，从而导致严重的排队延迟。其次，请求可能具有不同的输入和输出长度(图11)。直接的批处理技术填充请求的输入和输出，以平衡它们的长度，但会浪费GPU计算和内存。

为了解决这个问题，人们提出了细粒度的批处理机制，如cellular batching[16]和iteration-level scheduling[60]。与在请求级别工作的传统方法不同，这些技术在迭代级别操作。在每次迭代之后，完成的请求将从批处理中删除，并添加新的请求。因此，可以在等待单个迭代后处理新请求，而不是等待整个批处理完成。

此外，使用特殊的GPU内核，这些技术消除了填充输入和输出的需要。通过减少排队延迟和填充带来的低效率，细粒度批处理机制显著提高了LLM服务的吞吐量。

3. Memory Callenge in LLM Serving

尽管细粒度批处理减少了计算浪费，并使请求能够以更灵活的方式进行批处理，但可以批处理的请求数量仍然受到GPU内存容量的限制，特别是分配给存储KV cache 的空间。换句话说，服务系统的吞吐量受内存限制。克服这种内存限制需要解决内存管理中的以下挑战:

Large KV cache 。KV缓存大小随着请求数量的增加而快速增长。例如，对于13B参数的OPT模型[62]，单个令牌的KV缓存需要800 KB的空间，计算为2(key 和 value)× 5120(hidden state size)× 40(number of layers)× 2(bytes per FP16)。由于OPT可以生成多达2048个token的序列，因此存储一个请求的KV cache所需的内存可能高达1.6 GB。并发GPU的内存容量为几十GB。即使将所有可用内存分配给KV cache，也只能容纳几十个请求。此外，低效的内存管理会进一步减小批处理大小，如图2所示。此外，从目前的趋势来看，GPU的计算速度的增长速度超过了内存容量的增长速度。例如，从NVIDIA A100到H100, FLOPS提高了2倍以上，但GPU内存最大保持在80GB。因此，我们相信内存将成为越来越重要的瓶颈。

Complex decoding algorithms. LLM服务提供了一系列解码算法供用户选择，每种算法对内存管理复杂性的影响各不相同。例如，当用户从单个输入提示请求多个随机样本(程序建议中的典型用例)时，提示部分的KV cache 可以共享，在我们的实验中(§6.3)，它占总KV缓存的12%，以最小化内存使用。另一方面，在自回归生成阶段，由于不同的样本结果及其对环境和位置的依赖，KV cache 应该保持不共享。KV cache 共享的程度取决于所采用的具体解码算法。 在beam search [49]等更复杂的算法中，不同的请求可以共享其KV cache 的更大部分(高达55%的内存节省，参见§6.3)，并且共享模式随着解码过程的推进而变化。

Scheduling for unknown input & output lengths. LLM服务的请求在其输入和输出长度方面表现出可变性。这就要求内存管理系统能够适应各种提示长度。此外，随着解码时请求的输出长度增加，其KV cache 所需的内存也会扩展，并且可能耗尽用于传入请求或正在生成的现有提示的内存。系统需要做出调度决策，例如从GPU内存中删除或交换某些请求的KV缓存。

3.1 Memory Management in Existing Systems

由于当前深度学习框架中的大多数运算符要求将张量存储在连续内存中，以前的LLM服务系统也将一个请求的KV缓存存储为跨不同位置的连续张量。由于LLM的输出长度不可预测，因此它们根据请求的最大可能序列长度静态地为请求分配一块内存，而不考虑请求的实际输入或最终输出长度。

图3显示了两个请求:请求A的最大可能序列长度为2048，请求B的最大可能序列长度为512。现有系统中的块预分配方案有三个主要的内存浪费来源:为未来token 保留的内存、由于过度供应潜在的最大序列长度而导致的内部碎片，以及来自内存分配器(如buddy分配器)的外部碎片。外部碎片永远不会用于生成的令牌，这在服务请求之前是已知的。内部碎片也未被使用，但这只有在请求完成采样后才会实现。它们都是纯粹的内存浪费。虽然保留的内存最终会被使用，但是在整个请求期间保留这个空间，特别是当保留的空间很大时，会占用本来可以用来处理其他请求的空间。我们在图2中可视化了我们的实验中内存浪费的平均百分比，揭示了以前系统中的实际有效内存可以低至20.4%。

虽然compaction[54]已经被提出作为一种潜在的碎片解决方案，但由于大量KV缓存，在性能敏感的LLM服务系统中执行压缩是不切实际的。即使使用了压缩，为每个请求预先分配的块空间也会阻止现有内存管理系统中特定于解码算法的内存共享。

4. Method

在这项工作中，我们开发了一种新的注意力算法PagedAttention，并构建了一个LLM服务引擎vLLM，以解决§3中概述的挑战。vLLM的架构如图4所示。vLLM采用集中式调度程序来协调分布式GPU工作线程的执行。KV cache 管理器以分页方式有效地管理KV cache，由PagedAttention启用。具体来说，KV cache 管理器通过集中调度程序发送的指令来管理GPU工作线程上的物理KV cache。

接下来，我们在§4.1中描述PagedAttention算法。我们分别在§4.2中展示KV cache 管理器的设计以及它如何促进§4.3中的PagedAttention。然后，我们展示了这种设计如何促进各种解码方法(§4.4)的有效内存管理，并处理可变长度的输入和输出序列(§4.5)。最后，我们展示了vLLM的系统设计如何在分布式设置中工作(第4.6节)。

4.1 PageAttention

为了解决§3中的内存挑战，我们引入了PagedAttention，这是一种受操作系统中分页的经典思想启发的注意力算法[25]。与传统的注意力算法不同，PagedAttention允许在非连续的内存空间中存储连续的键和值。具体来说，PagedAttention将每个序列的KV cache 划分为KV block。每个块包含固定数量的token的key和value向量，我们将其记为KV block size。记key block为 $K_j=(k_{(j-1)B+1},\cdots ,k_{jB})$ ，记value block为$V_j=(v_{(j-1)B+1},\cdots ,v_{jB})$ ，注意力可以逐块计算为:
$$A_{ij}=\frac{exp(q_i^T{K}_j/\sqrt{d})}{{\sum }_{t=1}^{\lceil /B\rceil }exp(q_i^T{K}_t/\sqrt{d})},o_i=\sum _{j=1}^{\lceil i/B\rceil }{V}_j{A}_{ij}^T$$
其中$A_{ij}=(a_{i,(j-1)B+1},\cdots ,a_{i,j}B)$表示第j个KV block中的行向量。在注意力计算过程中，PagedAttention内核分别识别和提取不同的KV块，我们在图5中展示了一个PagedAttention的例子:键和值向量分布在三个块上，并且这三个块在物理内存上不是连续的。每次，内核将查询tokend 的Query向量和block中的key向量$K_j$ (例如，block 0中的key向量"Four score and seven")相乘，计算出注意力分数 $A_{ij}$， 然后将变量 $A_{ij}$ 和block中的value向量$V_j$相乘得到最终的注意力分数。

总之，PagedAttention算法允许将KV块存储在非连续的物理内存中，从而在vLLM中实现更灵活的分页内存管理。

4.2 KV Cache Manager

vLLM内存管理器背后的关键思想类似于操作系统中的虚拟内存[25]。操作系统将内存划分为固定大小的页面，并将用户程序的逻辑页面映射到物理页面。连续的逻辑页可以对应于非连续的物理内存页，允许用户程序访问内存，就好像它是连续的一样。此外，物理内存空间不需要提前完全预留，使操作系统可以根据需要动态分配物理页面。vLLM使用虚拟内存背后的思想来管理LLM服务中的KV缓存。通过PagedAttention，我们将KV缓存组织为固定大小的KV块，就像虚拟内存中的页面一样。

请求的KV cache 表示为一系列逻辑KV块，从左到右填充为生成的新token及其KV cache。最后一个KV区块的未填充位置保留给未来生成的token使用。在GPU节点上，块引擎分配一个连续的GPU DRAM块，并将其划分为物理KV块(这也在CPU RAM上完成，用于交换 ,§4.5)。KV块管理器还维护块表——每个请求的逻辑和物理KV块之间的映射。每个块表项记录一个逻辑块对应的物理块和填充位置的数量。分离逻辑和物理KV块允许vLLM动态增长KV缓存，而无需提前为所有位置保留它，这消除了现有系统中的大部分内存浪费，如图2所示。

4.3 Decoding with PagedAttention and vLLM

接下来，我们通过一个示例，如图6所示，来演示vLLM如何在单个输入序列的解码过程中执行PagedAttention并管理内存:

①就像在OS的虚拟内存中一样，vLLM不需要为最初可能生成的最大序列长度保留内存。相反，它只保留必要的KV块，以容纳在提示计算期间生成的KV缓存。在本例中，提示符有7个令牌，因此vLLM将前2个逻辑KV块(0和1)映射到2个物理KV块(分别为7和1)。在预填充步骤中，vLLM使用传统的自关注算法(例如[13])生成提示符和第一个输出令牌的KV缓存。然后，vLLM将前4个令牌的KV缓存存储在逻辑块0中，并将随后的3个令牌存储在逻辑块1中。剩余的槽保留给后续的自回归生成阶段。

②在第一个自回归解码步骤中，vLLM使用物理块7和1上的PagedAttention算法生成新的令牌。由于在最后一个逻辑块中仍然有一个槽可用，因此新生成的KV缓存存储在那里，并且块表的#filled记录被更新。

③在第二步解码时，由于最后一个逻辑块已满，vLLM将新生成的KV缓存存储在新的逻辑块中;vLLM为它分配一个新的物理块(物理块3)，并将这个映射存储在块表中。

全局而言，对于每次解码迭代，vLLM首先选择一组候选序列进行批处理(参见§4.5)，并为新需要的逻辑块分配物理块。然后,vLLM连接所有当前迭代的输入token作为一个序列，并将其输入LLM。在LLM的计算过程中，vLLM使用PagedAttention内核访问之前以逻辑KV块形式存储的KV缓存，并将新生成的KV缓存保存到物理KV块中。在KV块中存储多个令牌(块大小> 1)使PagedAttention内核能够跨多个位置并行处理KV缓存，从而增加硬件利用率并减少延迟。然而，更大的块大小也会增加内存碎片。我们在§7.2中研究了块大小的影响。

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

4.4 Application to Other Decoding Scenarios

§4.3展示了PagedAttention和vLLM如何处理基本的解码算法，例如贪婪解码和采样，将一个用户提示作为输入并生成单个输出序列。在许多成功的LLM应用程序中[18,34]，LLM服务必须提供更复杂的解码场景，表现出复杂的访问模式和更多的内存共享机会。在本节中，我们将展示vLLM对它们的一般适用性。

Parallel sampling. 并行采样。在基于LLM的程序助手中[6,18]，LLM为单个输入提示生成多个采样输出;用户可以从各种候选输出中选择自己喜欢的输出。到目前为止，我们已经隐含地假设了一个请求生成单个序列。在本文的其余部分中，我们假设一个请求生成多个序列的更一般的情况。在并行采样中，一个请求包含多个共享相同输入提示的样本，从而允许共享提示的KV缓存。通过它的PagedAttention和分页内存管理，vLLM可以很容易地实现这种共享并节省内存。在本文的其余部分中，我们假设一个请求生成多个序列的更一般的情况。在并行采样中，一个请求包含多个共享相同输入提示的样本，从而允许共享提示的KV缓存。通过它的PagedAttention和分页内存管理，vLLM可以很容易地实现这种共享并节省内存。

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

总之，**vLLM支持跨多个输出样本共享用于存储提示的KV缓存的大部分空间，但最后的逻辑块除外，该逻辑块由写时复制机制管理。**通过跨多个示例共享物理块，可以大大减少内存使用，特别是对于长输入提示。

Beam search. 在机器翻译等LLM任务中[59]，用户期望LLM输出的top-