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-𝑘最合适的翻译。Beam search 被广泛用于解码LLM最可能的输出序列，因为它降低了完全遍历样本空间的计算复杂度。该算法依赖于波束宽度参数𝑘，该参数决定了每一步保留的最佳候选数。在解码过程中，波束搜索通过考虑所有可能的标记来扩展波束中的每个候选序列，使用LLM计算它们各自的概率，并在候选序列(长度为$k\cdot |V|$)中保留最可能的𝑘个序列，其中 $|V|$ 是词汇表大小。

与并行解码不同，波束搜索工具不仅共享初始提示块，还共享不同候选块，并且共享模式随着解码过程的推进而动态变化，类似于复合分叉在操作系统中创建的进程树。图9显示了对于𝑘= 4的波束搜索示例，vLLM如何管理KV块。在用虚线表示的迭代之前，每个候选序列已经使用了4个完整的逻辑块。所有candidate共享第一个块0(即提示符)。candidate 3从第二部分开始离题。candidate0-2共用前3个block，并在第四个block分开。在随后的迭代中，前4个可能的候选项都来自候选项1和2。由于原来的候选0和3不再是最优候选，它们的逻辑块被释放，相应的物理块的引用计数被减少。vLLM释放所有引用计数达到0的物理块(block 2,4,5,8)，然后，vLLM分配新的物理块(block 9-12)来存储来自新候选对象的新KV缓存。现在，所有candidate共享0、1、3块;candidate0和1共享区块6，candidate2和3进一步共享block7。

以前的LLM服务系统需要在candidate上频繁地复制KV缓存。例如，在图9所示的情况下，在虚线之后，候选3将需要复制candidate2的KV缓存的大部分以继续生成。vLLM的物理块共享大大减少了这种频繁的内存复制开销。在vLLM中，不同candidate的大部分块可以共享。只有当新生成的令牌位于旧的共享块中时，才应用写时复制机制，就像并行解码一样。这只涉及复制一个数据块。

Shared prefix. 通常，LLM用户提供任务的(长)描述，包括指令和示例输入输出，也称为系统提示符[36]。描述与实际任务输入连接起来，形成请求提示。LLM生成的输出基于完整的提示符。图10显示了一个示例。此外，可以通过提示工程进一步调整共享前缀，以提高下游任务的准确性。

对于这种类型的应用程序，许多用户提示共享一个前缀，因此LLM服务提供商可以提前存储前缀的KV缓存，以减少在前缀上花费的冗余计算。在vLLM中，可以通过LLM服务提供者为一组预定义的共享前缀保留一组物理块来方便地实现这一点，正如操作系统如何跨进程处理共享库一样。带有共享前缀的用户输入提示符可以简单地将其逻辑块映射到缓存的物理块(最后一个块标记为写时复制)。提示阶段的计算只需要在用户的任务输入上执行。

Mixed decoding methods. 前面讨论的解码方法表现出不同的内存共享和访问模式。尽管如此，vLLM促进了具有不同解码偏好的请求的同时处理，这是现有系统无法有效做到的。这是因为vLLM通过将逻辑块转换为物理块的公共映射层隐藏了不同序列之间的复杂内存共享。LLM及其执行内核只看到每个序列的物理块id列表，不需要处理跨序列的共享模式。与现有系统相比，这种方法扩大了具有不同采样要求的请求的批处理机会，最终提高了系统的总体吞吐量。

4.5 Scheduling and Preemption

当请求流量超过系统容量时，vLLM必须优先考虑请求一个子集。在vLLM中，我们对所有请求采用先到先服务(FCFS)调度策略，确保公平性并防止饥饿。当vLLM需要抢占请求时，它确保首先服务最早到达的请求，并首先抢占最新的请求。

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

Swapping. 这是大多数虚拟内存实现使用的经典技术，它将被驱逐的页面复制到磁盘上的交换空间。在本例中，我们将被驱逐的块复制到CPU内存中。如图4所示，除了GPU块分配器之外，vLLM还包含一个CPU块分配器，用于管理交换到CPU RAM的物理块。当vLLM为新令牌耗尽空闲物理块时，它会选择一组序列来驱逐并将它们的KV缓存传输到CPU。一旦它抢占了一个序列并驱逐了它的块，vLLM就会停止接受新的请求，直到所有被抢占的序列都被完成。一旦请求完成，它的块就从内存中释放出来，并将被抢占序列的块带回来继续处理该序列。请注意，在这种设计中，交换到CPU RAM的块数量永远不会超过GPU RAM中的总物理块数量，因此CPU RAM上的交换空间受到分配给KV缓存的GPU内存的限制。

