Post

PagedAttention论文阅读

Posted
By jinhuix
31 min read
PagedAttention论文阅读

本文为论文 Efficient Memory Management for Large Language Model Serving with PagedAttention 的阅读笔记。

  1. 问题背景与挑战：大语言模型（LLMs）在高吞吐量的推理场景下，需要一次处理足够多的请求（即批量化）。但目前的系统在这一点上面临困难，原因主要是每个请求的KV缓存占用大量内存，并且其大小是动态变化的。这导致：内存碎片严重；存在重复的数据副本；限制了批处理请求的能力。
  2. 核心创新点：为解决上述问题，作者提出了PagedAttention，一种受操作系统虚拟内存分页机制启发的注意力机制，用于更高效地管理KV缓存。基于此机制，作者实现了一个LLM服务系统——vLLM。
  3. 系统特点：vLLM实现了两大关键目标：KV缓存几乎零浪费；在不同请求之间实现KV缓存的灵活共享，从而进一步节省内存。
  4. 性能优势：相比现有的主流系统（如FasterTransformer和Orca），vLLM在维持相同延迟的前提下，将吞吐量提升了2到4倍。提升效果在以下场景中更明显：序列长度较长；模型规模较大；解码算法更复杂。

1 Introduction

1. LLM推理服务的重要性与挑战

  • 随着GPT、PaLM等大型语言模型的广泛应用（如编程助手、通用聊天机器人），云服务商纷纷提供相关服务。
  • 然而，LLM推理成本极高，比传统查询高出约10倍，主要消耗在GPU资源上，因此提升吞吐量、降低单请求成本非常关键。

2. 推理流程中的KV缓存问题

  • Transformer模型通过自回归方式逐token生成输出，导致生成过程序列化且内存密集，计算资源反而被低效利用。
  • 批处理请求是提高吞吐量的主要方式，但受到内存限制：模型权重占65%的GPU内存；动态的KV缓存占近30%，是关键瓶颈；激活等其他部分占比很小；所以，KV缓存的管理效率决定了批量处理的上限。

PagedAttention-1

3. 现有系统的KV缓存管理问题

  • 多数系统使用连续内存存储KV缓存，这是由于主流深度学习框架的要求。
  • 但KV缓存具有动态性（生成时增长、完成后释放）和生命周期不可预测等特性，使得连续内存管理产生两大问题：
    1. 内外碎片严重：预分配大块连续空间（如2048 token），但实际请求短时会造成内部碎片；不同请求长度不同，导致空间分配不均，形成外部碎片；实验显示，现有系统中只有 20.4% ~ 38.2% 的KV缓存内存被真正用于存储token状态。
    2. 无法共享内存：如parallel sampling和beam search这类高级解码方式会生成多个序列，本可以共享一部分KV缓存；但因采用连续内存分配，不支持这些序列间的缓存共享。

PagedAttention-2

4. 解决方案：PagedAttention

  • 借鉴操作系统的虚拟内存和分页技术，提出PagedAttention：

    • 将KV缓存划分为若干小块（blocks），每块保存固定数量token的KV对；

    • 块之间不要求连续存储，像操作系统中的“页”一样动态分配和回收；

    • 优势包括：通过小块分配，缓解内部碎片；所有块大小相同，消除外部碎片；支持在同一请求不同序列之间，甚至跨请求的缓存块共享。

  • 基于PagedAttention构建了高吞吐量的LLM推理引擎 vLLM，其特性包括：

    • 支持主流大模型（GPT、OPT、LLaMA）；可支持单GPU显存无法容纳的大模型；结合块级内存管理和抢占式调度策略，极大提升性能；在不影响模型精度的前提下，吞吐量提升2-4倍，尤其在处理长序列、大模型、复杂解码时优势更明显。

2 Background

1. Transformer-Based LLMs 的基本生成原理

  • 语言建模的目标是估计一个 token 序列的联合概率，可以按时间顺序拆分为多个条件概率的乘积（自回归分解）。
  • Transformer 架构是主流大语言模型的核心框架，其主要模块是自注意力层（Self-Attention），用于建模当前 token 对之前 token 的依赖关系。
  • 自注意力机制会为每个位置计算 query、key 和 value 向量，并根据 query 与所有前面 key 的相似度加权求和得到当前的输出。
  • 除了注意力，Transformer 的其余模块（如前馈层、LayerNorm、残差连接等）都是逐位置独立计算的，因此可以并行处理。

