基于PagedAttention的大型语言模型高效内存管理研究学术报告

本报告旨在介绍加州大学伯克利分校Woosuk Kwon、Zhuohan Li、Siyuan Zhuang、Ying Sheng、Lianmin Zheng等研究人员于2023年10月23日发表于计算机系统领域顶级会议“第29届操作系统原理研讨会”（SOSP ‘23）的研究论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》。该研究提出了一种革命性的注意力算法——分页注意力（PagedAttention），并基于此构建了高性能的大型语言模型（Large Language Model， LLM）服务系统vLLM，旨在解决LLM推理服务中因动态键值缓存（Key-Value Cache， KV Cache）内存管理低效而导致的吞吐量瓶颈问题。

学术背景与研究目的 随着以GPT、PaLM、LLaMA等为代表的大语言模型在编程助手、通用聊天机器人等应用中的爆炸式增长，提供这些模型的高效云服务变得至关重要。然而，部署LLM服务成本高昂，其核心挑战在于提高吞吐量以降低单次请求的成本。LLM基于自回归的Transformer架构，以逐个生成令牌（Token）的方式运行。为了提高硬件利用率，服务系统通常会将多个用户请求批量处理。然而，这种批处理能力受到GPU内存容量的严重制约，尤其是用于存储每个请求上下文信息的键值缓存。KV Cache随生成令牌的数量动态增长，其大小在模型中占比可达30%（例如，一个13B参数的模型，单个请求最大序列长度下KV Cache可占1.6GB）。现有的服务系统（如FasterTransformer、Orca）由于深度学习框架的限制，需要为每个请求的KV Cache预分配一段连续的存储空间，最大可达模型支持的最大序列长度。这种管理方式导致了严重的内存碎片化（内部碎片和外部碎片）和内存冗余（无法在不同序列间共享KV Cache），使得实际用于存储有效令牌状态的内存利用率极低（论文中观测仅为20.4%-38.2%），从而限制了批处理规模，最终限制了系统吞吐量。为此，本研究旨在设计一种全新的内存管理机制，以近乎零浪费的方式管理KV Cache，并支持灵活的跨序列内存共享，从而显著提升LLM服务吞吐量。

详细研究流程与方法 本研究主要包含三个核心环节：1）问题剖析与量化；2）PagedAttention算法设计与vLLM系统构建；3）综合性能评估与消融研究。

第一环节：问题剖析与量化 研究团队首先深入分析了现有LLM服务系统（以FasterTransformer和Orca为代表）在内存管理上的根本缺陷。他们指出，KV Cache具有两个独特性质：其生命周期和长度在请求开始前未知，并且在不同请求或同一请求的不同输出序列间存在大量共享机会。而现有系统采用“连续内存块预分配”策略，这直接导致了三种内存浪费：为未来令牌预留（Reserved）的空间长时间占用内存；因实际长度远小于最大预留长度而产生的内部碎片（Internal Fragmentation）；以及因预分配块大小不一导致的内存分配器外部碎片（External Fragmentation）。此外，现有系统无法实现不同序列间的KV Cache共享。研究团队通过一个分析实验（文中Figure 2）量化了这些浪费，结果显示在现有系统中，仅有约1/5到1/3的内存真正用于存储令牌状态，其余均为各种形式的浪费。

第二环节：算法设计与系统构建 针对上述问题，研究团队借鉴操作系统中的经典虚拟内存和分页技术，提出了PagedAttention算法。这是本研究的核心创新点。与需要在连续内存上操作的传统注意力算法不同，PagedAttention允许将序列的KV Cache划分为固定大小的“KV块”（KV Block），这些块可以非连续地存储在物理内存中。注意力计算被相应地改造为块级计算。在此算法基础上，研究团队构建了vLLM服务引擎。

vLLM的工作流程如下： 1. 内存管理层：vLLM的设计灵感源于操作系统虚拟内存。系统将GPU DRAM划分为固定大小的物理KV块。每个用户请求的KV Cache被组织为一组逻辑KV块。一个中心化的“块表”（Block Table）维护着逻辑块到物理块的映射关系，并记录每个逻辑块中已填充令牌的数量。这种设计使得vLLm能够按需动态分配物理块，而无需为整个最大序列长度预留连续内存。 2. 服务执行流程： * 调度：中央调度器基于先到先服务（FCFS）策略选择一批请求。 * 内存准备：根据选中的请求，调度器确定每个序列需要的新逻辑块，并通过块管理器为它们分配物理块（若物理块不足，则触发基于“全或无”策略的请求抢占和交换/重计算机制）。 * 模型执行：调度器将输入令牌ID及每个序列对应的块表信息广播给所有GPU工作节点。 * PagedAttention内核计算：在各GPU节点上，定制的PagedAttention CUDA内核根据块表信息，从非连续的物理块中读取历史KV Cache，进行注意力计算，并将新生成的KV Cache写入对应的物理块中。这些内核进行了高度优化（如融合的块读取与注意力计算、融合的块写入等），以减少间接寻址带来的开销。 * 生成与回收：迭代完成后，生成的令牌返回给用户。当请求完成时，其占用的所有物理KV块被释放。

