【论文阅读】Efficient Memory Management for Large Language Model Serving with PagedAttention

Abstract

  大语言模型（LLMs）的高吞吐量服务需要一次批处理足够多的请求。然而，现有的系统面临挑战，因为每个请求的键值缓存（KV cache）占用内存非常庞大，并且会动态增长和缩小。当管理效率低下时，这部分内存会因碎片化和冗余重复被明显浪费，从而限制了批处理的大小。为了解决这个问题，作者提出了PagedAttention，一种受操作系统中经典虚拟内存和分页技术启发的注意力算法。在此基础上，作者构建了vLLM，一个能够实现以下目标的 LLM 服务系统：（1）KV缓存内存的近零浪费；（2）在请求内部和跨请求之间灵活共享KV缓存，以进一步减少内存使用。作者的评估显示，与目前最先进的系统（如FasterTransformer和Orca）相比，vLLM在相同延迟水平下将热门LLM的吞吐量提高了2-4倍。对于更长的序列、更大的模型和更复杂的解码算法，这种改进更为显著。vLLM的源代码已公开在https://github.com/vllm-project/vllm。

1 Introduction

  随着GPT [5, 37]和PaLM [9]等大语言模型（LLMs）的出现，新的应用如编程助手 [6, 18] 和通用聊天机器人 [19, 35] 开始深刻影响着我们的工作和日常生活。许多云服务公司 [34, 44] 正在竞争以托管服务来提供这些应用。然而，运行这些应用非常昂贵，需要大量的硬件加速器（如GPU）。根据最新的估算，处理一个LLM请求的成本可能是传统关键词查询的10倍 [43]。鉴于这些高成本，增加LLM服务系统的吞吐量——从而降低每个请求的成本——变得越来越重要。
  LLM的核心是一个自回归Transformer模型 [53]。该模型基于输入（提示）和其已生成的输出序列来逐词（token）生成。对于每个请求，这一昂贵的过程会重复进行，直到模型输出一个终止token。这种顺序生成过程使得工作负载受限于内存，未能充分利用GPU的计算能力，从而限制了服务吞吐量。
  通过将多个请求批量处理，可以提高吞吐量。然而，为了在一个批次中处理多个请求，需要高效管理每个请求的内存空间。例如，图1（左）展示了在配备40GB内存的NVIDIA A100 GPU上，一个13B参数LLM的内存分布情况。大约65%的内存分配给模型权重，这些权重在服务过程中保持不变。接近30%的内存用于存储请求的动态状态。对于Transformer模型，这些状态由与注意力机制相关的键（key）和值（value）张量组成，通常称为KV缓存 [41]，它们表示从先前token中提取的上下文，用于按顺序生成新的输出token。剩余的一小部分内存用于其他数据，包括激活值——在评估LLM时创建的临时张量。由于模型权重是固定的，而激活值只占GPU内存的一小部分，因此KV缓存的管理方式对确定最大批处理大小至关重要。当管理效率低下时，KV缓存内存可能会显著限制批处理大小，从而限制LLM的吞吐量，如图1（右）所示。

  在本文中，作者观察到现有的LLM服务系统 [31, 60] 在高效管理KV缓存内存方面存在不足。这主要是因为它们将请求的KV缓存存储在连续的内存空间中，由于大多数深度学习框架 [33, 39] 要求张量存储在连续内存中。然而，与传统深度学习工作负载中的张量不同，KV缓存具有独特的特性：它会随着模型生成新的token而动态增长和缩小，并且其生命周期和长度无法预先确定。这些特性使得现有系统的方法在以下两个方面显著低效：
  首先，现有系统 [31, 60] 存在内部和外部内存碎片问题。为了将请求的KV缓存存储在连续空间中，它们会预先分配一个具有请求最大长度（例如2048个token）的连续内存块。这可能导致严重的内部碎片，因为请求的实际长度可能远小于其最大长度（例如图11）。此外，即使预先知道实际长度，预分配仍然是低效的：由于整个内存块在请求的生命周期内都被保留，其他较短的请求无法利用当前未使用的任何部分。此外，外部内存碎片也可能非常显著，因为每个请求的预分配大小可能不同。事实上，作者在图2中的分析结果显示，现有系统中只有20.4% - 38.2%的KV缓存内存用于存储实际的token状态。

  其次，现有系统无法利用内存共享的机会。LLM服务通常使用高级解码算法，例如并行采样和束搜索，这些算法会为每个请求生成多个输出。在这些场景中，请求由多个序列组成，这些序列可以部分共享它们的KV缓存。然而，在现有系统中不可能进行内存共享，因为序列的KV缓存存储在单独的连续空间中。
  为了解决上述局限性，作者提出了PagedAttention，这是一种受操作系统（OS）解决内存碎片和共享问题启发的注意力算法：基于分页的虚拟内存。PagedAttention将请求的KV缓存划分为块，每个块可以包含固定数量token的注意力键和值。在PagedAttention中，KV缓存的块不必存储在连续空间中。因此，作者可以像操作系统的虚拟内存一样更灵活地管理KV缓存：可以将块视为页面，将token视为字节，将请求视为进程。这种设计通过使用相对较小的块并按需分配来缓解内部碎片问题。此外，由于所有块的大小相同，它消除了外部碎片。最后，它支持在块粒度上进行内存共享，无论是在同一请求的不同序列之间，甚至是在不同请求之间。
  在这项工作中，作者构建了vLLM，一种基于PagedAttention的高吞吐量的分布式LLM服务引擎，实现了KV缓存内存的近零浪费。vLLM采用了与PagedAttention共同设计的块级内存管理和抢占式请求调度。vLLM支持各种规模的热门LLM，如GPT [5]、OPT [62] 和LLaMA [52]，包括那些超过单个GPU内存容量的模型。作者在各种模型和工作负载上的评估表明，与现有最先进的系统 [31, 60] 相比，vLLM将LLM服务吞吐量提高了2-4倍，且完全不影响模型准确性。对于更长的序列、更大的模型和更复杂的解码算法，这种改进更为显著（§4.3）。总的来说，作者做出了以下贡献：
