Efficient Memory Management for Large Language Model Serving with PagedAttention

高效管理大型语言模型服务内存的 PagedAttention
作者/机构: Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica (UC Berkeley, Stanford University, Independent Researcher, UC San Diego)

A1 主要贡献

本文旨在解决大型语言模型（LLM）服务中的一个核心瓶颈：KV缓存（KV cache）的内存管理。现有系统由于将每个请求的KV缓存存储在连续的内存空间中，导致了严重的内存浪费，主要表现为内部碎片化（为请求预留最大长度的内存，但实际使用远小于此）和外部碎片化，这极大地限制了服务批处理大小（batch size），从而影响吞吐量。此外，现有系统无法有效利用并行采样、波束搜索等解码算法中存在的内存共享机会。

为解决这些问题，本文提出了以下核心贡献：
* 识别并量化了LLM服务中的内存分配挑战：论文指出，现有系统中KV缓存的有效内存利用率可能低至20.4%，大部分内存因碎片化和过度预留而被浪费（如图2所示）。
* 提出PagedAttention算法：受操作系统中虚拟内存和分页技术的启发，本文设计了一种新的注意力算法——PagedAttention。该算法允许将KV缓存存储在非连续的内存块（blocks）中，类似于操作系统中的物理页。这使得内存管理更加灵活，可以按需分配，从而基本消除了内存碎片。
* 设计并实现vLLM服务引擎：基于PagedAttention，本文构建了一个名为vLLM的分布式LLM服务引擎。vLLM通过其KV缓存管理器和调度器，实现了近乎零浪费的内存使用，并支持在请求内部及跨请求之间灵活共享KV缓存，例如在并行采样和波束搜索场景中通过写时复制（copy-on-write）机制实现。
* 显著提升服务性能：评估结果显示，与FasterTransformer和Orca等当前最先进的系统相比，vLLM在延迟水平相同的情况下，将流行LLM的吞吐量提升了2-4倍。在处理长序列、大模型和复杂解码算法时，性能提升尤为显著。

A3 背景知识与关键挑战

Transformer与自回归生成

语言模型的核心是自回归分解。语言建模的任务是计算一系列词元（token）的联合概率 $P(x_1, ..., x_n)$。通过自回归分解，这个联合概率可以被分解为条件概率的乘积，如公式1所示。

Transformer架构及其自注意力机制。Transformer已成为大规模语言建模的事实标准。其核心是自注意力层。对于一个输入隐藏状态序列 $(h_1, ..., h_i) \in \mathbb{R}^{i \times d}$，自注意力层首先对每个位置 $j$ 应用线性变换，得到查询（query）、键（key）和值（value）向量，如公式2所示。


接着，通过将位置 $i$ 的查询向量与之前所有位置的键向量相乘来计算注意力分数 $a_{i,j}$，并计算输出 $o_i$作为值向量的加权平均，如公式3和4所示。除了注意力计算，Transformer中的其他组件（如嵌入层、前馈层等）都是按位置独立应用的。

LLM服务的自回归生成过程。LLM服务通常以条件生成的形式提供，用户提供一个输入提示（prompt），模型根据该提示生成一系列输出词元。由于公式1的分解特性，新词元只能逐一生成，每个新词元的生成都依赖于序列中所有先前词元的键（key）和值（value）向量。这些向量通常被缓存起来以供未来词元生成使用，这被称为KV缓存。值得注意的是，一个词元的KV缓存依赖于其所有前序词元，因此同一词元在序列中不同位置的KV缓存是不同的。

