LLM在处理大量请求时面临内存使用效率低下的问题，每个请求需占用大量动态变化的内存，导致浪费并限制处理能力。本文提出了PagedAttention算法，受操作系统虚拟内存分页技术启发，开发了vLLM服务系统。vLLM通过灵活共享内存，实现几乎零内存浪费，显著降低了内存需求。测试表明，vLLM处理吞吐量比FasterTransformer和Orca提高2到4倍，且延迟保持相同，尤其在处理长文本、大模型和复杂解码时表现突出。

论文：(SOSP 2023)Efficient Memory Management for Large Language Model Serving with PagedAttention
代码：https://github.com/vllm-project/vllm
文章参考自：https://mp.weixin.qq.com/s/whsGK2gfVrIDNXTtxUUSOw

介绍

许多云服务公司正在争相提供LLM应用，但运行这些应用的成本非常高，需要大量的硬件加速器如GPU。据估计，处理一次大语言模型的请求，成本是传统关键词查询的 10 倍。由于成本如此之高，提升大语言模型的处理效率，从而降低每次请求的费用，变得越来越重要。

大语言模型（LLM）的核心是一个自回归的Transformer模型。这个模型根据输入的提示和之前生成的词（标记），一个一个地生成新的词。每次请求都需要重复这个复杂的过程，直到模型输出一个终止标记。这种逐字生成的过程需要大量内存，无法充分利用GPU的计算能力，限制了处理请求的效率。为了提高效率，可以将多个请求批量处理。但要做到这一点，需要有效管理每个请求的内存空间。

内存需求计算公式

  • 模型参数

    13B参数 ~ 13G 个参数 *  4个字节（FP 32） 或者 2个bytes （FP16，生产环境一般用FP16）

  • KV Cache

    2 (key and value vectors) * hidden_size  * layer_num *  head_num * batch_size * seq_len * 参数精度（4/2字节）

对于更大的模型，和更长的上下文长度，KV Cache的存储需求会显著增加。由于模型的参数是固定的，而激活占用的内存很少，KV缓存的管理方式决定了最大的批处理大小。如果管理不当，KV缓存会显著限制批处理的大小，从而影响LLM的处理效率。

作者发现现有的LLM服务系统在管理KV缓存内存方面效率不高。主要因为：

  • 内存碎片化

    KV Cache随着模型生成新词而 动态增长和收缩，并且它的生命周期和长度是未知的。现有系统为了在连续空间中存储请求的KV缓存，它们会 预先分配 一大块内存，按照请求的最大长度（比如2048个词）来分配，但请求的实际长度往往比最大长度短很多，且不能被其他短请求占用，导致严重的内部碎片化。而且，不同请求的预分配大小不同，还会导致 外部内存碎片化 。

    图中显示了两个请求：最大可能序列长度为2048的请求A和最大序列长度为512的请求B。三个主要的内存浪费来源：

    • 为未来令牌预留的插槽
    • 过度配置最大序列长度而导致的内部碎片
    • 来自内存分配器（如伙伴分配器）的外部碎片

    作者的分析结果显示，现有系统中只有20.4%-38.2%的KV缓存内存实际用于存储标记状态。

  • 无法利用内存共享

    LLM服务常常使用高级解码算法，比如并行采样和束搜索，这些算法会为每个请求生成多个输出。在这些场景中，请求包含的多个序列可以部分共享其KV缓存。但现有系统中KV Cache存储在不同的连续空间中，无法实现内存共享。

为了解决上述问题，作者受 操作系统虚拟内存分页 启发，提出了PagedAttention，将请求的KV缓存分成多个小块，每个小块包含固定数量的注意力键和值。这些小块不需要存储在连续的空间中，因此可以像操作系统管理虚拟内存那样灵活地管理KV缓存。通过使用相对较小的小块并按需分配，PagedAttention缓解了内部碎片化问题。此外，由于所有小块大小相同，它还消除了外部碎片化问题。它实现了内存共享，使同一请求的不同序列甚至不同请求之间可以在小块级别共享内存。

