| 项 | 信息 |
|---|---|
| 论文标题 | Efficient Memory Management for Large Language Model Serving with PagedAttention |
| 会议 | SOSP'23（系统顶会） |
| 一作 | Woosuk Kwon（UC Berkeley） |
| 通讯 | Ion Stoica（Spark/Ray 创始人） |
| arxiv | 2309.06180 |
| 代码 | github.com/vllm-project/vllm |
现有系统将每个请求的 KV Cache 存储在连续内存空间中，按最大可能序列长度预分配。这带来三种浪费：
| 浪费类型 | 原文术语 | 含义 |
|---|---|---|
| ① 预留槽位 | Reserved slots | 为未来可能生成的 token 预留空间。在请求整个生命周期内被独占，不能被其他请求利用。 |
| ② 内部碎片 | Internal fragmentation | 请求实际生成的 token 数少于最大长度，剩余预分配空间永远不会被使用（只有请求结束后才知道浪费了多少）。 |
| ③ 外部碎片 | External fragmentation | 不同请求的预分配块大小不同，释放后产生大小不一的空洞。新请求即使总空闲内存够也可能找不到足够大的连续空间。 |
将每个序列的 KV Cache 划分为固定大小的 KV block（默认每个 block 存 16 个 token 的 K 和 V），这些 block 在物理显存中不必连续。通过 block table（等价于 OS 页表）记录逻辑 block → 物理 block 的映射。
| OS 概念 | vLLM 对应 |
|---|---|
| 物理页 (physical page) | 物理 KV block |
| 虚拟地址空间 | 序列的逻辑 KV 位置 |
| 页表 (page table) | block table |
| 按需分页 (demand paging) | 按需分配 KV block |
| 字节 (byte) | token |
| 进程 (process) | 请求 (request) |
标准 attention kernel 假设 K/V 在内存中连续排列。PagedAttention kernel 多了一步间接寻址：
数学结果与标准 attention 完全等价。
$$A_{ij} = \frac{\exp(q_i^\top K_j / \sqrt{d})}{\sum_{l=1}^{\lceil T/B \rceil} \exp(q_i^\top K_l / \sqrt{d})} , \quad o_i = \sum_{j=1}^{\lceil T/B \rceil} V_j A_{ij}^\top$$
其中 $K_j$ 是第 j 个 block 的所有 Key 向量，$B$ 是 block size。
假设序列 10 个 token，block_size = 4，切成 3 个 block：
Block 0: token 0,1,2,3 的 K/V
Block 1: token 4,5,6,7 的 K/V
Block 2: token 8,9 的 K/V（没填满）
标准 attention 做的事：
scores = [q·k_0, q·k_1, q·k_2, ..., q·k_9] / √d ← 10 个 score
weights = softmax(scores) ← 对 10 个分数归一化
out = Σ weights[i] * v_i ← 加权求和
PagedAttention 做的事（等价改写）：
# Step 1：逐 block 算 raw score（不做 softmax！）
scores_block0 = [q·k_0, q·k_1, q·k_2, q·k_3] / √d
scores_block1 = [q·k_4, q·k_5, q·k_6, q·k_7] / √d
scores_block2 = [q·k_8, q·k_9] / √d
# Step 2：拼成一条，做一次全局 softmax
all_scores = concat(scores_block0, scores_block1, scores_block2)
weights = softmax(all_scores) ← 分母是全部 10 个 exp 之和
# Step 3：用全局 weights 对所有 V 加权
out = weights[0:4] · V_block0 + weights[4:8] · V_block1 + weights[8:10] · V_block2
| 步骤 | 发生什么 | Block 状态 |
|---|---|---|
| ① 新请求到达（Prefill） | prompt 的 KV 一次性计算，按 block_size 填入物理 block。最后一个 block 可能没填满。 | 从 pool 分配 N 个 block |
| ② Decode 生成中 | 新 token 的 KV 填入最后一个 block 的空位 | 最后一个 block 的 #filled + 1 |
| ③ 最后一个 block 满了 | 从 pool 申请新的物理 block，block table 追加一行 | 新 block 加入 |
| ④ 请求结束（EOS） | 所有物理 block 归还到 pool | 全部释放 |
唯一的内部碎片：最后一个 block 可能没填满（最多浪费 block_size - 1 = 15 个 token 的空间）。整体来看极小——相比之前 80% 的浪费，这点代价几乎可忽略。
同一 prompt 生成多个回答时：
多个 beam 共享不同长度的公共前缀，共享模式随解码动态变化（类似 OS 的进程树）。
CoW 触发条件：某条 beam 要向一个 ref_count > 1 的 block 写入新数据时。
触发后发生什么：
只 copy 被修改的那一个 block，不是整段 KV Cache——这就是 CoW 的核心价值。跟 Unix fork() 完全一致。
多用户共享相同 system prompt 时，前缀的 KV block 全局只存一份。新请求直接将逻辑 block 映射到已有物理 block。类似 OS 中进程间共享动态库。
| 策略 | 做法 | 适合场景 |
|---|---|---|
| Swapping（交换） | 将被抢占请求的 KV block 从 GPU 复制到 CPU 内存。GPU 空闲后再换回。 | CPU 内存充裕、KV cache 体量中等 |
| Recomputation（重计算） | 直接丢弃被抢占请求的 KV cache。等有资源后重新跑 Prefill（将原始 prompt + 已生成 token 合并重算 KV）。 | prompt 短（重算代价低）、CPU 内存紧张 |
| Continuous Batching | PagedAttention | |
|---|---|---|
| 解决什么 | 调度问题（请求什么时候进出 batch） | 内存管理问题（KV Cache 怎么存） |
| 核心机制 | 每 decode step 重新组 batch，请求结束立刻换上新的 | KV Cache 分 block + block table 映射 |
| 没有对方能不能工作 | 能（Orca 就是有 CB 没有 PA） | 能（可以不做 CB 只用 PA 管内存） |
| 但效果差 | 没 PA → 碎片严重 → batch 放不满 | 没 CB → 请求进出不灵活 → GPU 空转 |
| FlashAttention | PagedAttention | |
|---|---|---|
| 解决什么 | 单次 attention 计算的 HBM 读写量（IO 效率） | 多请求 KV Cache 的显存布局（内存管理） |
| 手段 | Tiling（分块计算 + 不保存中间结果到 HBM） | 分页（非连续存储 + block table 间接寻址） |
| 优化层面 | 一次 matmul 内部的数据搬运 | 多请求间的内存分配/释放/共享 |
| 能一起用吗？ | 能，且 vLLM 实际就是两者结合——PagedAttention kernel 内部采用 FlashAttention 风格的 tiling 减少 HBM 读写 | |
HBM（High Bandwidth Memory）就是 GPU 的主显存（"A100 有 80GB 显存"说的就是 HBM）。GPU 的内存层次类似 CPU：
| 层次 | GPU 上叫什么 | 大小 | 带宽 |
|---|---|---|---|
| 寄存器 | Register File | ~几 KB/SM | 最快 |
| L1 / 共享内存 | SRAM (Shared Memory) | 每 SM 192KB（A100） | ~19 TB/s |
| 主显存 | HBM | 40-80 GB | 2 TB/s |
关键差距：SRAM 比 HBM 快约 10 倍，但容量小 1000 倍。
做一次 softmax(Q @ K^T / √d) @ V 计算，标准实现的数据搬运过程：
1. 从 HBM 读出 Q (T×d)
2. 从 HBM 读出 K (T×d)
3. 计算 Q @ K^T → score (T×T)
4. 把 score 写回 HBM ← SRAM 放不下 T×T 矩阵！
5. 从 HBM 重新读出 score
6. 做 softmax → weights (T×T)
7. 把 weights 写回 HBM ← 又放不下
8. 从 HBM 重新读出 weights
9. 从 HBM 读出 V (T×d)
10. 计算 weights @ V → output
11. 把 output 写回 HBM
总 HBM 读写量：读 3Td + 2T²，写 2T² + Td → 总量 $\mathcal{O}(T^2 + Td)$
当 T=4096, d=128 时：$T^2 = 16M$ 元素，大部分时间 GPU 在等数据从 HBM 搬过来，算力空转。
核心 idea：不把中间结果（score、weights 这两个 T×T 矩阵）写回 HBM。
做法：Tiling（分块计算）——把 Q/K/V 切成小块，每次只在 SRAM 里处理一小块，在 SRAM 内部完成 softmax 的部分计算，中间结果不落地到 HBM。
| 标准 Attention | FlashAttention | |
|---|---|---|
| HBM 读写量 | $\mathcal{O}(T^2 + Td)$ | $\mathcal{O}(Td)$ |
| T=4096, d=128 | ~34M 元素 | ~2M 元素 |
| 加速比 | 基准 | 2-4×（主要是省了 HBM 搬运时间） |
| 对比维度 | 数据 |
|---|---|
| vLLM vs FasterTransformer | 吞吐高 14-24× |
| vLLM vs Orca (最优配置) | 吞吐高 2.2-4.3× |
| KV Cache 有效利用率 | vLLM 96.3% vs Orca(Max) 20.4% |
| Beam Search (width=6) | vLLM 比 Orca 吞吐高 2.3× |
| PagedAttention kernel 额外开销 | < 5% |
| 最优 block_size | 16（消融实验 16-32 最优区间） |
| OS 概念 | vLLM 对应 | 说明 |
|---|---|---|
| 物理页 | 物理 KV block | 固定大小（16 token），存储在 block pool 中 |
| 虚拟页 | 逻辑 KV block | 按序列位置编号：0, 1, 2, ... |
| 页表 | block table | 每个请求一张，逻辑→物理映射 |
| 按需分页 | 按需分配 block | 最后一个 block 满了才申请新 block |
| Copy-on-Write | Beam Search / Parallel Sampling 共享 | ref_count > 1 时写入触发 copy |
| Swap to disk | Swap to CPU memory | 被抢占请求的 KV block 暂存 CPU |
| 共享动态库 | Shared Prefix / Prefix Caching | 多请求共享同一 system prompt 的 KV block |
| 进程 | 请求 (request) | 独立的逻辑地址空间 |
| 字节 | token | 寻址的最小单位 |
vllm/v1/core/）vllm/v1/core/sched/）这是 Stage 1 / Week 2 的第一篇，对应 AI Infra 转型学习计划"精读 vLLM 论文"任务。