2. LLM 推理服务流程与 KV Cache 的作用

  • LLM 服务通常以“补全 API”或“聊天对话”形式提供，用户输入 prompt，模型按 token 逐个生成响应。
  • 推理流程可以分为两阶段：
    1. Prompt 阶段（并行）：输入用户的全部 prompt；计算第一个输出 token 的概率；同时构建对应的 key 和 value 向量（KV Cache）；此阶段可以用矩阵乘法并行完成，因此效率较高。
    2. 自回归生成阶段（串行）：
      • 每次只生成一个 token；当前 token 的生成依赖所有之前的 KV；旧的 KV 向量在前面的迭代中已缓存，新 token 只需计算当前 KV；
      • 由于依赖前文，无法并行处理，只能矩阵-向量乘法，效率较低；这是当前 LLM 服务中最耗时、最受限的阶段，容易导致 GPU 计算资源闲置、成为内存瓶颈。
  • KV Cache 的关键性：
    • 是自注意力每步都要用到的历史上下文信息；每个 token 的 KV 与其上下文有关，因此在同一序列中，同一个词出现在不同位置时，其 KV 也不同；因此 KV Cache 是动态增长的，不能共享或复用。

3. 提升推理吞吐量的批处理技术（Batching）

  • 批处理是提升 LLM 服务 GPU 利用率的关键手段：模型权重在多个请求间共享；大批量可以摊薄加载模型等开销。
  • 批处理的难点在于请求的异步性与长度差异：
    1. 请求到达时间不同，简单聚合会导致早到的请求等待，排队延迟增加；
    2. 请求长度差异大，统一填充（padding）会浪费大量计算和显存。
  • 解决方案是细粒度批处理（如cellular batching、迭代级调度）：
    • 按“生成迭代”而非“请求”来组织批次；每一迭代后，完成的请求可以剔除，新的请求可以立刻插入；请求只需等一次迭代即可加入推理，减少等待；
    • 使用定制GPU kernel 避免 padding，进一步提升效率；最终带来更高的吞吐量、更少的显存浪费。

3 Memory Challenges in LLM Serving

1. KV缓存主导内存瓶颈

  • 即使通过细粒度调度提升了计算资源利用率，批量请求数量仍受限于GPU内存容量，特别是KV cache占用。
  • KV缓存随请求和序列长度迅速膨胀：
    • 例如：OPT-13B模型中，每个token的KV缓存约为 800KB；若生成2048个token，单个请求就可能消耗1.6GB内存；即便用完40GB或80GB的GPU，也只能同时处理几十个请求；
  • 随着硬件发展，计算能力提升速度远快于显存容量（如A100到H100），内存瓶颈会日益严重。

2. 解码算法复杂性加剧内存管理难度

  • 不同解码策略对KV缓存共享有不同要求：
    • 采样（sampling）：多个样本可共用prompt部分的KV缓存（节省约12%内存）；
    • 自回归阶段因输出不同，不可共享KV；
    • Beam Search：不同beam间可以共享较大比例的KV缓存（最多节省55%）；
    • 内存共享的程度与解码算法密切相关，并且共享关系随解码过程不断演化。

3. 输入/输出长度不可预知带来调度难题

  • 请求输入/输出长度各异，系统必须支持动态增长的KV缓存需求；
  • 在生成阶段，如果KV增长超出预期，会抢占其他请求的显存资源；
  • 系统需具备动态调度能力，如移除或转存部分请求的KV缓存。

