Infra: PagedAttention

KV Cache 为什么会浪费显存

生成新 token 时，attention 要读取此前 token 的 key 和 value。在 causal attention 中，过去 token 的 KV 不会因为后面新增了 token 而改变，因此可以保存在显存里反复使用，这就是 KV cache。Prefill 处理 prompt 并建立 cache；随后 decode 逐 token 生成，cache 随之增长。

Serving 要同时处理许多请求，而每个请求会生成多长，事先并不知道。一种朴素做法是按最大长度预留一整段连续空间，避免 cache 增长时放不下。但这些还没存入 KV 的空间也被占住了，无法交给其他请求。

假设一个请求已缓存 3 个 token，最终只需缓存 5 个，却一开始就分配了 8 个 slot；一个 slot 表示存放一个 token 的 KV 所需的空间。按 CSE 291 讲义中的分类，浪费有三种：(Jain et al., 2026)

  • Reservation：之后会用到、现在还空着的 2 个 slot。即使事先知道最终长度，过早占住这部分空间，也会挤掉其他请求。
  • Internal fragmentation：分配区内最终用不上的 3 个 slot，来自按最大可能长度过量分配。
  • External fragmentation：分配区之间的空闲空间太零散，无法满足新请求的连续分配要求；总空闲量足够，也可能放不下。
KV cache memory: prompt, generated tokens, and three forms of waste从左到右是同一段 GPU KV 内存：请求 A 的 2 个 prompt token、1 个已生成 token、2 个之后才用到的 slot，以及 3 个直到结束也用不到的 slot。A 共分配 8 格，目前用了 3 格，最终用到 5 格。箭头指向当前生成位置。A 与请求 B 之间的 2 格空隙属于 external fragmentation；B 占据右侧区域。Largemodels2 prompttokenslearnGenerated1 tokenfromdata2 future slots(reservation)<resv><resv><resv>3 slots never used(internal fragmentation)……ExternalfragmentationLLMis…Request BRequest Acurrent step
左右滑动查看完整示意图
按 CSE 291 讲义 Figure 2 的横向布局重绘，数值沿用正文的 3／5／8 示例。每格代表一个 token 的 BF16 KV 容量，只画 KV 区域，不含权重或 workspace。词仅作示意；虚线 reservation 区的词是按最终输出回填的，其 KV 此时尚不存在。

前两类是按最终长度回看的划分：现在还不知道，剩下的 5 个 slot 中，哪些以后会用到。要减少这些浪费，就需要让请求按需拿空间，并且不要求拿到的空间都连在一起。

按 Block 分配

PagedAttention 是 vLLM 的核心设计之一。它把 KV cache 拆成固定大小的 block，每个 block 容纳 \(B\) 个 token 的 KV。vLLM 预先建立 GPU block pool，请求需要增长时，再从 pool 中领取空闲 block。这样就不用为每个请求预留最大长度。(Kwon et al., 2023)

为了让 attention 找到拆开的 KV，每个请求维护一张 block table，连接两种编号：

  • Logical block：按 token 顺序分组。例如 \(B=4\) 时，位置 0–3 属于第 0 块，4–7 属于第 1 块。
  • Physical block：KV 实际存放的显存块。它们可以散落在 pool 中，位置不必与 token 顺序一致。
Logical KV → block table → physical pool每块 4 个 token。请求 A 的 9 个 token 被切成 A0、A1、A2，经 block table 分别映射到 P3、P0、P4。P1 属于请求 B，P2 空闲。A2 只有一个 token，其余三个 slot 留空。Request A · 9 cached tokensA001230 → P3A145671 → P0A282 → P4P04567A1P1BBBBRequest BP2Free blockP30123A0P48A2GPU KV block pool · dashed cells = unused slots
左右滑动查看完整示意图
每格代表一个 token 的 KV，数字是 token 位置；图中 B = 4，仅示意布局，未按字节比例绘制。A 的 logical block 连续，physical block 可以分散；P2 可分给任意请求，A2 的空位留给 A 继续增长。

图中 A 的 block table 是 [3, 0, 4]：前三个 logical block 分别放在 P3、P0、P4。比如读取位置 8 的 token，就先找到 logical block 2，再查表到 P4，读取块内位置 0。Attention kernel 直接按这张表读取 KV，无须先把它们复制成连续的一整段。

在不共享、按需分配的简单情形下，长度 \(T\) 需要 \(\lceil T/B\rceil\) 个 block，占用 \(B\lceil T/B\rceil\) 个 slot。图中 9 个 token 占 3 块、共 12 个 slot，只有尾块的 3 个 slot 空着。新增 token 先填尾块，填满后再领一块；请求结束时释放引用，无人使用的块就能交给其他请求。

这同时缓解了三类浪费：按需领取减少 reservation，尾块的空位始终少于 \(B\)，等大小且可不连续的 block 则避免了 pool 内的 external fragmentation。Block 越小，尾部浪费越少，但 block table 更大，kernel 的寻址和管理开销也更高。

共享前缀，分叉时 Copy-on-write