LLM服务的两个计算阶段。LLM服务的生成计算可分为两个阶段：
1. 提示阶段（Prompt Phase）：处理整个用户提示 $(x_1, ..., x_i)$，计算第一个新词元 $x_{i+1}$ 的概率，并生成提示中所有词元的键向量和值向量。由于提示词元已知，此阶段的计算可以并行化（使用矩阵-矩阵乘法），能有效利用GPU的并行能力。
2. 自回归生成阶段（Autoregressive Generation Phase）：顺序生成剩余的新词元。在第 $j$ 次迭代中，模型输入一个词元 $x_{i+j}$，并利用已缓存的从1到 $i+j-1$ 位置的键和值向量，计算下一个词元 $x_{i+j+1}$ 的概率，同时生成位置 $i+j$ 的新键和值向量。由于数据依赖性，不同迭代间的计算无法并行化，通常使用效率较低的矩阵-向量乘法。因此，此阶段严重未充分利用GPU计算能力，成为内存密集型（memory-bound）操作，并占据了单个请求延迟的大部分。

LLM的批处理技术。通过批处理多个请求可以提高计算利用率，因为模型权重在批处理的请求间共享。然而，LLM服务的批处理面临两个挑战：一是请求在不同时间到达，简单的批处理会导致延迟增加；二是请求的输入和输出长度差异巨大，直接填充会导致计算和内存浪费。为解决此问题，研究者提出了细粒度批处理机制，如蜂窝批处理（cellular batching）【索引16, Low latency rnn inference with cellular batching，Pin Gao等人，2018年，EuroSys】和迭代级调度（iteration-level scheduling）【索引60, Orca: A Distributed Serving System for {Transformer-Based} Generative Models，Gyeong-In Yu等人，2022年，OSDI】。这些技术在迭代级别上操作，每次迭代后可以从批处理中移除已完成的请求并加入新请求，从而减少排队延迟和填充浪费，显著提升LLM服务吞吐量。

LLM服务中的内存挑战

尽管细粒度批处理提高了计算效率，但批处理大小仍受限于GPU内存容量，特别是存储KV缓存的空间，使得服务吞కి量受内存限制。

KV缓存规模巨大。KV缓存的大小随请求数量迅速增长。例如，对于OPT-13B模型，单个词元的KV缓存需要800KB空间。由于序列最长可达2048个词元，单个请求的KV缓存可能高达1.6GB。这限制了在几十GB显存的GPU上能同时容纳的请求数量。同时，GPU计算速度的增长快于内存容量的增长，使得内存瓶颈日益突出。

复杂的解码算法。LLM服务提供多种解码算法，每种算法对内存管理有不同要求。例如，并行采样中，同一提示生成的多个样本可以共享提示部分的KV缓存。在更复杂的波束搜索（beam search）中，不同分支（beam）可以共享更大部分的KV缓存，且共享模式会动态变化。

调度未知长度的输入输出。LLM请求的输入和输出长度变化很大。系统必须能适应不同长度的提示，并在输出过程中动态管理内存。当内存耗尽时，系统需要做出调度决策，如删除或换出某些请求的KV缓存。

现有系统的内存管理。由于深度学习框架要求张量存储在连续内存中，现有的LLM服务系统（如【索引31, FasterTransformer, NVIDIA, 2023】和【索引60, Orca: A Distributed Serving System for {Transformer-Based} Generative Models, Gyeong-In Yu等人, 2022, OSDI】）将每个请求的KV缓存存储为连续张量。由于输出长度不可预测，它们会为请求静态分配一个基于最大可能序列长度的内存块。这种预分配方案导致了三种主要的内存浪费：为未来词元预留（reserved）的空间、因过度分配导致的内部碎片（internal fragmentation），以及内存分配器（如伙伴分配器）产生的外部碎片（external fragmentation）。如图3所示，这些浪费使得其他请求无法使用空闲内存。分析表明（如图2），现有系统中实际用于存储词元状态的有效内存可能低至20.4%。尽管内存整理（compaction）是解决碎片化的一种方法，但由于KV缓存规模巨大，在对性能敏感的LLM服务中进行整理是不切实际的。此外，预分配机制也阻碍了特定解码算法中的内存共享。

