PagedAttention 笔记：为什么要像管理内存一样管理 KV 缓存

读 PagedAttention 时，我觉得最重要的不是“像操作系统分页”这个类比，而是它把 KV 缓存变成了需要明确管理的内存对象。

大模型推理中的 KV 缓存有几个棘手的特点：

• 请求长度不同。
• 解码长度无法提前知道。
• 请求随时结束，显存要回收。
• 并行采样和束搜索会共享提示词。
• 前缀缓存命中时，多个请求可能复用同一段历史。
• 调度器每一步都在重新组织批次。

如果仍然采用“每个请求占一大段连续 KV 缓冲区”的方式，显存很快会浪费在预留和碎片上。

先看一个请求怎么长大

一个请求进入推理引擎后，它需要的 KV 缓存不会一次性完整确定。

预填充提示词
  -> 为提示词 K/V 分配块
  -> 执行预填充注意力计算
  -> 保存 K/V 块

解码第 1 个 token
  -> 读取提示词对应的块
  -> 追加新的 K/V

解码第 2 个 token
  -> 读取提示词和第 1 个 token 对应的块
  -> 追加新的 K/V

这个请求像一条不断增长的序列。它需要保持逻辑 token 顺序，但物理上不一定非要连续。

PagedAttention 的核心就是把这两个问题拆开：

token 的逻辑顺序：
  由请求和调度器维护

KV 的物理位置：
  由块管理器维护

块表像一层小型 MMU

在 PagedAttention 中，请求的 KV 缓存被切成逻辑块，每个逻辑块再通过块表指向一个物理块。

request A:
  logical block 0 -> physical block 8
  logical block 1 -> physical block 2
  logical block 2 -> physical block 15

注意力内核仍然要按 token 顺序读取历史 K/V。区别在于，它不能再假设 K/V 位于一段连续地址，而要通过块表找到物理位置。

这层间接性换来两个能力：

1. 请求不需要连续显存。
2. 物理块可以被释放、复用和共享。

代价也很清楚：计算内核要理解块表，推理引擎要维护块的生命周期，调试时也多了一层映射。

块大小是系统参数

块太大，尾部浪费就会很明显。请求的最后一个块可能只用了几个 token，却占住整个物理块。

块太小，块表会更大，管理成本更高，计算内核的间接寻址也更频繁。

所以块大小不是论文里的装饰项，它与这些因素紧密相关：

• 平均提示词长度。
• 平均解码长度。
• 批次大小。
• 单头维度。
• 计算内核访问块布局的效率。
• 共享前缀比例。

这类参数不能脱离实际负载单独评价好坏。

共享前缀不是复制字符串

在并行采样或束搜索中，多个分支共享同一段提示词。直觉上，这段提示词的 K/V 不应该复制多份。

PagedAttention 可以让多个序列的逻辑块指向同一组物理块：

shared prompt blocks:
  P8 -> P2 -> P15

branch A:
  P8 -> P2 -> P15 -> P4

branch B:
  P8 -> P2 -> P15 -> P9

这里必须有引用计数。只要还有一个序列引用 P8，它就不能回到空闲链表。

这也是 KV 缓存与普通“字符串前缀相同”不同的地方。系统共享的是每层、每个注意力头、每个 token 对应的 K/V 块，而不是一段文本。

写时复制是正确性的边界

共享块只能读取。分支继续写入时，如果改动了一个被多个序列引用的块，就会污染其它分支。

所以需要写时复制：

P8 ref_count = 2
branch A wants to append into P8

allocate P11
copy P8 -> P11
branch A maps logical block to P11
P8 ref_count--
将新 token 写入 P11

这个过程如果做错，模型可能不会直接报错，而会悄悄读到其它分支写入的 K/V。输出变得异常，调用栈却未必留下明显错误。

调试 KV 缓存系统时，这类错误比内存泄漏更棘手，因为它污染的是上下文语义。

调度器不能只看剩余显存

PagedAttention 让块管理更灵活，但调度器仍然要做选择。

它要决定：

• 新请求能不能进来。
• 当前批次中哪些请求继续解码。
• 块不够时是抢占、等待、拒绝还是换出。
• 前缀命中是否值得等待。
• 长请求是否挤压短请求的尾延迟。

也就是说，块管理器只回答“物理块怎么放”，调度器还要回答“这一步应该服务谁”。

这点和 HiCache/Mooncake 的关系很像：有远端 KV 不代表一定要等它，能够预取也不代表当前请求应该阻塞。

PagedAttention 和 HiCache 不是同一层

我会把它们这样分开：

PagedAttention：
  位于单个推理服务运行时内部
  管理 GPU KV 物理块

HiCache：
  位于 SGLang 运行时内部
  管理 GPU、主机内存和存储层的三级缓存状态

Mooncake Store：
  位于运行时外部
  管理远端对象、副本和传输链路

PagedAttention 管理 GPU 中 KV 块的组织方式。HiCache 继续把前缀 KV 状态扩展到主机内存和外部存储。Mooncake Store 再提供远端对象存取和传输。

它们都和 KV Cache 有关，但状态对象不一样：

• PagedAttention 的对象是物理 KV 块。
• HiCache 的对象是前缀树节点的 GPU、主机内存和存储状态。
• Mooncake 的对象是对象、切片、副本和 Segment。

混在一起讲“KV 缓存命中”，很容易把读取路径说错。

我会重点查的几处实现

阅读 vLLM、SGLang 或类似推理引擎的实现时，我会直接找这些点：

• 块表怎样记录在序列上。
• 物理块的空闲链表怎样维护。
• 预填充后的块怎样提交。
• 解码追加数据时怎样申请新块。
• 共享前缀怎样增加引用计数。
• 请求结束或取消时怎样释放块。
• 写时复制发生在写入之前还是之后。
• 注意力内核怎样通过块表读取 K/V。
• 块不足时调度器采取什么策略。

PagedAttention 的价值不在于有一个好记的分页比喻，而在于把 KV 缓存从连续数组改成了有映射、有引用、有生命周期的内存对象。