PagedAttention 论文精读笔记

Efficient Memory Management for Large Language Model Serving with PagedAttention
SOSP 2023 · Koblenz, Germany
作者:Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng 等 机构:UC Berkeley, Stanford, UC San Diego
SOSP'23 vLLM KV Cache Paging Copy-on-Write Memory Management
论文核心贡献:提出 PagedAttention 算法,将操作系统经典的虚拟内存与分页技术引入 LLM 推理的 KV Cache 管理。在之上构建 vLLM serving 系统,实现 (1) 接近零浪费的 KV Cache 内存管理;(2) 灵活的 KV Cache 共享机制。相比 FasterTransformer 和 Orca,吞吐量提升 2-4×,且序列越长、模型越大、解码算法越复杂,优势越明显。

📋 精读目录

一、问题背景:为什么 KV Cache 是瓶颈

1.1 LLM 推理的成本焦虑

GPT、PaLM 等大模型催生了编程助手、通用聊天机器人等应用。云厂商竞相提供托管服务,但 LLM 推理极其昂贵——处理一个 LLM 请求的成本是传统关键词查询的 10 倍。提升吞吐、降低单请求成本成为核心诉求。

1.2 自回归生成的内存瓶颈

LLM 的核心是自回归 Transformer:每次基于已生成的 token 序列预测下一个 token。这个顺序生成过程使得工作负载成为 memory-bound,GPU 算力被严重低估用。

提升吞吐的关键是 batching——将多个请求一起处理。但 batching 的瓶颈在于 GPU 内存容量,尤其是 KV Cache 的存储空间。

NVIDIA A100 40GB
模型参数占 ~65%
(26GB for 13B model)
KV Cache
占 ~30%+
动态增长/收缩
Activations
仅占极小比例
临时张量
核心观察:模型权重是静态的,activation 很小且临时。因此,KV Cache 的管理方式直接决定了最大 batch size,进而决定了系统吞吐

1.3 现有系统的两大根本性缺陷

论文对 FasterTransformer [31] 和 Orca [60] 等现有系统做了详细剖析,发现它们将每个请求的 KV Cache 存储在连续的内存空间中。这在传统深度学习框架中很常见,但 KV Cache 有两个独特特征:

  1. 动态增长收缩:随模型生成新 token,KV Cache 不断增长,其生命周期和最终长度事先未知
  2. 共享机会:在并行采样、beam search 等场景中,多个序列可以部分共享 KV Cache

这两个特征使得"连续存储"方案在两方面严重低效:

缺陷一:严重的内存碎片(Fragmentation)

现有系统为每个请求预分配一段连续内存,大小为该请求的 maximum sequence length(如 2048 tokens)。这导致三种内存浪费:
缺陷二:无法利用内存共享机会

在并行采样(parallel sampling)中,一个 prompt 生成多个输出,prompt 部分的 KV Cache 本可共享。在 beam search 中,不同候选序列也可以共享部分前缀。但现有系统将各序列的 KV Cache 存放在独立的连续空间中,共享根本不可能
💡 论文的量化发现(Figure 2)

在实验中,现有系统的 KV Cache 内存中,实际存储 token states 的有效比例仅为 20.4% - 38.2%。也就是说,超过 60% 的 KV Cache 内存被浪费了!这是 vLLM 提升吞吐的根本空间。

二、LLM 推理基础(Background)

2.1 Transformer 的自回归分解

语言建模任务是对 token 序列 $(x_1, ..., x_n)$ 的联合概率建模。利用序列的自然顺序,可将其分解为条件概率的乘积:

$$P(x) = P(x_1) \cdot P(x_2 | x_1) \cdots P(x_n | x_1, \ldots, x_{n-1})$$

Transformer 通过自注意力层计算这一概率。对于输入隐藏状态序列 $(x_1, ..., x_n) \in \mathbb{R}^{n \times d}$,自注意力层首先对每个位置 $i$ 做线性变换得到 query、key、value 向量:

$$q_i = W_q x_i, \quad k_i = W_k x_i, \quad v_i = W_v x_i$$

然后计算注意力分数 $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})}, \quad o_i = \sum_{j=1}^{i} a_{ij} v_j$$