图3. 现有系统中的KV缓存内存管理。存在三种内存浪费——预留、内部碎片和外部碎片——这些浪费阻止了其他请求装入内存。每个内存槽中的词元代表其KV缓存。注意，相同的词元在不同位置时可以有不同的KV缓存。
图3. 现有系统中的KV缓存内存管理。存在三种内存浪费——预留、内部碎片和外部碎片——这些浪费阻止了其他请求装入内存。每个内存槽中的词元代表其KV缓存。注意，相同的词元在不同位置时可以有不同的KV缓存。

A2 方法细节

本文开发了PagedAttention算法和vLLM服务引擎以应对内存挑战。vLLM系统架构如图4所示，采用一个中心化调度器协调分布式GPU工作节点的执行。KV缓存管理器在PagedAttention的支持下，以分页方式高效管理KV缓存，并根据调度器的指令在GPU工作节点上管理物理KV缓存内存。

图4. vLLM系统概览。
图4. vLLM系统概览。

PagedAttention

PagedAttention算法。为解决内存挑战，本文引入了受操作系统中经典分页思想【索引25, One-level storage system, Tom Kilburn等人, 1962, IRE Transactions on Electronic Computers】启发的PagedAttention算法。与传统注意力算法不同，PagedAttention允许将连续的键（key）和值（value）存储在非连续的内存空间中。具体来说，PagedAttention将每个序列的KV缓存划分为多个KV块（KV blocks）。每个块包含固定数量词元的键和值向量，这个数量表示为KV块大小（$B$）。注意力计算（公式4）可以转化为块级计算，如公式5所示。


其中，$S_{i,k}$ 是在第 $k$ 个KV块上的注意力分数行向量。

PagedAttention的执行过程。在注意力计算期间，PagedAttention内核会分别识别并获取不同的KV块。如图5所示，一个序列的键和值向量分布在三个非连续的物理内存块中。内核会将查询词元（如"forth"）的查询向量 $q_i$ 与每个块中的键向量（如块0中"Four score and seven"的键向量）相乘，计算出注意力分数 $S_{i,k}$，然后用此分数与该块中的值向量相乘以得到最终的注意力输出 $o_i$。PagedAttention算法使得KV块可以存储在非连续的物理内存中，从而支持vLLM中更灵活的分页内存管理。

图5. PagedAttention算法示意图，其中注意力的键和值向量作为非连续块存储在内存中。
图5. PagedAttention算法示意图，其中注意力的键和值向量作为非连续块存储在内存中。

KV缓存管理器

vLLM内存管理器的设计思想。vLLM的内存管理器借鉴了操作系统中的虚拟内存【索引25, One-level storage system, Tom Kilburn等人, 1962, IRE Transactions on Electronic Computers】。操作系统将内存划分为固定大小的页，并将用户程序的逻辑页映射到物理页，使得连续的逻辑页可以对应非连续的物理页。vLLM利用这一思想来管理LLM服务中的KV缓存，将KV缓存组织为固定大小的KV块，类似于虚拟内存中的页。

逻辑块与物理块的分离。一个请求的KV缓存由一系列逻辑KV块表示，随着新词元的生成而从左到右填充。在GPU工作节点上，一个块引擎（block engine）分配一块连续的GPU DRAM并将其划分为物理KV块。KV块管理器为每个请求维护一个块表（block table），记录逻辑KV块到物理KV块的映射。每个表项记录了逻辑块对应的物理块以及已填充的位置数。这种逻辑与物理KV块的分离，使vLLM能够动态增长KV缓存内存，而无需预先保留所有位置的空间，从而消除了现有系统中的大部分内存浪费（如图2所示）。

图6. vLLM中的块表转换。
图6. vLLM中的块表转换。

使用PagedAttention和vLLM进行解码

