PagedAttention：操作系统级 KV 缓存虚拟内存管理

引言

在大型语言模型（LLM）的推理过程中，KV 缓存（Key-Value Cache） 是用来存储历史 token 的注意力键值对，避免每次生成新 token 时重复计算。然而，随着序列长度增长，KV 缓存的内存占用呈线性膨胀，且传统框架采用连续预分配方式，导致严重的内存碎片和显存浪费。PagedAttention 借鉴操作系统的分页内存管理思想，将 KV 缓存划分为固定大小的块，按需映射物理内存，实现了接近零碎片的显存利用与灵活的内存共享。本教程将带你从问题出发，深入理解 PagedAttention 的设计原理与实现细节。

一、LLM 推理的内存挑战

1.1 KV 缓存的作用

在自回归生成中，模型每生成一个新 token，就会计算当前输入序列的 Key 和 Value 张量，并与过去的 KV 拼接，用于注意力计算。把先前所有 token 的 K、V 保存下来，就形成了 KV 缓存：

Attention(Q, K, V) = softmax(QK^T / √d) V

如果不缓存，每次都需要对越来越长的完整序列重新计算 K、V，计算量随生成步数平方增长。缓存后每一步只需计算新 token 的 K、V，然后追加到缓存中。

1.2 连续内存分配带来的问题

传统推理引擎（如 Hugging Face Transformers、FasterTransformer 早期版本）为每个序列预先分配一块能容纳 最大可能长度 的连续 KV 缓存空间。这带来两大缺陷：

• 碎片化严重：不同请求可能先后到达，预分配导致已释放的空间无法被新请求利用，内部和外部碎片都很高，实际批处理大小受限。
• 无法共享：即使多个请求具有相同的前缀（如系统提示），仍然各自存储独立的 KV 缓存，造成显存冗余。

示例：假设最大长度为 2048，每个请求分配 2GB 连续显存。当一批请求结束或长度变化时，这些大块内存难以动态复用，导致显存利用率通常低于 30%。

二、PagedAttention 的核心思想

2.1 操作系统虚拟内存的灵感

操作系统将物理内存分成固定大小的页（Page），进程使用虚拟地址，通过页表映射到不连续的物理页。这带来了：

• 消除外部碎片；
• 按需分配物理内存；
• 多进程共享同一物理页（如动态库）。

PagedAttention 将同样的思路应用于 KV 缓存：把序列的 KV 缓存划分为固定大小的 KV 块（Block），每个块存储固定数量 token 的 K、V。序列的 KV 缓存由多个逻辑块组成，这些块在显存中可以不连续。

2.2 三个关键抽象

2.3 统一的块表（Block Table）

每个请求维护一张块表，记录其逻辑块序号到物理块地址的映射。计算注意力时，通过块表找到物理块，再按 token 偏移量索引。这与操作系统的页表解析过程非常相似。

三、PagedAttention 工作流程详解

3.1 块尺寸选择

块大小（block_size）是一个重要的可调参数。比如 block_size = 16，表示每个块存储 16 个 token 的 K 和 V。选择依据：

• 太大：内部碎片增加（最后一个块可能未填满）；
• 太小：块数量增多，块表变大，并行调度复杂度上升。
• 通常取值为 8 ~ 64，需结合 GPU 线程束大小和内存事务特性调优。

3.2 请求的生命周期

3.2.1 预填充阶段（Prefill）

用户提示（prompt）一次性送入模型。PagedAttention 会为其分配足够数量的物理块来容纳所有初始 token。如果提示长度不是 block_size 的整数倍，最后一个块仅部分填充，其余位置保留 padding。

此阶段需要计算整个提示的 KV 并写入分配的块中。由于操作是并行完成的，效率很高。

3.2.2 生成阶段（Decoding）

每生成一个新 token：

1. 计算该 token 的 K、V；
2. 检验当前序列最后一块是否已满：
   • 如果已满，向块管理器申请一个新物理块；
   • 更新块表，追加映射；
3. 将新 K、V 写入块的对应位置；
4. 执行注意力计算，需遍历块表中的所有逻辑块。

注意力计算需要访问所有块的 K、V，此时利用 GPU 并行性，一次性加载块表，从非连续物理地址聚集数据。