• 作者发现了在服务LLM时内存分配的挑战，并量化了它们对服务性能的影响。
• 作者提出了PagedAttention，这是一种将KV缓存存储在非连续分页内存中的注意力算法，其受操作系统中虚拟内存和分页技术的启发。
• 作者设计并实现了vLLM，这是一种基于PagedAttention构建的分布式LLM服务引擎。
• 作者在各种场景下评估了vLLM，并证明其显著优于之前的最先进解决方案，如FasterTransformer [31] 和 Orca [60]。

2 Background

在本节中，作者描述一般LLM的生成和服务过程，以及LLM服务中使用的迭代级调度。

2.1 Transformer-Based Large Language Models

  语言建模的任务是对一系列token $(x_1,...,x_n)$ 的概率进行建模。由于语言具有自然的顺序性，通常将整个序列的联合概率分解为条件概率的乘积（也称为自回归分解 [3]）:
$$P(x)=P(x_1)\cdot P(x_2|x_1)\cdot \cdot \cdot P(x_n|x_1,...,x_{n-1})\text{\hspace{1em}(1)}$$

  Transformer [53] 已成为大规模建模上述概率的事实标准架构。基于Transformer的语言模型中最重要的组件是其自注意力层。对于输入隐藏状态序列$(x_1，...，x_n)\in {\mathbb{R}}^{n\times d}$，自注意力层首先对每个位置 $i$ 应用线性变换，以获得查询（query）、键（key）和值（value）向量：
$$q_i=W_q{x}_i,k_i=W_k{x}_i,v_i=W_v{x}_i\text{\hspace{1em}(2)}$$

  然后，自注意力层通过将某个位置的查询向量与其之前所有位置的键向量相乘来计算注意力分数$a_{ij}$，并通过值向量的加权平均来计算输出$o_i$：
$$a_{ij}=\frac{exp(q_i^{\top }{k}_j/\sqrt{d})}{{\sum }_{t=1}^{i}exp(q_i^{\top }{k}_t/\sqrt{d})},o_i=\sum _{j=1}^i{a}_{ij}{v}_j\text{\hspace{1em}(3)}$$

  除了公式4中的计算外，Transformer模型中的所有其他组件，包括嵌入层、前馈层、层归一化 [2]、残差连接 [22]、输出logit计算以及公式2中的查询、键和值变换，都是以$y_i=f(x_i)$ 的形式独立应用于每个位置的。

2.2 LLM Service & Autoregressive Generation

  训练完成后，LLM通常被部署为有条件的生成服务（例如补全API [34] 或聊天机器人 [19, 35]）。对LLM服务的请求提供一系列输入提示token $(x_1,...,x_n)$，LLM服务根据公式1生成一系列输出token $(x_{n+1},...,x_{n+T})$。作者将提示与输出列表的连接称为序列。
  由于公式1中的分解，LLM只能逐个采样和生成新token，并且每个新token的生成过程都依赖于该序列中所有先前的token，特别是它们的键和值向量。在这个顺序生成过程中，现有token的键和值向量通常会被缓存以生成未来的token，这被称为KV缓存。需要注意的是，一个token的KV缓存依赖于其之前的所有token。这意味着出现在序列中不同位置的同一token的KV缓存将是不同的。
  给定一个请求提示，LLM服务中的生成计算可以分解为两个阶段：
  提示阶段将整个用户提示 $(x_1,...,x_n)$ 作为输入，并计算第一个新token的概率 $P(x_{n+1}|x_1,...,x_n)$。在这个过程中，还会生成键向量 $k_1,...,k_n$ 和值向量 $v_1,...,v_n$。由于提示token $x_1,...,x_n$ 都是已知的，提示阶段的计算可以使用矩阵乘法操作并行化。因此，这一阶段可以高效利用GPU固有的并行性。
  自回归生成阶段按顺序生成剩余的新token。在第 $t$ 次迭代中，模型将一个token $x_{n+t}$ 作为输入，并使用键向量 $k_1,...,k_{n+t}$ 和值向量 $v_1,...,v_{n+t}$ 计算概率 $P(x_{n+t+1}|x_1,...,x_{n+t})$。值得注意的是，位置 $1$ 到 $n+t-1$ 的键和值向量在先前的迭代中已被缓存，只有新的键和值向量 $k_{n+t}$ 和 $v_{n+t}$ 在此次迭代中计算。当序列达到最大长度（由用户指定或受LLM限制）或生成序列结束token（$<eos>$）时，此阶段完成。由于数据依赖性，不同迭代的计算无法并行化，并且通常使用矩阵-向量乘法，效率较低。因此，这一阶段严重未充分利用GPU计算能力，成为内存瓶颈，占据了单个请求延迟的大部分时间。

