大语言模型推理内存管理算法 paged attention

本文结合chatgpt生成。

Paged Attention 是 vLLM 提出的一个用于 高效管理 Transformer 推理阶段 KV Cache 的机制。它本质上不是一种新的 attention 数学形式，而是一种 KV Cache 的内存管理与访问机制。

1. 背景：推理时最大的瓶颈之一是 KV Cache

在自回归生成中，模型每生成一个 token，都要对历史 token 做 attention。

对于每一层 Transformer，历史 token 的 Key 和 Value 会被缓存下来：

\[K_{1:t}, V_{1:t} \]

这样生成第 t+1 个 token 时，就不需要重新计算前面所有 token 的 K/V。

这就是 KV Cache。

问题是：KV Cache 很大。

假设：

• batch size 很大；
• 每个请求生成长度不同；
• 有些请求很短，有些请求很长；
• 请求不断进来和结束；

那么 GPU 显存里会有大量动态变化的 KV Cache。

传统做法通常会给每个序列预分配一段连续显存，例如最大长度 2048 或 4096。这样会导致很严重的浪费。

例如一个请求只生成了 128 个 token，但可能为它预留了 2048 个 token 的 KV Cache 空间。

2. Paged Attention 的核心想法

Paged Attention 把每个序列的 KV Cache 切成固定大小的 block，也可以叫 page。

例如每个 block 存 16 个 token 的 KV Cache：

sequence A:
tokens:  [0 ... 15] [16 ... 31] [32 ... 47] ...
blocks:   block 0    block 1     block 2

这些 block 在物理显存中不要求连续。

逻辑上，一个序列的 token 是连续的；

物理上，它的 KV Cache 可以分散在不同显存块里。

这和操作系统的虚拟内存很像：

逻辑地址空间：
sequence A block 0 -> physical block 7
sequence A block 1 -> physical block 2
sequence A block 2 -> physical block 19

模型计算 attention 时，通过一个 block table 找到对应的物理 KV block。

3. 为什么叫 Paged Attention？

因为它借鉴了 paging 的思想：

操作系统虚拟内存	Paged Attention
虚拟页 page	KV cache block
页表 page table	block table
物理页框	GPU 中的 KV block
进程虚拟地址连续	序列 token 逻辑上连续
物理内存不连续	KV Cache 物理上不连续

所以 “paged” 指的是 KV Cache 按页组织，而不是 attention 权重本身分页。

---

4. 它解决的主要问题：显存碎片和浪费

传统 KV Cache 分配的问题主要有两个。

第一，内部碎片

如果为每个请求按最大长度预留空间，会浪费大量未使用 token 的 KV Cache。

例如：

请求实际长度：200 tokens
预留长度：2048 tokens
浪费：1848 tokens 的 KV cache 空间

Paged Attention 按 block 增长，需要多少分配多少。最多浪费一个 block 内的剩余空间。

如果 block size 是 16，那么一个序列最多浪费 15 个 token 的空间。

---

第二，外部碎片

如果要求每个序列的 KV Cache 物理连续，那么随着请求不断进入和结束，显存会出现碎片(比如某一块显存刚结束任务，空闲了下来)。

可能总空闲显存足够，但没有一整段连续空间可以给新请求。

Paged Attention 不要求物理连续，只需要找到若干空闲 block，所以显存利用率更高。

5. Attention 计算时怎么访问 KV？

普通 attention 大致是：

\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V \]

在推理阶段，对于当前 token 的 query \(Q_t\)，需要访问之前所有 token 的 \(K, V\)。

Paged Attention 中，历史 KV 不再是一段连续数组，而是一组 block：

logical token positions:
0 1 2 ... 15 | 16 17 ... 31 | 32 ...

block table:
logical block 0 -> physical block 12
logical block 1 -> physical block 5
logical block 2 -> physical block 27

计算 attention kernel 时，会根据 token 的位置找到对应 logical block，再通过 block table 映射到 physical block，从而读出对应的 K/V。

所以它需要一个专门实现的 CUDA kernel，否则频繁间接寻址会带来额外开销。

vLLM 的贡献之一就是把这种 block-based KV 访问做成高效 kernel。

---

6. 它和普通 attention 的区别

从数学上看，Paged Attention 没有改变 attention 公式

它改变的是 KV Cache 的存储布局和调度方式。

所以它和 FlashAttention 不完全是同一类东西。

技术	主要解决什么
FlashAttention	训练/推理中 attention 计算的 IO 优化
Paged Attention	推理阶段 KV Cache 的显存管理
Multi-Query Attention / Grouped-Query Attention	减少 KV head 数量，从而减少 KV Cache 大小
Paged Attention	让已有 KV Cache 更高效地分配、复用、共享

它们可以组合使用。

---

7. Paged Attention 对 serving 很重要

LLM serving 面临的是动态请求：

request A: prompt 200 tokens, generate 50 tokens
request B: prompt 1000 tokens, generate 500 tokens
request C: prompt 50 tokens, generate 2000 tokens
...

每个请求长度不同、结束时间不同。如果显存管理不灵活，batch 很难做大。

Paged Attention 让系统可以：

1. 按需为请求分配 KV block；
2. 请求结束后释放 block；
3. 把释放的 block 立刻复用给新请求；
4. 支持更多并发请求；
5. 降低显存浪费；
6. 提高吞吐量。

这也是 vLLM 吞吐量很高的重要原因之一。

8. 一个更具体的例子

假设 block size = 4。

一个序列目前有 10 个 token：

tokens: 0 1 2 3 | 4 5 6 7 | 8 9

它需要 3 个 block：

logical block 0: tokens 0-3
logical block 1: tokens 4-7
logical block 2: tokens 8-11

物理显存中可能是：

logical block 0 -> physical block 8
logical block 1 -> physical block 3
logical block 2 -> physical block 15

第 11 个 token 生成时，还可以继续写入 logical block 2。

当 token 数达到 12，再生成第 13 个 token 时，就需要分配一个新的 block。

请求结束后，这 3 或 4 个 physical block 都放回 free list，供其他请求复用。

9. Copy-on-Write：支持高效共享 KV Cache

Paged Attention 还有一个很重要的好处：可以支持 KV Cache 的共享。

比如 beam search，或者多个请求共享同一个 prompt：

prompt: "请总结下面这篇文章..."

多个生成分支的前缀相同，它们的 KV Cache 可以共享同一批 block。

当某个分支继续生成新 token，产生差异时，再复制需要修改的 block，这就是 copy-on-write。

类似操作系统里 fork 进程后的内存页共享。

这对 beam search（留多个候选序列，找到整体概率更高的序列的一种解码策略）、Parallel Sampling（对同一个 prompt 同时采样出多个不同回答）、prefix caching（缓存公共前缀 prompt 的 KV Cache） 都很有价值。