第4章 PagedAttention：虚拟内存的启示

"Good artists copy, great artists steal." -- Steve Jobs（引用 Picasso）

本章要点

• 理解 KV Cache 内存碎片问题的本质：为什么传统方案浪费 60-80% 的显存
• 从操作系统的虚拟内存机制类比理解 PagedAttention 的设计
• 深入 PagedAttention 内核的计算过程：如何在非连续内存上高效执行注意力
• 掌握块表（Block Table）的数据结构与映射机制
• 理解 Copy-on-Write 和块共享在并行采样中的应用

4.1 KV Cache：甜蜜的负担

要理解 PagedAttention，首先要理解它解决的问题。

Transformer 模型的注意力机制需要计算当前 Token 与所有历史 Token 的关系。为了避免重复计算，每个 Token 在第一次被处理时，会将其 Key 和 Value 向量缓存起来——这就是 KV Cache。

KV Cache 的大小与序列长度成正比。对于一个 Llama-2-13B 模型：

• 每个 Token 的 KV Cache 大小 = 2 × num_layers × num_kv_heads × head_dim × dtype_size
• = 2 × 40 × 40 × 128 × 2 bytes（FP16） = 800 KB

一个 2048 Token 的序列需要 1.6 GB 的 KV Cache。在一张 80 GB 的 A100 上，模型权重占约 26 GB，剩余 54 GB 中，KV Cache 最多能容纳 33 个并发序列。

但实际能容纳的远少于此。问题出在内存分配策略上。

传统分配：整块预留

传统的 KV Cache 管理方式是为每个请求预先分配最大可能长度的连续显存。如果最大长度设为 2048，即使一个请求最终只生成了 100 个 Token，它也占着 2048 Token 的显存。

这就像餐厅预订——你为 10 人订了一张大桌，但只来了 3 个人。剩下 7 个座位空着，其他客人也不能坐。

浪费分三类：

论文的测量结果触目惊心：在真实工作负载下，传统方案的显存浪费率高达 60-80%。这意味着，一张 80 GB 的 GPU，有 40-50 GB 的显存在"假装被使用"。

4.2 操作系统的答案

1961 年，曼彻斯特大学的 Atlas 计算机面对的是类似的问题——不同程序需要的内存大小不同，直接分配物理内存会导致碎片化。

答案是虚拟内存（Virtual Memory）：

1. 物理内存被划分为固定大小的页帧（Page Frame）
2. 每个程序看到的是连续的虚拟地址空间
3. 通过**页表（Page Table）**将虚拟页映射到物理页帧
4. 物理页帧不需要连续——离散分布在内存中也没关系

PagedAttention 将同样的思想应用到 KV Cache 管理上：

操作系统概念	PagedAttention 对应
物理页帧	KV Cache 块（Block）
虚拟页	Token 序列的逻辑分段
页表	块表（Block Table）
按需分页	KV Cache 块按需分配
页面回收	块的 LRU 驱逐
Copy-on-Write	共享块的写时复制

4.3 PagedAttention 的设计

块（Block）：固定大小的 KV Cache 单元

KV Cache 被划分为固定大小的块，每个块可以存储 B 个 Token 的 Key-Value 对（默认 B=16）。

一个块的物理大小 = 2 × num_kv_heads × head_dim × B × dtype_size（2 是因为 Key 和 Value 各一份）。

对于 Llama-2-13B（FP16），每个块 = 2 × 40 × 128 × 16 × 2 = 320 KB。

所有块在 GPU 显存中被预先分配为一个大的连续数组——BlockPool。这些块是物理上连续的（为了 GPU 内核的高效访问），但逻辑上可以被任意请求以任意顺序使用。

块表（Block Table）：逻辑到物理的映射

每个请求维护一张块表，记录它的第 i 个逻辑块映射到哪个物理块。

关键洞察：请求 A 的三个块在物理上是不连续的（块 7、块 3、块 12），但通过块表，请求 A"看到"的是连续的 Token 序列（0-47）。这就是虚拟内存的魔力——逻辑连续，物理离散。

按需分配：只用多少，占多少

与传统的"预留最大长度"不同，PagedAttention 按需分配块。一个新请求到达时，只为它的输入 Token 分配刚好够用的块数：

• 输入 100 Token → 分配 ⌈100/16⌉ = 7 个块
• 解码阶段，每生成 16 个 Token → 多分配 1 个块
• 请求结束 → 释放所有块

这就像按需租房——不需要提前预定整栋楼，每次只租一间，住满了再租一间。

内部碎片从整个序列缩小到了最后一个块。一个 100 Token 的请求，最后一个块只装了 4 个 Token（100 mod 16），浪费 12 个位置。对比传统方案浪费 2048-100=1948 个位置，改进了两个数量级。

外部碎片则完全消除了——因为块大小固定，任何空闲块都能被任何请求使用，不存在"空间够但不连续"的问题。

论文的测量结果

PagedAttention 将 KV Cache 的浪费率从 60-80% 降到了 4% 以下。这意味着同样的显存可以服务 2-4 倍的并发请求。在吞吐量测试中，vLLM 比 HuggingFace Transformers 高 14-24 倍，比当时的 SOTA 系统（如 Orca）高 2-4 倍。

4.4 注意力内核的挑战

有了分页的 KV Cache，注意力计算面临一个新挑战：标准的 FlashAttention 假设 KV Cache 在内存中是连续的。但现在，一个请求的 KV Cache 分散在不连续的物理块中。

PagedAttention 内核需要：