2.3 Batching Techniques for LLMs

  通过将多个请求批量处理，可以提高服务LLM的计算利用率。由于请求共享相同的模型权重，移动权重的开销在一个批次中的请求之间分摊，并且当批次大小足够大时，这一开销会被计算开销所覆盖。然而，将请求批量处理到LLM服务中并不简单，原因有两个。首先，请求可能在不同的时间到达。一种简单的批处理策略要么让较早的请求等待较晚的请求，要么延迟传入的请求直到较早的请求完成，从而导致显著的排队延迟。其次，请求可能具有截然不同的输入和输出长度（图11）。一种直接的批处理技术是将请求的输入和输出填充到相同的长度，这会浪费GPU计算和内存。
  为了解决这个问题，已经提出了细粒度的批处理机制，例如蜂窝批处理 [16] 和迭代级调度 [60]。与传统的按请求级别工作的方式不同，这些技术是在迭代级别上进行操作的。每次迭代后，完成的请求从批次中移除，并添加新的请求。因此，新请求可以在等待一次迭代后被处理，而无需等待整个批次完成。此外，通过使用特殊的GPU内核，这些技术消除了填充输入和输出的需求。通过减少排队延迟和填充带来的低效性，细粒度的批处理机制显著提高了LLM服务的吞吐量。

3 Memory Challenges in LLM Serving

  尽管细粒度的批处理减少了计算浪费，并使请求以更灵活的方式进行批量处理，但可以一起批处理的请求数量仍然受到GPU内存容量的限制，特别是分配用于存储KV缓存的空间。换句话说，服务系统的吞吐量受限于内存。要克服这一内存限制，需要解决内存管理中的以下挑战：
  庞大的KV缓存。KV缓存的大小随着请求数量的增加而迅速增长。例如，对于13B参数的OPT模型 [62]，单个token的KV缓存需要800 KB的空间，计算方式为 2（键和值向量） × 5120（隐藏状态大小） × 40（层数） × 2（FP16每个字节的位数）。由于OPT可以生成最多2048个token的序列，存储一个请求的KV缓存所需的内存可能高达1.6 GB。当前GPU具有几十GB的内存容量。即使将所有可用内存都分配给KV缓存，也只能容纳几十个请求。此外，效率低下的内存管理可以进一步降低批处理大小，如图2所示。另外，根据当前的趋势，GPU的计算速度增长快于内存容量 [17]。例如，从NVIDIA A100到H100，FLOPS增加了2倍以上，但GPU内存最大保持在80GB。因此，作者认为内存将成为一个越来越重要的瓶颈。
  复杂的解码算法。LLM服务提供了一系列解码算法供用户选择，每种算法对内存管理复杂度的影响各不相同。例如，当用户从单个输入提示中请求多个随机样本时（程序建议中的一般用例 [18]），提示部分的KV缓存（在作者的实验(§6.3)中占总KV缓存内存的12%）可以共享以最小化内存使用。另一方面，由于样本结果不同及其对上下文和位置的依赖，自回归生成阶段的KV缓存应保持不共享。KV缓存共享的程度取决于所采用的特定解码算法。在更复杂的算法（如束搜索 [49]）中，不同的请求束可以共享其KV缓存的较大部分（最多可节省55%的内存，见§6.3），并且共享模式随着解码过程的推进而变化。
  未知输入和输出长度的调度。对LLM服务的请求在输入和输出长度上表现出多样性。这要求内存管理系统能够适应各种提示长度。此外，随着请求的输出长度在解码过程中增长，其KV缓存所需的内存也会扩展，并可能耗尽可用内存，从而影响请求的传入或正在进行的现有提示的生成。该系统需要做出调度决策，例如将某些请求的KV缓存从GPU内存中删除或换出。

3.1 Memory Management in Existing Systems

  由于当前深度学习框架 [33, 39] 中的大多数操作符要求张量存储在连续内存中，之前的LLM服务系统 [31, 60] 也将一个请求的KV缓存存储为跨不同位置的连续张量。由于LLM的输出长度不可预测，它们根据请求的最大可能序列长度静态地为请求分配一块内存，而不管请求的实际输入或最终输出长度如何。
  图3展示了两个请求：请求A的最大可能序列长度为2048，请求B的最大长度为512。现有系统中的块预分配方案存在三个主要的内存浪费来源：为未来token保留的槽、由于为潜在最大序列长度过度分配而导致的内部碎片，以及来自内存分配器（如伙伴分配器）的外部碎片。外部碎片将永远不会用于生成的token，这在服务请求之前是已知的。内部碎片也仍然没有使用，但这只有在请求完成采样后才能意识到。它们都是纯粹的内存浪费。尽管保留的内存最终会被使用，但在整个请求期间保留此空间，尤其是当保留空间很大时，会占用本可用于处理其他请求的空间。作者在实验中可视化了内存浪费的平均百分比，如图2所示，揭示了之前系统中的实际有效内存可以低至20.4%。

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