解码过程示例。图6展示了vLLM在解码单个输入序列时如何执行PagedAttention和管理内存：
1. 初始分配：vLLM不为最大可能序列长度预留内存，仅为提示（prompt）计算所需的KV缓存分配必要的KV块。在此例中，提示有7个词元，vLLM将前2个逻辑KV块（0和1）映射到2个物理KV块（7和1）。在预填充步骤中，vLLM使用传统的自注意力算法生成提示和第一个输出词元的KV缓存，并将前4个词元的KV缓存存入逻辑块0，后3个存入逻辑块1。
2. 第一步自回归解码：vLLM在物理块7和1上使用PagedAttention算法生成新词元。由于最后一个逻辑块中仍有可用空间，新生成的KV缓存被存储在那里，并更新块表中该块的已填充记录。
3. 第二步解码：当最后一个逻辑块已满时，vLLM将新生成的KV缓存存储到一个新的逻辑块中。为此，它会分配一个新的物理块（物理块3），并将其映射关系存储在块表中。

全局解码流程。在每个解码迭代中，vLLM首先选择一组候选序列进行批处理，并为新需要的逻辑块分配物理块。然后，它将当前迭代的所有输入词元（提示阶段请求的所有词元和生成阶段请求的最新词元）连接成一个序列，并输入到LLM中。在LLM计算期间，vLLM使用PagedAttention内核访问以逻辑KV块形式存储的先前KV缓存，并将新生成的KV缓存保存到物理KV块中。在KV块中存储多个词元（块大小>1）可以提高硬件利用率和降低延迟，但较大的块大小会增加内存碎片。vLLM通过从左到右填充所有块，并且仅在所有先前块都已满时才分配新的物理块，从而将每个请求的内存浪费限制在一个块内，有效利用了内存。这使得更多请求可以装入内存进行批处理，从而提高吞吐量。请求完成后，其KV块被释放以供其他请求使用。图7展示了vLLM同时管理两个序列内存的例子，它们的逻辑块被映射到GPU工作节点中由块引擎保留的空间内的不同物理块。

图7. vLLM中同时存储两个请求的KV缓存。
图7. vLLM中同时存储两个请求的KV缓存。

对其他解码场景的应用

vLLM的PagedAttention和分页内存管理同样适用于更复杂的解码场景。

并行采样（Parallel Sampling）。在并行采样中，一个请求包含多个共享相同输入提示的样本，因此提示的KV缓存可以被共享。vLLM通过其分页内存管理可以轻松实现这种共享以节省内存。如图8所示，对于两个输出的并行解码，vLLM仅为提示状态保留一份空间，两个序列的提示逻辑块（逻辑块0和1）都映射到相同的物理块（物理块7和1）。为管理共享，vLLM为每个物理块引入了引用计数。在此例中，物理块7和1的引用计数均为2。在生成阶段，vLLM为需要被多个序列修改的物理块实现了块粒度的写时复制（copy-on-write）机制。当样本A1需要写入其最后一个逻辑块时，vLLM发现对应物理块1的引用计数大于1，于是分配一个新物理块3，复制物理块1的内容，并将原块的引用计数减1。之后当样本A2写入物理块1时，由于引用计数已为1，它可以直接写入。通过这种方式，vLLM在多样本间共享了大部分提示KV缓存，显著减少了内存使用。

图8. 并行采样示例。
图8. 并行采样示例。

波束搜索（Beam Search）。波束搜索【索引49, Sequence to sequence learning with neural networks, Ilya Sutskever等人, 2014, NeurIPS】不仅共享初始提示块，还能在不同候选项之间共享其他块，且共享模式会动态变化。图9展示了vLLM如何管理一个波束宽度 $N=4$ 的波束搜索示例。在虚线所示的迭代之前，所有候选项共享块0（提示）。在后续迭代中，排名靠前的4个候选项都源于候选项1和2。原候选项0和3被丢弃，它们的逻辑块被释放，相应物理块的引用计数减少。引用计数为0的物理块（块2、4、5、8）被vLLM释放。然后，vLLM分配新的物理块（块9-12）来存储新候选项的KV缓存。vLLM通过物理块共享显著减少了传统系统中频繁的KV缓存内存复制开销。

图9. 波束搜索示例。
图9. 波束搜索示例。