1. 查询当前请求的块表，获取物理块地址
2. 遍历所有物理块，对每个块中的 Key 计算注意力分数
3. 对所有块的分数做 softmax 归一化
4. 用归一化的权重加权 Value，得到注意力输出

这个过程在 CUDA 内核中通过**分块归约（Block Reduction）**实现：

对于 Query token q:
    对于每个物理块 b（通过块表间接访问）:
        加载 b 中的 K 向量: K[b*16 : (b+1)*16]
        计算注意力分数: score = Q × K^T / √d
        记录局部 max 和 sum（用于在线 softmax）
        加载 b 中的 V 向量
        累积加权 V

    全局归一化（合并所有块的局部统计量）
    输出最终注意力结果

在线 Softmax（Online Softmax） 是这里的关键技巧。标准 softmax 需要先遍历所有分数找到 max，再遍历一次做归一化。但在分块处理中，我们不能等到遍历完所有块才做归一化——那样就需要把所有中间结果存下来，丧失了分块的意义。

在线 softmax 允许我们在遍历每个块时就增量计算 softmax，最后只需要一次修正就能得到全局正确的结果。这个算法最早由 Milakov 和 Gimelshein 在 2018 年提出，后来被 FlashAttention 广泛采用。PagedAttention 在此基础上增加了块表间接寻址。

4.5 Copy-on-Write：块共享的优雅

PagedAttention 的分页机制还带来了一个额外的收益：块共享。

考虑并行采样场景——用户发送同一个 Prompt，要求生成 3 个不同的回复（n=3）。这 3 个采样共享同一个 Prompt 的 KV Cache，只是解码阶段生成的 Token 不同。

在传统方案中，每个采样需要独立复制一份完整的 Prompt KV Cache。但在 PagedAttention 中，3 个采样可以共享 Prompt 对应的物理块——它们的块表指向同一组物理块，引用计数为 3。

当某个采样需要修改共享块时（比如最后一个 Prompt 块还没填满，需要追加解码 Token），系统使用 Copy-on-Write（写时复制）：

1. 检查目标块的引用计数
2. 如果 > 1，复制该块到一个新的物理块，将引用计数减 1
3. 修改新块（追加 Token），只影响当前采样

这与 Linux 的 fork() 系统调用使用的 COW 机制一模一样——父子进程共享页面，只在写入时才真正复制。在 LLM 推理中，这可以节省大量显存：如果 Prompt 有 500 个 Token，3 个并行采样本需要 3 × 500 的 KV Cache，有了块共享只需要 500 + 3 × 解码部分。

4.6 从论文到工程

论文描述了核心思想，但工程实现要复杂得多。这里列举几个论文中没有详细讨论，但在源码中能看到的重要工程决策。

块大小的选择

块大小（默认 16）不是随意选的。它受到三个约束：

• 太小（如 1）→ 块表太长，间接寻址开销大，CUDA 内核的局部性差
• 太大（如 256）→ 内部碎片增大（最后一个块平均浪费 128 个位置），KV Cache 分配粒度粗
• GPU 硬件对齐→ 块大小需要是 warp size（32）的因子或倍数，以便 CUDA 内核高效处理

16 是一个经过基准测试验证的平衡点：内部碎片控制在平均 8 个 Token（可接受），块表长度适中，CUDA 内核能高效处理。

预分配与惰性释放

系统启动时，vLLM 会根据可用显存大小预分配所有物理块。这个大数组在整个运行期间不会被释放或重新分配——避免了 CUDA 内存分配器的开销和碎片。

块的"分配"和"释放"只是逻辑操作：从空闲链表中取出一个块号，或把一个块号放回空闲链表。没有任何 cudaMalloc 或 cudaFree 调用。

多头注意力的块布局

实际的 KV Cache 块不是简单的一维数组。对于多头注意力（Multi-Head Attention），一个块的布局是：

Key 块: [num_kv_heads, block_size, head_dim]
Value 块: [num_kv_heads, block_size, head_dim]

Key 和 Value 分开存储（而不是交错），因为注意力计算时先访问所有 Key 计算分数，再访问所有 Value 做加权。分开存储可以让 GPU 在读取 Key 时获得连续的内存访问模式，提高缓存命中率。

4.7 本章小结

PagedAttention 是 vLLM 最核心的创新，它的智慧不在于发明新算法，而在于将一个 62 年前的操作系统思想移植到 GPU 显存管理中：

• 问题——传统 KV Cache 分配导致 60-80% 的显存浪费（预留浪费、内部碎片、外部碎片）
• 灵感——操作系统的虚拟内存分页机制
• 设计——固定大小的块、块表映射、按需分配、引用计数
• 内核实现——在非连续物理块上通过块表间接寻址执行注意力计算，配合在线 softmax
• 额外收益——Copy-on-Write 块共享，在并行采样中节省大量显存
• 工程细节——块大小选择、预分配策略、内存布局优化

这个设计将 vLLM 的吞吐量提升到了同期系统的 2-24 倍。更重要的是，它为后续的所有优化（前缀缓存、分块预填充、投机解码）奠定了基础——因为这些优化都依赖于 KV Cache 能被灵活地分配、共享和复用。

下一章，我们将深入 KV Cache 管理器——那个管理所有物理块的"房东"。它处理着块的生老病死：分配、释放、缓存、驱逐、共享。

---

延伸阅读

• Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023 (arXiv:2309.06180)
• Milakov & Gimelshein, "Online normalizer calculation for softmax", arXiv:1805.02867
• Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", NeurIPS 2022

源码导航

• V1 注意力后端：vllm/v1/attention/backends/
• PagedAttention 算子：vllm/attention/ops/
• FlashAttention 集成：vllm/vllm_flash_attn/