除注意力计算外,Embedding、Feed-forward、LayerNorm、Residual Connection、输出 logit 计算以及 Q/K/V 变换等,都是以 $y_i = f(x_i)$ 的形式独立作用于每个位置。

2.2 LLM 服务的两阶段生成

LLM 部署为条件生成服务(如 Completion API、Chatbot)。给定输入 prompt $(x_1, ..., x_n)$,LLM 按上述分解方式自回归地生成输出 token $(x_{n+1}, ..., x_{n+T})$。

由于每生成一个新 token 都依赖之前所有 token 的 key 和 value 向量,系统会将这些向量缓存起来,称为 KV Cache。注意:同一个 token 出现在序列不同位置时,其 KV Cache 是不同的(因为依赖前面所有 token)。

Phase 1: Prompt Phase(预填充阶段)

将完整用户 prompt $(x_1, ..., x_n)$ 作为输入,计算第一个新 token 的概率 $P(x_{n+1} | x_1, ..., x_n)$。同时生成所有 prompt token 的 key 向量 $k_1, ..., k_n$ 和 value 向量 $v_1, ..., v_n$。

由于所有 prompt token 已知,此阶段可通过 matrix-matrix multiplication 并行计算,能高效利用 GPU 并行性。

Phase 2: Autoregressive Generation Phase(自回归生成阶段)

在第 $t$ 次迭代,模型输入单个 token $x_{n+t}$,利用已缓存的 $k_1, ..., k_{n+t-1}$ 和 $v_1, ..., v_{n+t-1}$,计算 $P(x_{n+t+1} | x_1, ..., x_{n+t})$。只需计算新的 $k_{n+t}$ 和 $v_{n+t}$。

不同迭代之间因数据依赖无法并行,使用 matrix-vector multiplication,效率较低。此阶段严重低估用 GPU 计算力,成为 memory-bound,贡献了单请求延迟的大部分

2.3 Batching 技术演进

Batching 能提升 GPU 利用率,但 LLM serving 中的 batching 面临两个难题:

  1. 请求到达时间不同——要么让早到的等,要么让新来的等
  2. 输入/输出长度差异巨大——padding 会浪费计算和内存

Iteration-level scheduling(迭代级调度) 解决了这些问题:每次迭代后,完成的请求被移除,新请求被加入。新请求只需等待一个迭代周期,而非整个 batch 完成。配合专用 GPU kernel 消除 padding 需求。

关键认知:虽然 iteration-level scheduling 让计算更高效、batching 更灵活,但batch size 仍受 GPU 内存容量限制,特别是 KV Cache 的存储空间。因此,内存管理成为新的瓶颈。

三、现有系统的内存管理困境

3.1 KV Cache 的规模有多大?

以 OPT-13B 模型为例,单个 token 的 KV Cache 需求:

$$2 \text{ (key + value)} \times 5120 \text{ (hidden size)} \times 40 \text{ (layers)} \times 2 \text{ (FP16 bytes)} = 800 \text{ KB}$$

OPT 最大序列长度 2048 tokens,因此单个请求的 KV Cache 最高可达 1.6 GB。A100 40GB 即使全部内存给 KV Cache,也只能容纳几十个请求。

更严峻的是,GPU 计算速度增长快于内存容量增长(A100 → H100,FLOPS 提升 >2×,内存仍 80GB)。内存将越来越成为瓶颈

3.2 三种内存浪费的详细分析

Figure 3 展示了两个请求的内存布局:

浪费类型原因特点
Reserved为未来 token 预留最终会被使用,但长期占用阻止其他请求使用
Internal Fragmentation预分配 max length,实际更短请求结束后才知浪费,纯浪费
External Fragmentation不同请求预分配大小不同分配器(如 buddy)产生的间隙,纯浪费
Compaction(内存整理)不可行:虽然有研究提出 compaction 来解决碎片,但在性能敏感的 LLM serving 系统中,对庞大的 KV Cache 做 compaction 不现实。即使做了 compaction,预分配的 chunk 空间仍阻止了解码算法特有的内存共享。

3.3 复杂解码算法的额外挑战

LLM 服务提供多种解码算法,每种对内存管理有不同要求:

现有系统完全无法利用这些共享机会。