共享前缀（Shared Prefix）。在许多应用中，用户请求的提示（prompt）包含一个共享的系统提示（system prompt），如图10所示。vLLM可以预先存储这些共享前缀的KV缓存，从而减少重复计算。这通过为一组预定义的共享前缀保留一组物理块来实现，类似于操作系统处理共享库。带有共享前缀的用户输入可以将其逻辑块映射到这些缓存的物理块（最后一个块标记为写时复制），提示阶段的计算只需在用户的任务输入上执行。

图10. 机器翻译的共享提示示例。示例采纳自【索引5, Language models are few-shot learners, Tom Brown等人, 2020, NeurIPS】。
图10. 机器翻译的共享提示示例。示例采纳自【索引5, Language models are few-shot learners, Tom Brown等人, 2020, NeurIPS】。

混合解码方法。vLLM能够同时处理具有不同解码偏好的请求，这是现有系统难以高效实现的。vLLM通过一个将逻辑块转换为物理块的通用映射层，隐藏了不同序列之间复杂的内存共享。LLM及其执行内核只看到每个序列的物理块ID列表，无需处理跨序列的共享模式。这扩大了对具有不同采样需求的请求进行批处理的机会，最终提高了系统的总吞吐量。

调度与抢占

调度策略。当请求流量超过系统容量时，vLLM采用先到先服务（FCFS）的调度策略，确保公平性并防止饥饿。当需要抢占请求时，vLLM确保最早到达的请求被优先服务，而最晚到达的请求被首先抢占。

内存耗尽时的处理。当vLLM耗尽GPU物理块时，它需要决定驱逐哪些块以及如何恢复它们。vLLM采用“全有或全无”（all-or-nothing）的驱逐策略，即要么驱逐一个序列的所有块，要么一个都不驱逐。此外，同一请求内的多个序列（如波束搜索中的候选项）作为一个序列组被“组调度”（gang-scheduled），它们总是被一起抢占或重新调度。对于恢复被驱逐的块，vLLM考虑两种技术：
1. 交换（Swapping）：将被驱逐的块复制到CPU内存中的交换空间。vLLM包含一个CPU块分配器来管理交换到CPU RAM的物理块。当一个请求完成并释放其块后，被抢占序列的块会被重新加载回GPU内存以继续处理。
2. 重计算（Recomputation）：当被抢占的序列被重新调度时，重新计算其KV缓存。重计算的延迟可能远低于原始延迟，因为已生成的词元可以与原始用户提示连接成一个新的提示，其所有位置的KV缓存可以在一个提示阶段迭代中生成。
这两种方法的性能取决于CPU RAM和GPU内存之间的带宽以及GPU的计算能力。

分布式执行

模型并行支持。vLLM支持Megatron-LM风格的张量模型并行策略【索引47, Megatron-lm: Training multibillion parameter language models using model parallelism, Mohammad Shoeybi等人, 2019】，以在分布式GPU上执行大型模型。该策略遵循单程序多数据（SPMD）执行模式，其中线性层被分区以进行分块矩阵乘法，GPU通过all-reduce操作同步中间结果。

分布式内存管理。即使在模型并行执行中，每个模型分片也处理相同的输入词元集，因此需要相同位置的KV缓存。因此，vLLM在中心化调度器中设有一个单一的KV缓存管理器，如图4所示。不同的GPU工作节点共享该管理器以及从逻辑块到物理块的映射。在每个步骤中，调度器首先准备包含批处理中每个请求的输入词元ID和块表的消息，然后将此控制消息广播给GPU工作节点。工作节点根据块表读取KV缓存并执行模型，期间通过all-reduce同步中间结果。最后，工作节点将本轮迭代采样的词元发送回调度器。这种设计下，GPU工作节点无需在内存管理上进行同步。

A4 实验环境

A5 实验结果

本节评估vLLM在多种工作负载下的性能，主要关注服务吞吐量，通过测量不同请求速率下的归一化延迟（请求端到端延迟除以其输出长度）来评估。

A6 结论

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

参考文献引用说明