4. 现有系统的内存管理机制存在浪费

  • 当前主流框架（如PyTorch、TensorFlow）要求张量是连续内存块，导致：
    • LLM服务系统将整个KV缓存为每个请求预分配为连续内存；
    • 预留的KV缓存大小基于最大可能序列长度（如2048 token），即使实际很短；
  • 这带来三种内存浪费：
    1. 未来token的预留空间：整个生命周期都占用内存，即使暂时未使用；
    2. 内部碎片：过度分配但未使用的空间，直到生成结束后才能确认；
    3. 外部碎片：由底层内存分配器（如buddy allocator）引起的不连续空间，永远无法被生成token使用；
  • 实验表明：现有系统实际有效利用的KV缓存内存仅为20.4%~38.2%，其余被浪费。
  • 现有方法无法有效缓解问题
    • 内存压缩理论上可解决碎片，但在大模型推理中代价过高，影响性能；即便压缩，也无法解决因连续分配导致的KV缓存无法共享的问题；所以现有系统仍缺乏可扩展、灵活的KV内存管理方案。

PagedAttention-3

4 Method

1. vLLM 系统架构概述

  • vLLM 由一个中央调度器和多个GPU worker组成；
  • 中央调度器：负责任务调度和协调，决定哪些请求运行、KV cache如何管理；
  • KV Cache 管理器：位于GPU worker侧，按照中央调度器的指令进行实际的物理内存分配和回收，采用“分页方式”管理KV缓存；
  • 核心创新：引入类操作系统中“虚拟内存分页（paging）”的思想来重新设计注意力机制，使KV缓存无需连续分配，解决内存碎片与共享问题。

PagedAttention-4

4.1 PagedAttention

PagedAttention 是受操作系统分页机制启发而设计的一种块式注意力算法，具有以下核心特性：

1. 非连续KV缓存

  • 不像传统Attention机制那样要求KV向量必须是连续存储；PagedAttention 将每个序列的KV Cache 分成多个 KV块，每块包含 固定数量的token 的 Key 和 Value 向量；每个块可以独立存储在物理内存的任意位置，无需紧挨。

2. 块级别注意力计算

  • 假设块大小为 B，第 j 块的 Key 和 Value 分别为：$K_j = (k_{(j-1)B+1}, …, k_{jB})$，$V_j = (v_{(j-1)B+1}, …, v_{jB})$
  • 查询向量 $q_i$ 对每个块 $K_j$ 计算注意力得分 $A_{ij}$，再与对应的 $V_j$ 相乘，累加得到输出：$o_i = \sum_{j=1}^{\lceil i/B \rceil} V_j A_{ij}^\top$

3. 运行示例

  • 某次注意力计算中，输入的 Query 是单词 “forth”；
  • 系统从非连续的三块物理内存中读取 Key 和 Value 向量，如第 0 块对应的内容是 “Four score and seven”；
  • 内核逐块进行注意力打分和输出计算，无需一次性载入完整历史KV序列。

4.2 KV Cache Manager

vLLM 的 KV Cache 管理思想源自操作系统中的虚拟内存机制：

  • 虚拟内存类比：操作系统将用户进程的逻辑内存划分为页（pages），并映射到非连续的物理页，用户仍然感觉在使用“连续”内存。vLLM 也采用类似方式来管理 KV cache。
  • KV块机制：vLLM 将每个请求的 KV cache 分成固定大小的 KV 块（类似虚拟页），每个请求的 KV cache 逻辑上是一个序列，每生成一部分 token，就在右侧顺序填充一个 KV 块。
  • 物理块分配：GPU worker 的 block engine 预先分配一大段 GPU 显存，切成多个物理 KV 块，用来接纳不同请求的 KV 块。
  • 映射表（Block Table）：每个请求维护一个逻辑块 → 物理块的映射表，每个条目记录：映射到的物理块编号；当前填入了多少 token（用于判断是否还有空间继续写入）。
  • 动态增长、无需预留全部内存：与传统系统不同，vLLM 不需为每个请求预留整个最大长度的内存，而是随着 token 生成过程逐步动态分配新块，有效降低内存浪费。

4.3 Decoding with PagedAttention and vLLM

通过一个解码实例，展示 vLLM 如何配合 PagedAttention 在推理过程中动态分配与管理 KV cache。