四、PagedAttention 核心算法

4.1 核心思想:操作系统分页的启示

PagedAttention 的灵感来自操作系统经典的 virtual memory with paging。类比关系如下:

操作系统LLM Serving (vLLM)
Page(页)KV Block(固定大小的 KV 块)
Byte(字节)Token(token 的 KV 向量)
Process(进程)Request(推理请求)
Virtual Address(虚拟地址)Logical KV Block(逻辑块)
Physical Address(物理地址)Physical KV Block(物理块)
Page Table(页表)Block Table(块表)

4.2 分块注意力计算

传统注意力要求 KV Cache 连续存储。PagedAttention 打破了这一限制——允许连续的 key 和 value 存储在非连续的内存空间中

具体做法:将每个序列的 KV Cache 划分为 KV blocks,每个 block 包含固定数量 token 的 key 和 value 向量。设 block size 为 $B$:

Key block: $K_j = (k_{(j-1)B+1}, \ldots, k_{jB})$

Value block: $V_j = (v_{(j-1)B+1}, \ldots, v_{jB})$

注意力计算可转化为 block-wise 形式:

$$A_{ij} = \frac{\exp(q_i^\top K_j / \sqrt{d})}{\sum_{t=1}^{\lceil i/B \rceil} \exp(q_i^\top K_t / \sqrt{d})}, \quad o_i = \sum_{j=1}^{\lceil i/B \rceil} V_j A_{ij}^\top$$

其中 $A_{ij} = (a_{i,(j-1)B+1}, \ldots, a_{i,jB})$ 是第 $j$ 个 KV block 上的注意力分数行向量。

GPU Kernel 的工作方式:PagedAttention kernel 在计算时,根据 block table 分别定位和获取不同的 KV block。例如查询 token "forth" 时,kernel 逐个读取 block 中的 key 向量计算注意力分数,再与对应 block 的 value 向量相乘得到输出。

4.3 解决内存问题的三大机制

1. 消除内部碎片(Internal Fragmentation)

使用相对小的 block(默认 16 tokens),按需分配。未用满的 block 最多浪费一个 block 的空间,而非整个 max-length chunk。

2. 消除外部碎片(External Fragmentation)

所有 block 大小相同,分配器无需处理不同大小的请求,从根本上消除了外部碎片。

3. 实现细粒度内存共享

共享可以在 block 粒度进行——同一个物理 block 可被映射到多个序列的逻辑 block 中。

五、vLLM 系统设计

5.1 系统架构概览

vLLM 采用中心化调度器(Centralized Scheduler)协调分布式 GPU Worker 的执行。KV Cache Manager 以分页方式管理 KV Cache。

组件职责
Centralized Scheduler协调分布式 GPU worker 执行,决定每个迭代处理哪些请求
KV Cache Manager管理物理 KV Cache 内存,维护 block table 映射
CPU Block Allocator管理被 swap 到 CPU RAM 的 block(用于抢占恢复)
GPU Block Allocator管理 GPU DRAM 上的物理 KV block
Block Engine在每个 GPU worker 上,分配连续 GPU DRAM 并划分为物理 KV blocks

5.2 Block Table:逻辑到物理的映射

vLLM 将 KV Cache 组织为固定大小的 KV blocks。每个请求的 KV Cache 表示为一系列逻辑 KV blocks,从左到右填充。最后一个 block 的未填充位置为后续生成预留。

Block engine 在 GPU worker 上分配连续的 GPU DRAM chunk,划分为物理 KV blocks。KV block manager 维护 block tables——每个请求的逻辑 block 到物理 block 的映射。

分离逻辑与物理 KV blocks 的关键收益:

5.3 Decoding 执行流程示例

Figure 6 展示了单个序列的完整执行过程:

Step 1: Prompt 阶段初始化

不预先预留最大序列长度的内存。只为 prompt 计算所需的 KV block 做映射。Prompt 有 7 个 token,映射前 2 个逻辑 block(0 和 1)到 2 个物理 block(7 和 1)。

Prefill 时,vLLM 用传统 self-attention 生成 prompt 的 KV Cache 和第一个输出 token。前 4 个 token 的 KV 存在逻辑 block 0(物理 block 7),后 3 个存在逻辑 block 1(物理 block 1)。