3.3 内存管理与碎片消除

块管理器维护空闲物理块列表，分配策略类似于操作系统物理内存分配：

• 当请求完成或退出，其占用块被回收，重新加入空闲池；
• 任何空闲块都可以分配给任何新请求，不受原先连续约束；
• 内部碎片仅存在于每个序列的最后一个块，整体显存利用率可达 99% 以上。

3.4 前缀共享（Prefix Sharing）

多个请求可能共享相同的前缀（例如，同一个系统提示：“You are a helpful assistant.”）。传统方式每个请求各自存储这部分 KV，造成 N 倍冗余。

PagedAttention 通过引用计数实现物理块共享：

• 系统维护一个前缀块池，存储常见前缀的 KV 缓存；
• 当请求命中前缀，其初始块表中的逻辑块直接映射到共享的物理块；
• 这些物理块被标记为只读，当请求后续需要写入（例如生成产生了分歧），采用 Copy-on-Write 策略：在写入前复制该块，并更新该请求的块表映射到新块，原共享块引用计数减一。

这极大地减少了冗余，尤其对多轮对话和批量类似请求的场景显存节省可达数十倍。

四、GPU 内核实现要点

4.1 分块注意力计算

由于 K、V 物理上不连续，注意力内核必须能处理块表访问模式。典型实现采用专门编写的 CUDA Kernel：

• 每个线程块负责一个注意力头的部分计算；
• 从块表中读取该序列所有块地址，循环加载每个块的 K、V；
• 使用寄存器/共享内存缓存当前 token 的 Q，逐步完成 softmax 收缩。

vLLM 等框架的 PagedAttention 内核利用 FlashAttention 风格的分块 softmax 算法，进一步优化 IO。

4.2 块访问的合并优化

块地址不连续，但每个块内部是连续的显存区域。内核可以在块内使用向量化加载，且块大小对齐 GPU 缓存行（128 字节），最大化访存效率。

4.3 动态块调度

批处理时，多个请求的序列长度不同，所需的块数量也不同。PagedAttention 允许动态分配，配合 continuous batching，新的请求可以立即复用刚释放的块，实现更高吞吐。

五、与操作系统的类比总结

• 虚拟地址 = 逻辑 token 位置：通过（块索引, 块内偏移）定位序列中的任意 token。
• 页表 = 块表：记录逻辑块到物理块的映射，并支持按需调入（分配新块）。
• 页面置换：在显存使用接近容量时，可采用交换（offloading）策略将不活跃的块移出到 CPU 内存，类似 OS 的页面换出。
• Copy-on-Write：用于前缀共享，保证多个序列独立修改时不相互影响。

六、与竞品方案对比

使用 PagedAttention 的推理引擎（如 vLLM）在实际用例中将 batch size 可提升 3~5 倍，吞吐量提升 2~4 倍，同时显存占用大幅降低。

七、快速上手指南（基于 vLLM）

如果你希望立即体验 PagedAttention，通过 vLLM 只需几行代码：

from vllm import LLM, SamplingParams

# 模型会自动使用 PagedAttention 管理 KV 缓存
llm = LLM(model="meta-llama/Llama-2-7b-hf", block_size=16)

prompts = ["Explain quantum computing in simple terms."]
sampling_params = SamplingParams(temperature=0.8, max_tokens=200)
outputs = llm.generate(prompts, sampling_params)
print(outputs[0].outputs[0].text)

vLLM 内部为你完成了所有分页管理、前缀共享和高效内核调度，无需额外配置。

八、总结与展望

PagedAttention 将操作系统的成熟设计范式引入 LLM 推理，实现了 KV 缓存细粒度的动态管理，解决了困扰推理系统的内存碎片和共享难题。这不仅提高了单 GPU 的吞吐，也让长上下文、大批量服务成为可能。

未来方向包括：

• 分层块管理：利用 GPU 显存、CPU 内存和 NVMe 形成三级存储，自动 swap；
• 自适应块大小：根据序列特性动态调整块尺寸，进一步减少碎片；
• 分布式跨 GPU 共享：在多 GPU 推理中利用块表实现全局 KV 缓存池。

参考资料

• vLLM 官方文档
• Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP '23)