4 Method

  在这项工作中，作者开发了一种新的注意力算法PagedAttention，并构建了一个LLM服务引擎vLLM，以应对§3中概述的挑战。vLLM的架构如图4所示。vLLM 采用了一个集中式调度器来协调分布式 GPU 工作节点的执行。KV缓存管理器通过PagedAttention以分页方式高效管理KV缓存。具体来说，KV缓存管理器通过集中式调度器发送的指令来管理GPU工作节点上的物理KV缓存内存。
  接下来，作者将在§4.1中描述PagedAttention算法。在此基础上，作者在§4.2中展示 KV 缓存管理器的设计，并在§4.3中介绍它如何促进 PagedAttention 的实现。然后，作者展示这一设计如何促进各种解码方法（§4.4）的有效内存管理，并处理可变长度的输入和输出序列（§4.5）。最后，作者展示vLLM的系统设计如何在分布式环境中工作（§4.6）。

4.1 PagedAttention

  为了解决§3中的内存挑战，作者引入PagedAttention，这是一种受操作系统经典分页思想 [25] 启发的注意力算法。与传统的注意力算法不同，PagedAttention允许将连续的键和值存储在非连续的内存空间中。具体来说，PagedAttention将每个序列的KV缓存划分为KV块。每个块包含固定数量token的键和值向量，作者将其表示为KV块大小（B）。记键块 $K_j=(k_{(j-1)B+1},...,k_{jB})$ 和值块 $V_j=(v_{(j-1)B+1},...,v_{jB})$ 。公式4中的注意力计算可以转换为以下分块计算：
$$A_{ij}=\frac{exp(q_i^{\top }{K}_j/\sqrt{d})}{{\sum }_{t=1}^{\lceil i/B\rceil }exp(q_i^{\top }{K}_{t}1/\sqrt{d})},o_i=\sum _{j=1}^{\lceil i/B\rceil }{V}_j{A}_{ij}^{\top },\text{\hspace{1em}(4)}$$

其中$A_{ij}=(a_{i,(j-1)B+1},...,a_{i,jB})$是第 $j$ 个KV块上的注意力分数行向量。
  在注意力计算过程中，PagedAttention内核会分别识别和获取不同的KV块。作者在图5中展示了PagedAttention的一个示例：键和值向量分布在三个块上，而这三个块在物理内存中并不连续。每次，内核将查询token（“forth”）的查询向量 $q_i$ 与块中的键向量 $K_j$（例如，块0中“Four score and seven”的键向量）相乘，以计算注意力分数 $A_{ij}$，然后将 $A_{ij}$ 与块中的值向量 $V_j$ 相乘，以得出最终的注意力输出 $o_i$。
  总之，PagedAttention算法允许KV块存储在非连续的物理内存中，这使得vLLM中的分页内存管理更加灵活。

4.2 KV Cache Manager

  vLLM内存管理器背后的关键思想类似于操作系统中的虚拟内存 [25]。操作系统将内存划分为固定大小的页面，并将用户程序的逻辑页面映射到物理页面。连续的逻辑页面可以对应于非连续的物理内存页面，从而允许用户程序像访问连续内存一样访问内存。此外，物理内存空间不需要预先完全保留，使操作系统能够根据需要动态分配物理页面。vLLM利用虚拟内存背后的思想来管理LLM服务中的KV缓存。通过PagedAttention的支持，作者将KV缓存组织为固定大小的KV块，就像虚拟内存中的页面一样。
  一个请求的KV缓存被表示为一组逻辑KV块，随着新token及其KV缓存的生成，从左到右填充。最后一个KV块中未填充的位置保留给未来的生成。在GPU工作节点上，块引擎分配一块连续的GPU DRAM并将其划分为物理KV块（在CPU RAM上也进行此操作以进行交换；见§4.5）。KV块管理器还维护块表——每个请求的逻辑KV块与物理KV块之间的映射。每个块表条目记录逻辑块所对应的物理块以及已填充位置的数量。将逻辑KV块与物理KV块分离，使vLLM能够动态增长KV缓存内存，而无需预先为所有位置保留内存，从而消除了现有系统中的大部分内存浪费，如图2所示。