1. 解码步骤详解

  • Prompt 阶段：输入 prompt 含 7 个 token；vLLM 映射出两个逻辑块（block 0、1），并分配两个物理块（例如物理块 7 和 1）；将前 4 个 token 写入 block 0，接下来的 3 个写入 block 1；block 1 尚有空位，为未来生成的 token 预留。
  • 第一次生成：使用 PagedAttention，加载已有的物理块 7 和 1；新生成的 token 的 KV 写入 block 1 的剩余槽位；映射表更新 block 1 的填充数量。
  • 第二次生成：由于 block 1 已满，分配一个新的逻辑块（block 2）；对应分配一个新的物理块（如 block 3），并建立映射；KV cache 写入新的块中。

PagedAttention-5

2. 多请求批处理中的行为（全局调度逻辑）

  • 每一轮迭代中，vLLM：筛选一批可以同时处理的请求；分配它们本轮所需的逻辑块并映射到物理块；将所有参与本轮推理的 token 组成一个大批次输入；使用 PagedAttention 访问历史 KV 块并更新最新的 KV cache。

3. 性能考虑：块大小的取舍

  • 如果一个 KV 块能容纳多个 token：
    • 优点：GPU 并行计算更高效（提升吞吐、降低延迟）；
    • 缺点：可能会导致更大的未使用空间（碎片）；

4. 处理两个序列示例

  • 展示两个请求的逻辑 KV 块如何映射到非连续的物理块；
  • 块之间不需要邻接，每个物理块空间都被充分利用；
  • 逻辑块之间、请求之间不会产生干扰。

4.4 Scheduling and Preemption

1. 调度策略：确保公平性与可抢占性

  • vLLM 采用 先来先服务（FCFS） 策略，保证请求公平处理，避免饥饿。
  • 当 GPU 资源不足，必须进行请求抢占时：优先保留早到的请求，后到的请求先被抢占；符合公平性原则，有利于长 prompt 或长序列的处理连续性。

2. 动态 KV 缓存增长带来的挑战

  • LLM 请求的 prompt 长度不定，输出长度不可预测；
  • 多个请求同时运行，生成过程会不断扩大其 KV cache；
  • 当 GPU 显存空间（KV cache 空间）不足，必须决定：
    1. 抢占哪些块？
    2. 如何恢复被抢占的块？

3. 块抢占策略：以序列为最小单位

  • KV cache 访问具有“全序列访问”特性，即一个序列在生成中会反复使用其所有历史 KV 块；因此采用 全有或全无（all-or-nothing）抢占策略：一个序列的所有块要么一起保留，要么一起抢占；
  • 对于如 beam search 这类多条候选序列共享 KV cache 的情形，以“序列组”为单位调度和抢占，保证内存一致性。

4. 块恢复策略：两种可选机制

  • Swapping（交换到CPU）

    • 被抢占的 KV 块从 GPU 显存拷贝到 CPU 内存中；vLLM 包含一个 CPU block allocator，管理这些交换出去的块；当 GPU 内存不足时，选择一个序列组将其 KV cache 整体迁移到 CPU；

    • 被抢占后，vLLM 会暂停接收新请求，直到现有请求完成并释放显存；一旦 GPU 中释放出足够空间，被抢占序列的 KV cache 重新调入继续生成；

    • 此机制下 CPU RAM 交换区大小被限制在 GPU 中分配给 KV cache 的大小之内，不会无限增长。

  • Recomputation（重计算）

    • 不保存 KV cache，而是在序列被重新调度时重新计算 KV cache；重计算只需将已生成 token 拼接到 prompt 中，在 prompt phase 中一次性并行计算；

    • 相比于 autoregressive phase 的逐 token 生成，重计算更高效；适用于对 latency 要求不高、GPU 计算资源充裕的场景。

5 Evaluation

5.1 Experimental Setup

1. 评估模型与硬件环境：

  • 使用 OPT 系列（13B、66B、175B 参数）和 LLaMA 13B 作为评估模型；所有实验在 Google Cloud 的 A2 实例（NVIDIA A100 GPU）上完成；