作者基于PagedAttention构建了一个高吞吐量的分布式LLM服务引擎vLLM，该引擎实现了KV缓存内存的近乎零浪费，支持各种流行的大语言模型，如GPT、OPT和LLaMA，并且支持超过单个GPU内存容量的模型。评估结果显示，vLLM在各种模型和工作负载下，相比最先进的系统，其LLM服务吞吐量提高了2-4倍。且这种改进在处理较长序列、较大模型和更复杂的解码算法时更加显著。

解决方案

PageAttention算法

PagedAttention将每个序列的键值（KV）缓存划分为KV块。每个块包含一定数量的令牌的键和值向量。例如，键块𝐾𝑗包含的是第𝑗个块中的键向量，而值块𝑉𝑗包含的是第𝑗个块中的值向量。在注意力计算过程中，PagedAttention内核分别识别和获取不同的KV块。

例如，当处理查询标记为“forth”时，内核将查询向量𝑞𝑖与每个块中的键向量𝐾𝑗相乘，以计算注意力得分𝐴𝑖𝑗。然后，它将这些得分与每个块中的值向量𝑉𝑗相乘，得出最终的注意力输出。

KV Cache管理器

KV Cache管理器将KV缓存组织成固定大小的KV块，类似于虚拟内存中的页面。一个请求的KV缓存被表示为一系列逻辑KV块，当新的令牌及其KV缓存生成时，从左到右填充。最后一个KV块的未填充位置保留供将来使用。

在GPU工作器上，block engine分配一块连续的GPU DRAM，并将其划分为物理KV块。KV块管理器还维护块表——每个请求的逻辑和物理KV块之间的映射关系。每个块表条目记录了逻辑块的相应物理块及其填充位置的数量。将逻辑和物理KV块分开使得vLLM可以动态增长KV缓存内存，而无需提前为所有位置预留空间，这消除了现有系统中大部分的内存浪费。

基本的推理场景

本例中，提示有7个标记，因此vLLM将前两个逻辑KV块（0和1）映射到两个物理KV块（7和1）。

在prefill步骤中，vLLM使用传统的自注意力算法生成KV缓存和第一个token。然后，vLLM将前4个标记的KV缓存存储在逻辑块0中，后续3个标记存储在逻辑块1中。剩余的位置留待后续自回归生成阶段使用。

在第1个decode步骤中，vLLM使用PagedAttention算法在物理块7和1上生成新的标记。由于最后一个逻辑块中还有一个空位，新生成的KV缓存被存储在那里，并且更新了块表中的已填充记录。

在第2个decode步骤中，由于最后一个逻辑块已满，vLLM将新生成的KV缓存存储在一个新的逻辑块中；vLLM为其分配一个新的物理块（物理块3）并将此映射存储在块表中。

对于每个decode迭代，vLLM首先选择一组候选序列进行批处理(原理同Orca选择性批处理)，并为新需求的逻辑块分配物理块。vLLM动态地为逻辑块分配新的物理块，将请求的所有内存浪费限制在一个块内，因此可以有效地利用所有内存，

这两个序列的逻辑块被映射到GPU工作器中块引擎预留的不同物理块中。这两个序列的相邻逻辑块在物理GPU内存中不需要连续，物理块的空间可以被两个序列有效地利用。

不同的推理场景

Parallel sampling（并行采样）

在 LLM 应用中，一个输入提示可能生成多个输出，用户可以从多个候选项中选择最喜欢的结果。一个请求包含多个共享相同输入提示的样本，因此提示的KV缓存也可以共享。

图中展示了两个输出的并行解码示例。为了管理共享，vLLM为每个物理块引入了引用计数，物理块7和1的引用计数都是2。在生成阶段，这两个输出生成不同的输出标记，因此需要单独的KV缓存存储。