Step 2: 第一个 Decode 迭代

用 PagedAttention 在物理 block 7 和 1 上生成新 token。最后一个逻辑 block 还有空位,新 KV 直接存入,更新 block table 的 #filled 记录。

Step 3: 第二个 Decode 迭代

最后一个逻辑 block 已满,新 KV 存入新逻辑 block。vLLM 分配新物理 block(物理 block 3),在 block table 中记录映射。

全局流程:每个 decode 迭代,vLLM 首先选择候选序列组成 batch,为新需要的逻辑 block 分配物理 block。然后将当前迭代的所有输入 token(prompt 阶段的全部 token + generation 阶段的最新 token)拼接为一个序列输入 LLM。LLM 计算时,PagedAttention kernel 按逻辑 block 读取之前的 KV Cache,将新生成的 KV 存入物理 block。

5.4 Block Size 的权衡

一个 block 中存储多个 token(block size > 1)能让 PagedAttention kernel 在更多位置上并行处理 KV Cache,提升硬件利用率、降低延迟。但 block size 过大也会增加内部碎片。

设计选择:实验中 block size 16-128 在 ShareGPT 上表现最好;Alpaca 上 16-32 最好(序列较短)。默认 block size = 16,足够利用 GPU 并行,又避免显著内部碎片。

六、复杂解码场景的内存共享

6.1 Parallel Sampling:并行采样

场景:一个 prompt 生成多个输出样本,用户从中选择。

由于所有样本共享同一个 prompt,prompt 的 KV Cache 可以共享。vLLM 通过 PagedAttention 和分页内存管理轻松实现这一点:

Copy-on-Write 机制(块粒度)

当样本 A1 需要写入其最后一个逻辑 block 时,vLLM 发现对应物理 block 的引用计数 > 1,于是:

  1. 分配一个新物理 block
  2. 将原物理 block 的数据复制到新 block
  3. 原物理 block 引用计数减 1
  4. A1 在新物理 block 上写入

当样本 A2 后续写入原物理 block 时,引用计数已为 1,直接写入。

收益:prompt 的 KV Cache 空间在多个样本间共享,仅最后一个逻辑 block 需要 copy-on-write。对于长输入 prompt,内存节省非常可观。

6.2 Beam Search:束搜索

Beam search 用于机器翻译等任务,用户期望 top-k 最佳输出。算法维护 $k$ 个候选序列,每步扩展所有候选并保留概率最高的 $k$ 个。

与 parallel sampling 不同,beam search 不仅可共享 prompt block,还可共享更多前缀 block,且共享模式随解码过程动态变化。

Figure 9 展示了 $k=4$ 的例子:

💡 Beam Search 的内存收益

现有系统需要在 beam 候选之间频繁复制 KV Cache。vLLM 通过物理 block 共享,大部分 block 可在不同候选间共享,copy-on-write 仅在旧共享 block 内生成新 token 时触发,只涉及一个 block 的复制。

6.3 Shared Prefix:共享前缀

常见场景:用户提供任务描述(system prompt / instruction / few-shot examples),与实际输入拼接形成完整 prompt。多个用户请求共享相同前缀。

vLLM 的处理方式类似于 OS 处理进程间的共享库:

  1. 服务提供方预先为共享前缀保留一组物理 block
  2. 用户请求的 prompt 逻辑 block 直接映射到这些缓存的物理 block(最后一个 block 标记为 copy-on-write)
  3. Prompt 阶段只需对用户的任务输入部分执行计算

6.4 Mixed Decoding:混合解码方法

不同解码方法(greedy sampling, parallel sampling, beam search, shared prefix)具有不同的内存共享和访问模式。vLLM 通过统一的逻辑→物理 block 映射层隐藏了这些复杂性。

LLM 和执行 kernel 只看到每个序列的物理 block ID 列表,无需处理序列间的共享模式。这扩大了不同采样需求请求的 batching 机会,最终提升整体吞吐。

七、调度与抢占:Swapping vs Recomputation

7.1 FCFS 调度策略