4.3 Decoding with PagedAttention and vLLM

  接下来，作者通过一个示例（如图6所示）来演示vLLM在单个输入序列的解码过程中如何执行PagedAttention和管理内存：①像操作系统的虚拟内存一样，vLLM 不需要在一开始就为最大可能生成的序列长度预留内存。相反，它只保留必要的KV块以容纳提示计算期间生成的KV缓存。在本例中，提示有7个token，因此vLLM将前2个逻辑KV块（0和1）映射到2个物理KV块（分别为7和1）。在预填充步骤中，vLLM使用传统的自注意力算法（例如 [13]）生成提示的KV缓存和第一个输出token。然后，vLLM将前4个token的KV缓存存储在逻辑块0中，接下来的3个token存储在逻辑块1中。剩余的槽位保留给后续的自回归生成阶段。②在第一个自回归解码步骤中，vLLM在物理块7和1上使用PagedAttention算法生成新的token。由于最后一个逻辑块中仍然有一个可用槽位，新生成的KV缓存被存储在那里，并更新块表中的#filled记录。③在第二个解码步骤中，由于最后一个逻辑块已满，vLLM将新生成的KV缓存存储在新的逻辑块中；vLLM为其分配一个新的物理块（物理块3），并将此映射存储在块表中。

  在全局范围内，对于每次解码迭代，vLLM首先选择一组候选序列进行批处理（更多内容见§4.5），并为新需要的逻辑块分配物理块。然后，vLLM将当前迭代的所有输入token（即提示阶段请求的所有token和生成阶段请求的最新token）连接为一个序列，并将其输入到LLM中。在LLM的计算过程中，vLLM使用PagedAttention内核访问以逻辑KV块形式存储的先前KV缓存，并将新生成的KV缓存保存到物理KV块中。在一个KV块（块大小 > 1）中存储多个token使PagedAttention内核能够在更多位置上并行处理KV缓存，从而提高硬件利用率并减少延迟。然而，较大的块大小也会增加内存碎片。作者在§7.2中研究了块大小的影响。
  再次强调，随着更多token及其KV缓存的生成，vLLM会动态地将新的物理块分配给逻辑块。由于所有块都是从左到右填充的，并且只有当所有先前的块都满时才会分配一个新的物理块，vLLM将请求的所有内存浪费限制在一个块内，因此可以有效地利用所有内存，如图2所示。这使得更多的请求适应内存进行批处理，从而提高了吞吐量。一旦请求完成生成，其KV块可以被释放以存储其他请求的KV缓存。在图7中，作者展示了一个vLLM为两个序列管理内存的示例。两个序列的逻辑块被映射到GPU工作节点中由块引擎保留的空间内的不同物理块。两个序列的相邻逻辑块不需要在物理GPU内存中连续，并且物理块的空间可以被两个序列有效利用。

4.4 Application to Other Decoding Scenarios

  §4.3展示了PagedAttention和vLLM如何处理基本的解码算法，例如贪婪解码和采样，这些算法将一个用户提示作为输入并生成单个输出序列。在许多成功的LLM应用 [18, 34] 中，LLM服务必须提供更复杂的解码场景，这些场景表现出复杂的访问模式和更多的内存共享机会。作者在本节中展示vLLM在这些场景中的普遍适用性。
  并行采样。在基于LLM的程序助手 [6, 18] 中，LLM为单个输入提示生成多个采样输出；用户可以从多个候选中选择一个最喜欢的输出。到目前为止，作者已经隐含地假设一个请求生成单个序列。在本文的其余部分，作者假设更一般的情况，即一个请求生成多个序列。在并行采样中，一个请求包含多个共享相同输入提示的样本，从而允许该提示的KV缓存也被共享。通过其PagedAttention和分页内存管理，vLLM可以轻松实现这种共享并节省内存。
  图8展示了一个两个输出的并行解码示例。由于两个输出共享相同的提示，作者在提示阶段只为提示状态保留一份副本的空间；两个序列提示的逻辑块映射到相同的物理块：两个序列的逻辑块0和1分别映射到物理块7和1。由于单个物理块可以映射到多个逻辑块，作者为每个物理块引入了一个引用计数。在这种情况下，物理块7和1的引用计数均为2。在生成阶段，两个输出采样不同的输出token，并需要单独的KV缓存存储。vLLM在块粒度上实现了写时复制机制，用于需要由多个序列修改的物理块，类似于操作系统虚拟内存中的写时复制技术（例如，在fork进程时）。具体来说，在图8中，当样本A1需要写入其最后一个逻辑块（逻辑块1）时，vLLM识别到相应物理块（物理块1）的引用计数大于1；它分配一个新的物理块（物理块3），指示块引擎从物理块1中复制信息，并将引用计数减少到1。接下来，当样本A2写入物理块1时，引用计数已经减少到1；因此A2直接将其新生成的KV缓存写入物理块1。

  总之，vLLM允许多个输出样本共享用于存储提示KV缓存的大部分空间，除了被写时复制机制管理的最后一个逻辑块。通过在多个样本之间共享物理块，可以大大减少内存使用，特别是对于长输入提示。
  束搜索。在机器翻译 [59] 等LLM任务中，用户期望LLM输出前 $k$ 个最合适的翻译。束搜索 [49] 被广泛用于从LLM中解码最可能的输出序列，因为它降低了完全遍历样本空间的计算复杂度。该算法依赖于束宽度参数 $k$ ，它决定每一步保留的top候选数量。在解码过程中，束搜索通过考虑所有可能的token来扩展束中的每个候选序列，使用LLM计算它们各自的概率，并从 $k\cdot |V|$ 个候选中保留前 $k$ 个最可能的序列，其中 $|V|$ 是词汇表大小。
  与并行解码不同，束搜索不仅支持初始提示块的共享，还支持不同候选项之间其他块的共享。随着解码过程的推进，共享模式会动态变化，类似于操作系统中由复合fork创建的进程树。图 9 展示了 vLLM 在k = 4的束搜索示例中如何管理KV块。在虚线所示的迭代之前，每个候选序列已使用了4个完整的逻辑块。所有束候选共享第一个块0（即提示）。候选3从第二个块开始与其他候选分离。候选0-2共享前3个块，并在第4个块处分叉。在后续迭代中，前4个可能的候选都源自候选1和2。由于原始候选0和3不再是最优候选，它们的逻辑块被释放，相应物理块的引用计数减少。vLLM释放所有引用计数达到0的物理块（块2、4、5、8）。然后，vLLM分配新的物理块（块9-12）以存储来自新候选的新KV缓存。现在，所有候选共享块0、1、3；候选0和1共享块6，候选2和3进一步共享块7。

  以前的LLM服务系统需要在束候选之间进行KV缓存的频繁内存复制。例如，在图9所示的情况下，在虚线之后，候选3需要复制大部分候选2的KV缓存以继续生成。vLLM的物理块共享显著减少了这种频繁内存复制的开销。在vLLM中，不同束候选的大多数块可以共享。写时复制机制仅在新生成的 tokens 位于旧的共享块内时应用，类似于并行解码。这仅涉及复制一个数据块。
  共享前缀。通常，LLM用户会提供任务的（长）描述，包括指令和示例输入输出，也称为系统提示 [36]。该描述与实际任务输入连接在一起，形成请求的提示。LLM基于完整提示生成输出。图10展示了一个示例。此外，可以通过提示工程进一步调整共享前缀，以提高下游任务的准确性 [26, 27]。

  对于这种应用类型，许多用户提示共享一个前缀，因此LLM服务提供商可以预先存储前缀的KV缓存，以减少在前缀上花费的冗余计算。在 vLLM 中，可以通过为 LLM 服务提供商预定义的一组共享前缀保留一组物理块来方便地实现这一点，类似于操作系统在进程间处理共享库的方式。带有共享前缀的用户输入提示可以简单地将其逻辑块映射到缓存的物理块（最后一个块标记为写时复制）。提示阶段的计算只需要在用户的任务输入上执行。
  混合解码方法。前面讨论的解码方法表现出不同的内存共享和访问模式。尽管如此，vLLM推动了具有不同解码偏好的请求的同时处理，这是现有系统无法有效做到的。这是因为vLLM通过一个将逻辑块转换为物理块的通用映射层，隐藏了不同序列之间的复杂内存共享。LLM及其执行内核只看到每个序列的物理块ID列表，不需要处理跨序列的共享模式。与现有系统相比，这种方法拓宽了具有不同采样要求的请求的批处理机会，最终增加了系统的整体吞吐量。