Recomputation. 在这种情况下，当被抢占的序列被重新调度时，我们只需重新计算KV缓存。请注意，重新计算延迟可以显著低于原始延迟，因为在解码时生成的token可以与原始用户提示连接为一个新的token----它们在所有位置的KV缓存可以在一个提示阶段迭代中生成。

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

4.6 Distributed Execution

许多llm的参数大小超过了单个GPU的容量[5,9]。因此，有必要将它们划分到分布式gpu上，并以模型并行的方式执行[28,63]。这需要能够处理分布式内存的内存管理器。vLLM通过支持 transformer 上广泛使用的Megatron-LM风格张量模型并行策略，在分布式设置中是有效的。

该策略遵循SPMD(单程序多数据)调度，其中线性层是分区的，以执行逐块矩阵乘法，gpu通过allreduce操作不断同步中间结果。具体来说，注意力 算子在注意头维度上被分割，每个SPMD过程在多头注意中负责注意头的一个子集。我们观察到，即使模型并行执行，每个模型分片仍然处理相同的一组输入令牌，因此需要KV缓存用于相同的位置。因此，vLLM在集中式调度器中具有单个KV缓存管理器，如图4所示。不同的GPU工作者共享管理器，以及从逻辑块到物理块的映射。这种公共映射允许GPU worker使用调度程序为每个输入请求提供的物理块来执行模型。虽然每个GPU worker都有相同的物理块id，但一个工作线程只为其相应的注意头存储一部分KV缓存。

在每个步骤中，调度器首先为批处理中的每个请求准备带有输入token id的消息，并为每个请求准备块表。接下来，调度器将这个控制消息广播给GPU worker。然后，GPU worker 开始使用输入token id执行模型。在注意层，GPU工作人员根据控制消息中的块表读取KV缓存。在执行过程中，GPU worker 中间结果与all-reduce通信原语同步，而不需要调度程序的协调。最后，GPU工作器将这次迭代的采样token发送回调度器。总之，GPU工作人员不需要同步内存管理，因为他们只需要在每次解码迭代开始时接收所有内存管理信息以及step输入。

5. Implementation

vLLM是一个端到端服务系统，采用FastAPI[15]前端和基于gpu的推理引擎。前端扩展了OpenAI API[34]接口，允许用户为每个请求定制采样参数，如最大序列长度和beam width 𝑘。vLLM引擎是用8.5k行Python和2K行c++ /CUDA代码编写的。我们在Python中开发控制相关组件，包括调度器和块管理器，同时为关键操作(如PagedAttention)开发自定义CUDA内核。对于模型执行器，我们使用PyTorch和Transformer实现流行的llm，如GPT[5]、OPT[62]和LLaMA [52]我们使用NCCL[32]在分布式GPU worker之间进行张量通信。

5.1 Kernel-level Optimization

由于PagedAttention引入了现有系统无法有效支持的内存访问模式，我们开发了几个GPU内核来优化它。

(1)Fused reshape and block write. 在每个Transformer层中，新的KV缓存被分割成块，重新塑造为块读取优化的内存布局，然后保存在块表指定的位置。为了最小化内核启动开销，我们将它们融合到一个内核中。

(2)Fusing block read and attention.我们采用FasterTransformer[31]中的注意力内核，根据块表读取KV缓存，并动态执行注意力操作。为了确保合并内存访问，我们分配了一个GPU warp来读取每个块。此外，我们还增加了对请求批处理中可变序列长度的支持。

(3)Fused block copy.由写时拷贝机制发出的块拷贝操作可以在不连续的块上操作。如果我们使用cudamempyasync API，这可能导致大量的小数据移动调用。为了减少开销，我们实现了一个内核，它将不同块的复制操作批处理到单个内核启动中。

5.2 Supporting Various Decoding Algrithms

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