vLLM对需要被多个序列修改的物理块实现了块级的copy-on-write机制（类似于操作系统中的写时复制技术）。当样本A1需要写入其最后一个逻辑块（逻辑块1）时，vLLM检测到相应的物理块（物理块1）的引用计数大于1，于是分配一个新的物理块（物理块3），并指示块引擎将信息从物理块1复制过去，同时将引用计数减为1。接下来，当样本A2写入物理块1时，引用计数已经减少到1，因此A2可以直接将新生成的KV缓存写入物理块1。

Beam search（束状搜索）

在LLM（大型语言模型）任务如机器翻译中，用户通常期望得到最合适的前𝑘个翻译结果。Beam search是一种常用的解码算法，用于从LLM中生成最可能的输出序列。算法依赖于一个参数𝑘，该参数决定每一步保留的最顶尖的候选序列数量。在解码过程中，Beam search会扩展每个候选序列，考虑所有可能的标记，计算它们各自的概率，并从𝑘 · |𝑉|个候选中保留前𝑘个最可能的序列，其中|𝑉|是词汇表的大小。

与并行解码不同，Beam search不仅可以共享初始提示块，还可以在不同的候选序列之间共享其他块，并且这些共享模式会随着解码过程的推进动态变化，类似于操作系统中由复合分叉创建的进程树。

图中展示了vLLM如何管理一个𝑘=4的Beam search示例中的KV块。在图中虚线之前，每个候选序列已经使用了4个完整的逻辑块。所有Beam候选共享第一个块0（即提示）。候选3从第二块开始与其他候选分开。候选0-2共享前三个块，并在第四块分开。在随后的迭代中，最可能的前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服务系统需要在Beam候选之间频繁复制KV缓存，这会带来很大的内存复制开销。例如，在图所示的情况下，虚线之后，候选3需要复制候选2的大部分KV缓存以继续生成。vLLM通过物理块共享显著减少了这种频繁的内存复制开销。在vLLM中，不同Beam候选的大多数块可以共享。只有当新生成的标记位于旧的共享块内时，才会应用写时复制机制，这仅涉及复制一个块的数据。

总结来说，vLLM通过有效的物理块共享和写时复制机制，大大提高了Beam search解码过程中的内存利用效率，减少了内存复制开销。

Shared prefix（共享前缀）

许多用户提示可能共享一个前缀，因此LLM服务提供者可以提前存储这个前缀的KV缓存，以减少在前缀上的重复计算。在vLLM中，可以通过为一组预定义的共享前缀保留一组物理块来方便地实现这一点，就像操作系统处理跨进程共享库一样。

混合解码

基于上面的内存管理方法，vLLM能够同时处理具有不同甚至是混合解码偏好的请求，这是现有系统无法高效完成的。这是因为vLLM通过一个公共映射层隐藏了不同序列之间复杂的内存共享，该映射层将逻辑块转换为物理块。

LLM及其执行内核只需看到每个序列的一组物理块ID，而不需要处理跨序列的共享模式。与现有系统相比，这种方法拓宽了具有不同采样需求的请求的批处理机会，从而最终提高了系统的整体吞吐量。

调度与抢占

当请求流量超过系统容量时，采用先到先服务（FCFS）调度策略，确保公平性并防止饥饿。

  1. 驱逐
    通常，驱逐策略使用启发式方法来预测哪个块将来访问的最远，并驱逐该块。由于一个序列的所有块都是一起访问的，所以我们实施了一种全有或全无的驱逐策略，即要么驱逐序列的所有块，要么一个也不驱逐。

  2. 恢复
    交换：将被驱逐的块复制到 CPU 内存中。当 vLLM 为新令牌耗尽自由物理块时，它选择一组序列进行驱逐，并将它们的 KV 缓存传输到 CPU。一旦抢占了一个序列并驱逐了它的块，vLLM 就停止接受新请求，直到所有抢占的序列完成为止。一旦一个请求完成，其块就会从内存中释放，抢占的序列的块就会被重新引入以继续处理该序列。
    重新计算：重新计算的延迟可能低于交换延迟，性能取决于 CPU RAM 和 GPU 内存之间的带宽以及 GPU 的计算能力。为了了解两种方法之间的权衡，作者评估了它们的端到端性能，并对它们的开销进行微基准测试。结果显示，交换在小块较小时产生了过多的开销。这是因为小块通常导致CPU和GPU之间大量的小数据传输，从而限制了有效的PCIe带宽。