4.5 Scheduling and Preemption

  当请求流量超过系统容量时，vLLM必须优先处理部分请求。在vLLM中，作者采用先到先服务（FCFS）调度策略来处理所有请求，以确保公平性并防止饥饿现象。当vLLM需要抢占请求时，它确保最早到达的请求优先得到服务，而最新的请求优先被抢占。
  LLM服务面临一个独特的挑战：LLM的输入提示在长度上可能差异很大，并且所得的输出长度事先未知，取决于输入提示和模型。随着请求数量及其输出的增长，vLLM可能会用尽GPU的物理块来存储新生成的KV缓存。在这种情况下，vLLM需要回答两个经典问题：（1）应该驱逐哪些块？（2）如果需要再次使用，如何恢复被驱逐的块？通常，驱逐策略使用启发式方法预测哪个块将在未来最远时间被访问，然后驱逐该块。由于在作者的案例中，作者知道一个序列的所有块都是一起访问的，因此作者实现了一种全有或全无的驱逐策略，即要么驱逐一个序列的所有块，要么一个都不驱逐。此外，一个请求中的多个序列（例如，一个束搜索请求中的束候选）作为一个序列组进行组调度。由于这些序列之间可能存在内存共享，一个序列组中的所有序列总是一起被抢占或重新调度。为了回答如何恢复被驱逐的块的第二个问题，作者考虑两种技术：
  交换。这是大多数虚拟内存实现使用的经典技术，它将驱逐的页面复制到磁盘上的交换空间。在作者的案例中，作者将驱逐的块复制到CPU内存中。如图4所示，除了GPU块分配器外，vLLM还包括一个CPU块分配器来管理交换到CPU RAM的物理块。当vLLM耗尽用于新token的可用物理块时，它会选择一组序列进行驱逐，并将其KV缓存传输到CPU。一旦抢占某个序列并驱逐其块，vLLM 就会停止接受新请求，直到所有被抢占的序列完成为止。一旦请求完成，其块将从内存中释放，并将被抢占序列的块重新加载以继续处理该序列。请注意，通过这种设计，交换到CPU RAM的块数永远不会超过GPU RAM中的总物理块数，因此CPU RAM上的交换空间受限于为KV缓存分配的GPU内存。
  重新计算。在这种情况下，当被抢占的序列重新调度时，作者简单地重新计算其KV缓存。需要注意的是，重新计算的延迟可能显著低于原始延迟，因为解码时生成的token可以与原始用户提示拼接为新的提示——它们在所有位置的KV缓存可以在一次提示阶段迭代中生成。
  交换和重新计算的性能取决于CPU RAM和GPU内存之间的带宽以及GPU的计算能力。作者在§7.3中测试了交换和重新计算的速度。