同一 prompt 采样多条输出时，prompt 的 KV 完全相同。让各序列的 block table 指向同一份 physical block，就能省掉重复存储。

问题出在尚未填满的共享尾块：A 和 B 接下来生成的 token 可能不同，如果都写进同一个空位，就会覆盖彼此的 KV。原论文用 reference count 记录有多少序列正在使用一块；写入时若引用数大于 1，就先复制该块，将写入者的映射改到新块，再追加自己的 KV。这就是 copy-on-write。已经填满的前缀块无需修改，可以继续共享。

Shared prefix · copy only the partial tail分叉前 A、B 都指向 P3 和 P0，分别保存位置 0 到 3 和位置 4 到 5 的 KV。A 写入不同的新 token 前，将 P0 复制到 P4。之后 A 指向 P3、P4，B 指向 P3、P0；两条输出仍共享完整的 P3。After A and B append different tokensP3 · shared prefix0123refcount = 2A: [P3, P4]45aP4Copied tail · refcount = 1B: [P3, P0]45bP0Original tail · refcount = 1
左右滑动查看完整示意图
原论文的 copy-on-write 示例：分叉前两条序列共享 6 个 token，block table 都是 [P3, P0]。A 复制尾块 P0 到 P4 后写入 a；B 独占原来的 P0 后写入 b。a、b 表示新 token 的 KV；完整前缀 P3 始终只存一份。

后来的 Automatic Prefix Caching 把复用扩展到不同时间到达的请求。例如，多个请求以相同的 system prompt 开头，就可以复用仍在缓存中的完整前缀块，省去这部分的 prefill 计算。(vLLM contributors, n.d.)

这里必须匹配整个前缀：同一句话接在不同的上文之后，其 KV 通常也不同。因此 vLLM 用于查找缓存的 hash 包含前序 block 的 hash、当前块的 token，以及 LoRA ID 等必要的上下文标识。当前这套设计只缓存完整 block；命中前缀省下的是重复计算，这些 KV 仍要参与后续的 attention。

Continuous Batching

显存容得下更多请求，还需要调度把它们用起来。在普通 decode 中，一轮计算为 batch 中每个活跃序列生成一个 token。各序列需要的轮数不同：如果固定整个 batch 的成员，短请求结束后，新请求仍要等最长的请求完成才能加入。

Continuous batching 每轮重新组织 batch：完成的请求退出，资源允许时接纳等待中的请求。这种 iteration-level scheduling 在 Orca 中已有，Stanford CS336 的 inference lecture 也将它与 PagedAttention 放在一起讨论。(Liang, 2025; Yu et al., 2022)

Continuous batching · refill at iteration boundaries两个并发名额，A 需要三轮 decode，B 需要一轮。C 在第一轮结束后就绪，需要两轮。静态 batch 在 B 结束后留下空位，C 等待 A 结束；continuous batching 在第二轮让 C 加入，与 A 一起执行两轮。Static batchStep 1ABStep 2AidleStep 3AidleC waits until this batch finishesContinuous batchStep 1ABStep 2ACStep 3ACB finishes → C joins at step 2
左右滑动查看完整示意图
只比较 decode 调度，假设最多并发两条序列，且 C 在第 1 轮结束时已经完成 prefill、显存足够。方格是一轮 decode，并非等长的时间单位；实际 prefill 成本与调度策略未画出。

图中 B 只需一轮，A 需要三轮。Continuous batching 让 C 在第二轮就加入，与 A 一起执行。为单独展示 decode 调度，图中假设 C 已完成 prefill；实际接纳新请求时，还要为 prefill 安排计算预算。

两者配合：分页让同一份显存容纳更多活跃请求，continuous batching 及时补入工作，提高 serving throughput。但单个请求的每一步 decode 未必更快，并发量仍受 KV 容量、计算预算和延迟目标约束。

PagedAttention 与 FlashAttention 也可以配合：前者管理跨生成步骤保留的 KV，让 attention 读取分页布局；后者将 attention 计算分块并融合，减少中间结果在显存中的读写。两者都保留原来的 dense attention 语义，分页本身不会缩短需要关注的上下文。

References

Jain, M., Kandar, T., Srivastava, S., Zhu, J., Chaudhary, J., Srivastava, D., & Kong, J. (2026). CSE 291A/DSC 291, Lecture 15: KV Cache Management and Flash Attention. haoailab.com
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Proceedings of the 29th Symposium on Operating Systems Principles, 611–626. doi.org
Liang, P. (2025). CS336: Language Modeling from Scratch, Lecture 10: Inference. github.com
vLLM contributors. (n.d.). Automatic Prefix Caching. docs.vllm.ai
Yu, G.-I., Jeong, J. S., Kim, G.-W., Kim, S., & Chun, B.-G. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 521–538. usenix.org

Cite this post

@misc{pu2026mlmlrevisitinfrapagedattention,
  author = {Pu, Fanyi},
  title  = {Infra: PagedAttention},
  year   = {2026},
  month  = {9},
  url    = {https://pufanyi.com/blog/ml/ml-revisit/infra/paged-attention}
}