2. 工作负载构建方法：

  • 数据来源包括 ShareGPT（用户对话）和 Alpaca（GPT-3.5 自动生成的指令数据）；ShareGPT 的输入平均比 Alpaca 长 8.4 倍，输出也长 5.8 倍，且波动更大；采用泊松分布合成请求到达时间，不同请求率模拟不同负载。

3. 对比系统：

  • FasterTransformer：是一个为低延迟优化的分布式推理引擎；
  • Orca：是目前 SOTA 的高吞吐量推理系统，但未开源，论文复现了三种近似版本

PagedAttention-6

5.2 Basic Sampling

作者评估了 vLLM 在基本采样场景（每个请求生成一个样本）下的性能，分别在三种模型和两种数据集上进行实验。

  • 吞吐能力显著提升：在 ShareGPT 数据集上，随着请求速率增加，系统延迟先缓慢增长，再突然爆炸（系统过载所致）。vLLM 能在不显著增加延迟的前提下，支持比 Orca (Oracle) 高 1.7–2.7 倍的请求速率，比 Orca (Max) 高 2.7–8 倍。这是因为 vLLM 的 PagedAttention 能更高效地管理内存，从而允许更大的批量处理。

PagedAttention-7

  • 并发能力显著提升：vLLM 对于 OPT-13B 同时处理的请求数是 Orca (Oracle) 的 2.2 倍，是 Orca (Max) 的 4.3 倍。相比之下，FasterTransformer 的请求承载能力更低，vLLM 的最大请求速率是其 22 倍，因为后者缺乏精细调度机制且内存管理低效。
  • 在 Alpaca 数据集上的表现：趋势与 ShareGPT 类似，但 vLLM 相对于 Orca 的优势略小，尤其是图 12(f) 中。原因在于：
    • 使用 OPT-175B 时 GPU 显存充足；
    • Alpaca 的输入/输出序列较短，内存压力较小；
    • Orca 的内存低效问题在此情况下影响较小，系统瓶颈转向计算能力而非内存容量。

PagedAttention-8

5.3 Ablation Studies

1. Impact of Block Size

block size（块大小）选择对性能影响显著。块太小会降低并行利用率，太大则引发内存碎片、降低共享效率。在 ShareGPT 数据集下，块大小为 16 到 128 表现较佳；而在 Alpaca 数据集（序列更短）中，过大的块严重拖慢性能。综合考虑 GPU 利用率与碎片控制，vLLM 默认设置块大小为 16。

2. Comparing Recomputation and Swapping

在小块场景下，swapping 因频繁的小块数据传输造成 PCIe 带宽瓶颈，性能下降明显；而 recomputation 的延迟较稳定、不依赖块大小，因此更适用于小块。对于中等块大小（16–64），两者表现相近；swapping 仅在块较大时略占优，但其延迟始终高于 recomputation 的 5 倍左右。因此，vLLM 可根据工作负载灵活选择策略。

6 Other

vLLM 通过理解 LLM 推理的内存访问模式与性能瓶颈，重新构造并优化了分页机制。虽然此类方法不适合所有 GPU 工作负载，但在满足类似内存动态增长和共享需求的场景中具有广泛潜力。

1. 适用性受限于特定场景： 虚拟内存和分页技术在 LLM 推理中效果显著，是因为其输出长度动态且常受限于 GPU 显存。但这种需求并不普遍。例如，在 DNN 训练中，张量大小通常固定，可预先优化内存分配；在传统 DNN 推理中，性能往往由计算而非内存瓶颈决定，因此使用分页反而会带来额外开销，可能降低性能。

2. vLLM 做了特定优化以适配 LLM 特性： vLLM 并非简单照搬操作系统的分页机制，而是结合 LLM 服务的语义进行了重设计。例如：

  • 实现了“全或无”式的 swap-out 策略，因为每次请求推理都需要其所有 token 的 KV 状态同时驻留在显存中；
  • 引入了 重计算（recomputation） 作为替代传统 swap-in 的恢复方式，进一步减少了内存占用（而这在 OS 中不可行）；
  • 将内存访问操作与 attention 等 GPU kernel 融合，减小内存间接访问带来的开销。
大模型
KVCache 大模型推理 vLLM1
This post is licensed under CC BY 4.0 by the author.