4.6 Distributed Execution

  许多LLM的参数大小超过了单个GPU的容量 [5, 9]。因此，有必要将它们划分到分布式 GPU 上，并以模型并行的方式执行 [28, 63]。这需要一个能够处理分布式内存的内存管理器。vLLM通过支持 Transformer 上广泛使用的 Megatron-LM 样式的张量模型并行策略 [47]，在分布式环境中非常有效。该策略遵循SPMD（单程序多数据）执行调度，其中线性层被分区以执行块状矩阵乘法，并且GPU通过全规约操作不断同步中间结果。具体来说，注意力操作在注意力头维度上分割，每个SPMD进程负责多头注意力中的一部分注意力头。
  作者观察到，即使使用模型并行执行，每个模型分片仍然处理相同的输入token集，因此需要相同位置的KV缓存。因此，vLLM 在集中式调度器中配备了一个单一的 KV 缓存管理器，如图4所示。不同的GPU工作节点共享该管理器以及从逻辑块到物理块的映射。这种通用映射允许GPU工作节点使用调度器为每个输入请求提供的物理块来执行模型。尽管每个GPU工作节点具有相同的物理块ID，但每个工作节点仅存储其对应注意力头的部分KV缓存。
  在每个步骤中，调度器首先为批处理中的每个请求准备包含输入token ID的消息以及每个请求的块表。接下来，调度器将该控制消息广播给 GPU 工作节点。然后，GPU 工作节点开始使用输入token ID执行模型。在注意力层中，GPU工作节点根据控制消息中的块表读取KV缓存。在执行过程中，GPU工作节点使用全规约通信原语同步中间结果，无需调度器的协调，如 [47] 中所述。最后，GPU工作节点将该迭代中采样的token发送回调度器。总之，GPU工作节点不需要在内存管理上进行同步，因为它们只需要在每次解码迭代开始时接收所有内存管理信息以及步骤输入。

5 Implementation

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

5.1 Kernel-level Optimization

  由于PagedAttention引入了现有系统无法有效支持的内存访问模式，作者开发了几个GPU内核来优化它。（1）融合的reshape和块写入。在每个Transformer层中，新的KV缓存被分割成块，reshape为优化块读取的内存布局，然后保存在块表指定的位置。为了最小化内核启动开销，作者将它们融合到一个内核中。（2）融合的块读取和注意力。作者调整了FasterTransformer [31] 中的注意力内核，以根据块表读取KV缓存并即时执行注意力操作。为了确保合并的内存访问，作者分配一个GPU warp来读取每个块。此外，作者增加了对请求批次中可变序列长度的支持。（3）融合的块复制。由写时复制机制发出的块复制操作可能在不连续的块上操作。如果作者使用cudaMemcpyAsync API，这可能导致大量小数据移动的调用。为了减轻开销，作者实现了一个内核，将不同块的复制操作批量处理为单个内核启动。

5.2 Supporting Various Decoding Algorithms

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

6 Evaluation

在本节中，作者将在多种工作负载下评估 vLLM 的性能。

6.1 Experimental Setup

  模型和服务器配置。作者使用13B、66B和175B参数的OPT [62] 模型以及13B参数的LLaMA [52] 进行评估。13B 和 66B 是 LLM 中常见的模型规模，如 LLM 排行榜 [38] 所示，而 175B 则是著名的 GPT-3 [5] 模型的规模。对于所有的实验，作者使用 Google Cloud Platform 上配备 NVIDIA A100 GPU 的 A2 实例。详细的模型规模和服务器配置如表 1 所示。

  工作负载。作者基于ShareGPT [51] 和Alpaca [50] 数据集合成工作负载，这些数据集包含真实LLM服务的输入和输出文本。ShareGPT 数据集是由用户共享的与 ChatGPT 对话的集合 [35] 。Alpaca数据集是通过GPT3.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] 是一个优化吞吐量的最先进LLM服务系统。由于Orca未公开可用，作者实现了自己的Orca版本。作者假设 Orca 使用伙伴分配算法来确定存储 KV 缓存的内存地址。作者基于请求输出的预留空间程度实现了 Orca 的三个版本：
  • Orca (Oracle)。作者假设系统已知将为请求实际生成的输出长度。这显示 Orca 的上限性能，在实际中实现是不可行的。
  • Orca (Pow2)。作者假设系统为输出预留最多 2 倍的空间。例如，如果实际输出长度为 25，则为输出预留 32 个位置。
  • Orca (Max)。作者假设系统始终将空间预留到模型的最大序列长度，即 2048 个 token。
  关键指标。作者关注服务吞吐量。具体而言，在使用不同请求率的工作负载下，作者测量系统的标准化延迟，即每个请求的端到端延迟除以其输出长度的平均值，如Orca [60] 中所述。高吞吐量的服务系统在高请求率下应保持较低的标准化延迟。在大多数实验中，作者使用 1 小时的请求轨迹进行评估。作为例外，由于成本限制，对于 OPT-175B 模型，作者使用 15 分钟的请求轨迹。

6.2 Basic Sampling

  作者在三个模型和两个数据集上使用基本采样（每个请求一个样本）评估了 vLLM 的性能。图 12 的第一行展示了 ShareGPT 数据集上的结果。曲线表明，随着请求率的增加，延迟最初呈现逐步上升的趋势，但随后突然激增。这可以归因于当请求率超过服务系统的容量时，队列长度持续无限增长，从而导致请求延迟也不断增加。

  在 ShareGPT 数据集上，vLLM 在保持相似延迟的情况下，能够支持比 Orca (Oracle) 高 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 Parallel Sampling and Beam Search

  作者使用两种流行的采样方法评估PagedAttention中内存共享的有效性：并行采样和束搜索。在并行采样中，请求中的所有并行序列可以共享提示的KV缓存。如图14的第一行所示，随着采样序列数量的增加，vLLM相较于Orca基线带来了更大的性能提升。类似地，图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 Shared prefix

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