vLLM 对所有请求采用 First-Come-First-Served (FCFS) 调度,确保公平性和防止饥饿。当需要抢占时,最早到达的请求优先服务,最晚到达的请求优先被抢占。

7.2 内存耗尽时的困境

LLM 服务面临的独特挑战:输入 prompt 长度差异巨大,输出长度事先未知。随着请求数量和输出增长,GPU 物理 block 可能耗尽。此时需要回答两个经典问题:

  1. 哪些 block 应该被逐出(evict)?
  2. 被逐出的 block 如何恢复?

7.3 All-or-Nothing 逐出策略

传统逐出策略用启发式预测哪个 block 最远被访问。但 vLLM 知道一个序列的所有 block 是一起被访问的,因此采用 all-or-nothing 策略:要么逐出一个序列的全部 block,要么不逐出。

此外,同一请求内的多个序列(如 beam search 的候选)作为 sequence group 一起调度,因它们之间可能存在内存共享。

7.4 两种恢复技术

技术一:Swapping(交换)

经典虚拟内存技术:将被逐出的 page 复制到磁盘 swap 空间。vLLM 中,被逐出的 block 复制到 CPU RAM

vLLM 包含 CPU block allocator 管理被 swap 到 CPU RAM 的物理 block。当 GPU 物理 block 耗尽时,选择一组序列逐出,将其 KV Cache 转移到 CPU。

一旦请求完成,其 block 被释放,被抢占序列的 block 被重新加载回 GPU 继续处理。

关键保证:CPU swap 空间上的 block 数不超过 GPU RAM 上分配的物理 block 总数,因此 swap 空间有界。

技术二:Recomputation(重计算)

当被抢占序列被重新调度时,简单地重新计算其 KV Cache

重计算延迟可以显著低于原始延迟:decode 阶段生成的 token 可与原始用户 prompt 拼接为新 prompt,所有位置的 KV Cache 可在一次 prompt phase 迭代中生成。

7.5 Swapping vs Recomputation 的权衡

维度SwappingRecomputation
小 block size大量小块数据传输,PCIe 带宽受限开销恒定,不受 block size 影响
大 block size传输效率高,更优仍可行,但不如 swapping
中 block size (16-64)两者相当两者相当
绝对开销较大(受 PCIe 带宽限制)不超过 swapping 的 20%
实验结论(§7.3):对于中等 block size(16-64),两种方法端到端性能相当。小 block size 下 recomputation 更优;大 block size 下 swapping 更优。但 recomputation 的开销从未超过 swapping 的 20%

八、分布式执行

8.1 Tensor Parallelism 支持

许多 LLM 参数规模超过单卡容量,需要在分布式 GPU 上分区执行。vLLM 支持广泛使用的 Megatron-LM style tensor model parallelism

8.2 分布式 KV Cache 管理

关键观察:即使使用模型并行执行,每个模型 shard 仍处理相同的输入 token 集合,因此需要相同位置的 KV Cache。

vLLM 在中心化调度器中维护单个 KV Cache Manager,所有 GPU worker 共享该管理器以及逻辑 block 到物理 block 的映射。每个 GPU worker 有相同的物理 block ID,但只存储对应 attention head 的 KV Cache 部分。

分布式执行流程
  1. 调度器准备每个请求在 batch 中的输入 token ID,以及每个请求的 block table
  2. 调度器将控制消息广播给所有 GPU worker
  3. GPU worker 用输入 token ID 执行模型
  4. Attention 层中,GPU worker 按控制消息中的 block table 读取 KV Cache
  5. 执行中,GPU worker 通过 all-reduce 同步中间结果(无需调度器协调)
  6. 最后,GPU worker 将本迭代的采样 token 发回调度器
重要设计:GPU worker 无需在内存管理上同步——它们只需在每个 decode 迭代开始时接收所有内存管理信息(随 step inputs 一起)。

九、实验评估

9.1 实验设置

模型GPU总 GPU 内存参数大小KV Cache 可用最大 KV slots
OPT-13B1×A10040 GB26 GB12 GB15.7K
OPT-66B4×A100160 GB132 GB21 GB9.7K
OPT-175B8×A100-80GB640 GB346 GB264 GB60.1K

工作负载:

Baseline:

关键指标:Normalized latency = 端到端延迟 / 输出长度。高吞吐系统应在高请求率下保持低 normalized latency。

9.2 Basic Sampling:基本采样

ShareGPT 结果(Figure 12 第一行):

Alpaca 结果(Figure 12 第二行):

核心结论:序列越长、模型越大、解码算法越复杂,vLLM 的优势越明显。因为在这些场景下,内存瓶颈越严重,PagedAttention 的内存效率提升价值越大。

9.3 Parallel Sampling & Beam Search

Parallel Sampling(Figure 14 第一行):

并行样本数越多,vLLM 相对 Orca 的优势越大。因为 prompt KV Cache 的共享比例随样本数增加而提升。

Beam Search(Figure 14 第二行):

Beam search 允许更多共享,vLLM 收益更大。在 OPT-13B + Alpaca 上:

内存节省量化(Figure 15):

方法AlpacaShareGPT
Parallel sampling (2-6 outputs)6.1% - 9.8%16.2% - 30.5%
Beam search (width 2-6)37.6% - 55.2%44.3% - 66.3%

9.4 Shared Prefix:共享前缀

使用 LLaMA-13B + WMT16 英德翻译数据集:

前缀越长、共享比例越高,vLLM 的 prefix caching 收益越大。

9.5 Chatbot:聊天机器人

使用 ShareGPT 数据集合成聊天历史。因 OPT-13B 上下文限制,将 prompt 截断到最后 1024 tokens。

十、消融实验与关键设计选择

10.1 Kernel Microbenchmark

PagedAttention 的动态 block 映射给 GPU 操作带来额外开销:

Figure 18a 显示,相比高度优化的 FasterTransformer,PagedAttention 的 attention kernel 延迟高 20-26%

但:这一开销仅影响 attention 算子,不影响 Linear 等其他算子。尽管有 kernel 开销,PagedAttention 仍使 vLLM 在端到端性能上远超 FasterTransformer——因为内存效率的提升远远超过了 kernel 的额外开销。

10.2 Block Size 的影响

Block size 是 vLLM 的关键超参数:

实验结论(Figure 18b):

10.3 Swapping vs Recomputation 深度对比

Figure 19a(Microbenchmark):

Figure 19b(End-to-end):

十一、核心洞察与论文价值

11.1 论文的核心方法论

将成熟 OS 技术重新应用到新领域:PagedAttention 并非发明了全新的算法,而是敏锐地识别出 LLM serving 中 KV Cache 的管理问题与操作系统虚拟内存问题的高度相似性,将 paging、copy-on-write 等经过数十年验证的技术适配到 GPU 内存管理中。

11.2 与 Orca 的互补关系

论文在 Related Work 中明确论述了与 Orca 的关系:

维度OrcavLLM
优化目标提升 GPU 利用率提升内存利用率
核心手段Iteration-level scheduling + 请求交错PagedAttention + 分页内存管理
互补性两者正交互补!Orca 的细粒度调度让内存管理更具挑战,vLLM 的技术因此更有价值

11.3 LLM-Specific 的优化增强

vLLM 并非简单照搬 OS 技术,而是结合 LLM serving 的特点做了关键增强:

  1. All-or-nothing swap-out:利用"处理请求需要其全部 token states"的特性,序列级而非 block 级逐出
  2. Recomputation 恢复:利用 LLM 可将 decode token 拼接为 prompt 重新计算的特性,OS 中不可行
  3. Kernel fusion:将内存访问操作与 attention 等计算融合为单个 kernel,降低 paging 的间接开销

11.4 技术适用性的边界

论文在 Discussion 中坦诚讨论了技术边界:

不适用场景: 适用条件:动态内存分配(输出长度未知)+ 性能受 GPU 内存容量限制。LLM serving 恰好同时满足。

11.5 为什么这篇论文重要

🏆 学术与工程的双重价值

11.6 一句话总结

PagedAttention 的核心洞见:LLM 推理中 KV Cache 的管理问题,本质上是操作系统虚拟内存问题的 GPU 版本。将 paging 和 copy-on-write 适配到注意力计算中,就能在几乎零内存浪费的前提下,实现灵活的内存共享,从而将 batch size 提升数倍,最终带来 2-4× 的吞吐提升