相比之下，重新计算的开销在不同的块大小下保持不变，因为重新计算不利用KV块。因此，当块大小较小时，重新计算更有效率，而当块大小较大时，交换更有效率，尽管重新计算的开销从未高于交换的延迟的20%。对于从16到64的中等块大小，这两种方法表现出相当的端到端性能。

分布式执行

vLLM使用的Megatron-LM张量并行策略。每个模型分片处理相同的输入，因此需要相同位置的KV缓存。因此，vLLM在集中式调度器中具有一个单一的KV缓存管理器。不同的GPU工作节点共享该管理器，以及逻辑块到物理块的映射。

系统实现

vLLM基于FastAPI的前端和基于GPU的推理引擎。前端扩展了OpenAI API接口。

vLLM引擎由8.5K行Python代码和2K行C++/CUDA代码编写而成。使用Python开发了调度器和块管理器等控制相关组件，同时为PagedAttention等关键操作开发了自定义CUDA核心。在分布式GPU工作节点之间的张量通信中，使用NCCL。

核心级优化

由于PagedAttention引入了现有系统不高效支持的内存访问模式，开发了几个GPU Kernel进行优化。

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

讨论

虚拟内存和分页技术在其他GPU工作负载中的应用：

虚拟内存和分页在LLM（大型语言模型）推理中的有效性源于其动态内存分配需求（因为输出长度不可预知）及其对GPU内存容量的依赖。然而，这并不适用于所有GPU工作负载。例如，在DNN（深度神经网络）训练中，张量形状通常是静态的，因此可以预先优化内存分配。在非LLM的DNN推理中，内存效率的提升可能不会带来性能提升，因为性能主要受计算能力限制。在这些场景中，vLLM技术可能反而会因为内存重定向和非连续内存块带来的额外开销而降低性能。

然而，对于具有与LLM推理类似特性的其他工作负载，vLLM技术的应用仍具有潜力。

相关工作

通用模型服务系统

Clipper[11]、TensorFlow Serving[33]、Nexus[45]、InferLine[10]和Clockwork[20]是一些早期的通用模型服务系统。他们研究为单个或多个模型提供服务的批处理、缓存、放置和调度。最近，DVABatch[12]引入了多入口多出口批处理。REEF[21]和Shepherd[61]建议优先抢占。AlpaServe[28]利用模型并行性进行统计复用。然而，这些通用系统未能考虑LLM推理的自回归特性和令牌状态，导致错过了优化机会。

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中提出的技术变得更加关键。

内存优化

GPU的计算能力和内存容量之间的差距越来越大，导致内存成为训练和推理的瓶颈。交换[23,42,55]、重新计算[7,24]及其组合[40]已被用于减少训练的峰值内存。值得注意的是，FlexGen[46]研究了在GPU内存有限情况下如何通过交换LLM推理的权重和令牌状态，但它不针对在线服务设置。OLLA[48]优化了张量的生存期和位置以减少碎片化，但它不进行细粒度块级管理或在线服务。FlashAttention[13]应用平铺和内核优化来减少注意力计算的峰值内存并降低I/O成本。本文介绍了一种在线服务环境下块级内存管理的新思路。

总结

本文提出了一种新的注意力算法PagedAttention，该算法允许在非连续的分页内存中存储注意力键和值，并介绍了vLLM，一种通过PagedAttention实现高效内存管理的高吞吐量大语言模型（LLM）服务系统。受操作系统的启发，展示了如何改编诸如虚拟内存和写时复制等成熟技术，以高效管理KV缓存并处理LLM服务中的各种解码算法。实验结果表明，vLLM相比现有的最先进系统实现了2-4倍的吞吐量提升。