6.5 Chatbot

  聊天机器人 [8, 19, 35] 是LLM最重要的应用之一。为了实现聊天机器人，作者让模型通过将聊天历史和最后一个用户查询连接为提示来生成响应。作者使用ShareGPT数据集合成了聊天历史和用户查询。由于OPT-13B模型的上下文长度有限，作者将提示截断到最后1024个token，并让模型生成最多1024个token。作者不在不同对话轮次之间存储KV缓存，因为这样做会占用对话轮次之间其他请求的空间。
  图17显示，与三种Orca基线相比，vLLM可以维持2倍高的请求率。由于ShareGPT数据集包含许多长对话，大多数请求的输入提示有1024个token。由于伙伴分配算法，Orca基线为请求输出保留了1024个token的空间，无论它们如何预测输出长度。因此，这三种 Orca 基线表现相似。相比之下，vLLM可以有效地处理长提示，因为PagedAttention解决了内存碎片和保留的问题。

7 Ablation Studies

在本节中，作者研究了vLLM的各个方面，并通过消融实验评估了作者的设计选择。

7.1 Kernel Microbenchmark

  PagedAttention中的动态块映射会影响涉及存储KV缓存的GPU操作的性能，即块读/写和注意力。与现有系统相比，作者的GPU内核（§5）涉及访问块表、执行额外分支和处理可变序列长度的额外开销。如图18a所示，这导致了 20%-26% 的注意力内核延迟提升，相较于高度优化的 FasterTransformer 实现。作者认为开销很小，因为它只影响注意力操作符，而不影响模型中的其他操作符，例如线性。尽管有开销，PagedAttention使vLLM在端到端性能上显著优于FasterTransformer（§6）。

7.2 Impact of Block Size

  块大小的选择对vLLM的性能有重大影响。如果块大小太小，vLLM可能无法充分利用GPU的并行性来读取和处理KV缓存。如果块大小太大，内部碎片增加，共享的概率降低。
  在图18b中，作者在固定请求率下使用 ShareGPT 和 Alpaca基本采样轨迹评估了vLLM在不同块大小下的性能。在ShareGPT轨迹中，块大小从16到128导致最佳性能。在Alpaca轨迹中，虽然块大小16和32效果很好，但较大的块大小显著降低了性能，因为序列变得比块大小更短。在实践中，作者发现块大小16足够大以有效利用GPU，并且足够小以避免在大多数工作负载中出现显著的内部碎片。因此，vLLM将其默认块大小设置为16。

7.3 Comparing Recomputation and Swapping

  vLLM支持重新计算和交换作为其恢复机制。为了理解这两种方法之间的权衡，作者评估了它们的端到端性能，并对它们的开销进行了微基准测试，如图19所示。作者的结果表明，交换在小块大小时会产生过多的开销。这是因为小块大小通常会导致CPU和GPU之间的大量小数据传输，这限制了有效的PCIe带宽。相比之下，重新计算的开销在不同块大小下保持不变，因为重新计算不使用KV块。因此，当块大小较小时，重新计算更有效，而当块大小较大时，交换更有效，尽管重新计算的开销从未超过交换延迟的20%。对于16到64的中等块大小，这两种方法表现出相当的端到端性能。

8 Discussion

  将虚拟内存和分页技术应用于其他GPU工作负载。虚拟内存和分页的思想对于管理LLM服务中的KV缓存是有效的，因为工作负载需要动态内存分配（由于输出长度无法预知），并且其性能受GPU内存容量的限制。然而，这并不普遍适用于每个GPU工作负载。例如，在DNN训练中，张量形状通常是静态的，因此可以提前优化内存分配。再例如，在服务非LLM的DNN时，内存效率的提高可能不会带来任何性能改进，因为性能主要受计算限制。在这种情况下，引入vLLM的技术可能会由于内存间接和非连续块内存的额外开销而降低性能。但是，作者很高兴看到vLLM的技术被应用于其他具有与LLM服务相似属性的工作负载。
  在应用虚拟内存和分页时的LLM特定优化。vLLM通过利用特定应用语义重新解释并增强了虚拟内存和分页的思想。一个例子是vLLM的全有或全无的换出策略，它利用了处理请求需要将其所有对应的token状态存储在GPU内存中的事实。另一个例子是重新计算方法来恢复被驱逐的块，这在操作系统中是不可行的。此外，vLLM 通过将用于内存访问操作的 GPU 内核与注意力等其他操作的内核融合，降低了分页中内存间接访问的开销。

9 Related Work

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

10 Conclusion

  本文提出了PagedAttention，一种新的注意力算法，允许注意力键和值存储在非连续的分页内存中，并提出了vLLM，一个由PagedAttention实现高效内存管理的高吞吐量LLM服务系统。受操作系统的启发，作者展示了如何将虚拟内存和写时复制等成熟技术应用于 LLM 服务中，以高效管理KV缓存并处理各种解码算法。作者的实验表明，vLLM比最先进的系统实现了2-4倍的吞吐量提升。

论文链接：https://dl.acm.org/doi/abs/10.1145/3600006.3613165