vLLM对复杂解码算法的支持是其另一大亮点。系统通过“分叉（Fork）”、“追加（Append）”、“释放（Free）”三个基本操作，配合块级引用计数和写时复制（Copy-on-Write）机制，高效支持了多种高级解码场景： * 并行采样：同一提示词（Prompt）生成的多个输出序列可以共享提示词对应的物理KV块，仅在不同序列产生分歧时才进行复制，显著节省内存。 * 束搜索：不同的束（Beam）候选者之间可以共享大量历史KV块（包括部分生成内容），共享模式随解码过程动态演化。相比现有系统需要频繁复制整个KV Cache，vLLm仅在写共享块时才复制一个块的数据，开销大幅降低。 * 共享前缀：对于拥有相同系统提示词或任务描述的请求，vLLm可以像操作系统处理共享库一样，预先计算并缓存共享前缀的KV块，后续请求直接映射即可，避免了重复计算。 * 混合解码：vLLm能够同时高效批处理使用不同解码算法的请求，这是现有系统难以做到的。

第三环节：综合性能评估 研究团队在多种模型（OPT-13B/66B/175B, LLaMA-13B）和真实数据集（ShareGPT, Alpaca）构成的合成工作负载上，对比评估了vLLm与现有最优系统（FasterTransformer及论文自行实现的Orca多个变体）的性能。关键实验设计与结果如下： * 基础采样性能：在ShareGPT和Alpaca数据集上，vLLm在维持相同延迟水平下，能支撑的请求率（吞吐量）比Orca（Oracle，知晓实际输出长度的理论上界）高1.7-2.7倍，比Orca（Max，总是预留最大长度）高2.7-8倍，比FasterTransformer最高可达22倍。这直接归因于vLLm高效的内存管理允许更大的批次规模（如Figure 13所示，vLLm同时处理的请求数可达Orca的2-4倍）。 * 复杂解码算法性能：在并行采样和束搜索实验中，随着并行序列数或束宽的增加，vLLm相对于Orca的优势更加显著。在束宽为6时，vLLm的吞吐量提升达到2.3倍。内存共享节省方面，束搜索最高可节省55.2%的内存（ShareGPT数据集上可达66.3%）。 * 共享前缀场景：在机器翻译任务中，当输入提示共享一个包含5个示例的较长前缀时，vLLm的吞吐量比Orca（Oracle）高出3.58倍，有效利用了前缀缓存。 * 消融研究： * 内核开销：尽管PagedAttention内核因块表访问和分支处理引入了额外开销，其延迟仅比高度优化的FasterTransformer注意力内核高20-26%，但端到端吞吐量优势巨大。 * 块大小影响：实验确定了默认块大小16是一个在GPU并行利用率和内存碎片间取得良好平衡的选择。 * 交换与重计算：对比了当GPU内存不足时，将KV块换出到CPU内存（交换）与直接丢弃并在需要时重新计算（重计算）两种恢复机制。研究发现，对于较小的块大小，重计算更有效；对于中等块大小（16-64），两者性能相当。vLLm支持这两种机制。

主要研究结果与结论 本研究通过系统的实验验证，取得了以下关键结果： 1. 内存利用率接近理论最优：vLLm成功消除了由预分配和连续存储导致的内存碎片，在评估中实现了接近零浪费的KV Cache内存管理。 2. 吞吐量显著提升：在各种模型、工作负载和解码算法下，vLLm将LLM服务吞吐量提升了2至4倍。 3. 高效支持复杂解码：通过块级内存共享和写时复制机制，vLLm高效支持了并行采样、束搜索等高级解码功能，大幅减少了内存冗余。 4. 灵活应对动态负载：结合迭代级调度和基于FCFS的抢占策略，vLLm能够有效处理输入输出长度未知且变化的请求流。

基于上述结果，本研究得出结论：通过将操作系统中的虚拟内存和分页思想创造性地应用于LLM服务中的KV Cache管理，所提出的PagedAttention算法和vLLm系统，能够从根本上解决现有系统的内存效率瓶颈，为高吞吐、低成本的LLM服务提供了一种强大且灵活的解决方案。

研究价值与亮点 本研究的科学价值与应用价值极高。在科学上，它成功地将经典系统领域（操作系统）的思想迁移并适配到新兴的AI系统领域，为解决特定领域（自回归生成模型服务）的性能瓶颈提供了一个跨学科融合的优秀范例。在应用上，vLLm作为一个开源系统，已被业界广泛采用，直接推动了LLM服务基础设施的进步，降低了服务成本，使得更多用户能够访问强大的LLM能力。

本研究的亮点在于： 1. 核心算法创新：提出PagedAttention，这是第一个允许在非连续内存上进行注意力计算的算法，是后续所有优化的基石。 2. 系统设计精巧：将虚拟内存、分页、块表、写时复制、引用计数等OS概念无缝融入LLM服务上下文，设计出完整且高效的vLLm系统架构。 3. 性能提升显著且普适：实验证明，其带来的2-4倍吞吐量提升在不同模型规模、序列长度和解码复杂度下均成立。 4. 解决实际问题：精准定位并量化了LLM服务中的核心内存瓶颈，提出的方案具有明确的工程落地价值和广泛的行业影响力。

其他有价值内容 论文还讨论了将虚拟内存和分页技术应用于其他GPU工作负载的潜力与局限性，指出该技术特别适用于那些内存需求动态变化且性能受内存容量限制的场景（如LLM服务），而对于计算密集型或张量形状固定的工作负载则可能不适用甚至带来额外开销。同时，论文强调了vLLm在应用经典思想时所做的LLM特定优化，例如“全或无”的换出策略和利用请求语义进行重计算的恢复机制，这些是通用操作系统所不具备的，体现了领域特定系统设计的深度。