前一篇写完之后,KV Cache 在我脑子里的形象已经变了——它不再是一次 forward 里的临时中间量,而是一个跨多轮存活、随 Decode 持续增长、每个请求各自独立的状态。
但"知道它是有状态的"和"知道怎么管它"是两件事。当我开始想具体的内存管理问题的时候,发现事情比想象的复杂。
同一块 GPU 上的三个请求#
假设有一块 GPU 同时在服务三个请求:
- Request A:Prompt 1000 Token,正在 Decode,已经生成了 200 多个 Token,还没结束。
- Request B:Prompt 500 Token,生成了 20 个 Token 后遇到 EOS,已经结束。
- Request C:Prompt 2000 Token,刚完成 Prefill,正准备开始 Decode。
三个请求的 KV Cache 大小不一样,增长速度不一样,生命周期也不一样。A 还在持续增长,B 已经死了,C 刚出生。
对比一下传统推理(比如我之前做的推荐模型 Serving):一个请求进来,分配临时 workspace 做一次 forward,结束后释放。所有请求的 activation 生命周期都是"用完即弃",大小在 forward 前就可以确定。
KV Cache 完全不是这种模式:
- 跨多轮 forward 保留。 Decode 的每一步都要读取历史所有位置的 K 和 V——前一篇推导过这一点。
- 持续增长。 每生成一个 Token,KV Cache 就多一行。谁都不知道最终会生成多少个 Token。
- 请求间完全独立。 每个请求有自己的 prompt、自己的生成进度、自己的 KV Cache。
- 结束后需要回收。 B 结束了,它占用的显存应该能让 C 用上。
- 新请求随时可能加入。 B 结束后空出的空间可能正好够塞一个新请求进来。
在不知道最终序列长度的情况下,怎么为大量并发请求高效地管理 KV Cache?
这是我在前一篇末尾提到 PagedAttention 时留下的问题。这篇就把它展开。
连续分配为什么撑不住#
在看 PagedAttention 之前,我先自己想了几种朴素的分配策略。
方案 A:按最大 Context Length 预分配。
比如模型支持最大 4096 Token 的上下文。那每来一个请求,直接分配一整块能容纳 4096 Token KV Cache 的连续显存。
好处很明显:地址计算简单,不用在 Decode 过程中扩容。Token position \(i\) 在第 \(l\) 层的 KV 地址直接用 base + offset 就能算出来,不需要间接查表。
问题也很明显:一个实际只生成了 100 Token 的请求,白白预留了 4096 Token 的空间。前一篇算过,单个请求在典型配置下的 KV Cache 可以到几个 GiB 的量级。按最大长度预分配,大量显存被没有实际数据的预留空间占据,能同时服务的请求数就少了。
方案 B:按当前长度分配,增长时重新申请。
Prefill 结束后按 prompt 长度分配一段连续空间,后续 Decode 时空间不够就申请更大的连续块。
这比方案 A 节省了预留浪费。但连续内存的扩容有一个潜在问题:当原有空间后面的地址已经被其他请求占据、无法原地扩展时,可能需要申请一块新的更大连续区域,把已有数据搬过去,更新引用。搬迁有性能开销,而且在显存碎片化的情况下,可能连一大块连续空间都找不到。
方案 C:更复杂的分配器可以缓解,但不改变根本假设。
真实系统不会停在方案 B 的朴素版本上。PyTorch 的 Caching Allocator 会预留额外空间来减少频繁分配和搬迁;CUDA VMM 可以在虚拟地址连续的前提下按需映射物理页。这些机制都能部分缓解上面的问题。
但我从另一个方向想了想:方案 A 和 B 的根本假设都是 KV Cache 的逻辑序列需要存储在一段物理连续的显存中。
有没有办法让 KV Cache 在逻辑上看起来是连续的——Token 0, 1, 2, … 按顺序排列——但底层的物理存储不要求彼此相邻?
分页:逻辑连续,物理散落#
想到这一步的时候,我脑子里冒出来的第一个类比是操作系统的虚拟内存。
进程看到的地址空间是连续的,但物理内存是按固定大小的页(比如 4 KiB)分配的。逻辑上相邻的虚拟页,在物理内存中完全可以不相邻。操作系统用页表维护虚拟页号到物理页帧的映射。
PagedAttention 做的事情在思路上和这个很像。它把 KV Cache 的管理单位从"一整段连续缓冲区"变成了"固定大小的 Block"。
三个核心概念:
- Logical Block:每个请求的 KV 序列按照固定大小(比如每 Block 容纳 4 个 Token 的 KV)划分成逻辑块。Logical Block 0 存 Token 0-3 的 KV,Logical Block 1 存 Token 4-7 的 KV,依此类推。
- Physical Block:GPU 显存中真正保存 KV 数据的物理单元。系统初始化时会把一部分显存预先划分成一个 Physical Block Pool,每个 Block 大小固定。
- Block Table:为每个请求建立 Logical Block Index 到 Physical Block ID 的映射。
有了这层映射,逻辑上相邻的 Block 在物理显存中就完全不需要相邻了。
用一个具体例子:
1Block Size = 4 Token
2
3Request A(共 16 Token 的 KV Cache):
4 Logical Block 0 → Physical Block 7
5 Logical Block 1 → Physical Block 1
6 Logical Block 2 → Physical Block 9
7 Logical Block 3 → Physical Block 5Physical Block 7, 1, 9, 5 在显存中的位置完全不连续。但对于 Request A 来说,它的 KV 序列在逻辑上仍然是 Token 0, 1, 2, …, 15 有序排列的。
要定位某个 Token 的 KV 数据,需要通过 Block Table 做一次间接寻址。比如 Token 10:
1Logical Block Index = 10 // 4 = 2
2Offset in Block = 10 % 4 = 2
3Physical Block ID = BlockTable[2] = 9Token 10 的 KV 存储在 Physical Block 9 的第 2 个 slot(从 0 开始计数)。
这个过程有几个容易混淆的地方:
Block Table 是地址映射元数据,不是 Attention 的 Key Tensor。 Block Table 的每个条目存的是一个 Physical Block ID(本质上是一个索引或指针),不是 KV 数据本身。它的角色类似操作系统里的页表条目。
这是软件层的管理,不是 GPU 硬件页表。 GPU 有自己的 MMU 和页表,那是管理 GPU 虚拟地址到物理地址映射的硬件机制。PagedAttention 的 Block Table 是 Serving 软件在应用层维护的数据结构。两者处于不同的系统层次。
固定大小的 Block 减少了外部碎片和过度预分配,但最后一个 Block 仍然可能有内部碎片。 比如一个请求只有 14 个 Token,Block Size 是 4,那么 Logical Block 3 只用了 2 个 slot,另外 2 个是空的。这和操作系统分页中最后一页的内部碎片是同一回事。
Attention Kernel 也得跟着改#
到这里我停下来想了一个问题:KV 存储布局变了,原本的 Attention Kernel 还能直接用吗?
传统的 Attention 实现假设 KV Cache 在物理上是连续的。给定 base address 和 token index,可以直接算出地址:
1address = base + token_index * stride一次简单的线性寻址。
但在分页存储下,逻辑 token index 不再直接对应物理地址。Kernel 需要先查 Block Table,找到对应的 Physical Block,再在 Block 内部按 offset 访问:
1block_idx = token_index // block_size
2offset = token_index % block_size
3physical_block = block_table[block_idx]
4address = physical_block_base + offset * stride这里的伪代码是概念模型,具体 stride 和物理布局因实现而异。但核心变化是确定的:Attention Kernel 内部需要增加一层间接寻址逻辑——把 Block Table 传进去,在遍历历史 Token 的 K/V 时按 Block 为单位查表、跳转。
PagedAttention 没有修改 Attention 的数学语义。 \(\text{Softmax}(QK^T/\sqrt{d_k})V\) 还是那个公式,每个 Token 还是要和所有历史 Token 做 Attention。改变的是 KV 数据在物理存储中的组织方式和访问路径。
但这个变化不是没有代价。连续存储下的线性遍历对 GPU 的内存访问模式更友好——顺序读取容易被 HBM burst 和 L2 Cache 优化。分页存储下,不同 Block 可能分布在显存的不同位置,访问模式变得不那么规整。这给 Kernel 实现带来额外的工程复杂度,也可能带来一些寻址开销。
后来我意识到,这其实是一个在系统设计中反复出现的 pattern:存储布局的改变不只影响分配器本身,还会传导到所有读取这块存储的计算逻辑。 内存管理从来不只是分配器内部的问题。
不过这篇不打算深入 CUDA Kernel 的具体实现——Thread Block 划分、Warp 调度、Tiling 策略、Online Softmax——那些更偏算子工程,是另一条线了。
谁分配、谁回收、谁计算#
PagedAttention 不只是一个 Kernel,它是一个内存管理架构。我在读 vLLM 的设计时,发现系统中有三个组件各自承担不同的职责:
Scheduler:决定哪些请求在当前 step 执行、每个请求处理多少 Token。它是全局的调度决策者,但不直接管理 Block 的分配细节。
KV Cache Manager:负责 Physical Block 的分配、引用计数、缓存和回收。它维护一个 Block Pool,知道哪些 Block 已被占用、哪些空闲、哪些可以被淘汰。当 Scheduler 决定让一个请求继续执行时,KV Cache Manager 判断是否有可用 Block 来容纳新增的 KV。
PagedAttention Kernel:根据每个请求的 Block Table,访问分散在各个 Physical Block 中的 KV 数据,完成 Attention 计算。它不关心 Block 是怎么分配的,只管按 Block Table 寻址。
用一个具体案例来说明这三者的协作。假设 GPU 上的 KV Block Pool 有 8 个 Physical Block,Block Size 为 4 Token:
时刻 1:Request A 和 B 加入。
A 的 Prompt 有 6 个 Token(Prefill 后需要 2 个 Block),B 的 Prompt 有 5 个 Token(需要 2 个 Block)。
1A: Logical Block 0 → Physical Block 0 (Token 0-3)
2 Logical Block 1 → Physical Block 1 (Token 4-5, 2 slots 空闲)
3
4B: Logical Block 0 → Physical Block 2 (Token 0-3)
5 Logical Block 1 → Physical Block 3 (Token 4, 3 slots 空闲)
6
7已占用: 0, 1, 2, 3 空闲: 4, 5, 6, 7时刻 2:A 继续 Decode,生成了 3 个 Token。
A 的 KV Cache 从 6 Token 增长到 9 Token。Logical Block 1 先被填满(原来有 2 个空 slot,填入 Token 6 和 7),然后需要一个新的 Logical Block 2。KV Cache Manager 从空闲池分配 Physical Block 4。
1A: Logical Block 0 → Physical Block 0 (Token 0-3)
2 Logical Block 1 → Physical Block 1 (Token 4-7, 满)
3 Logical Block 2 → Physical Block 4 (Token 8, 3 slots 空闲)
4
5B: 不变
6
7已占用: 0, 1, 2, 3, 4 空闲: 5, 6, 7时刻 3:B 结束,Request C 加入。
B 结束后,它的 Physical Block 2 和 3 被释放回空闲池。C 的 Prompt 有 7 个 Token(需要 2 个 Block)。
1B 释放 → Block 2, 3 回归空闲池
2
3C: Logical Block 0 → Physical Block 2 (Token 0-3)
4 Logical Block 1 → Physical Block 3 (Token 4-6, 1 slot 空闲)
5
6已占用: 0, 1, 2, 3, 4 空闲: 5, 6, 7C 拿到的 Physical Block 2 和 3 就是 B 刚释放的。
时刻 4:A 继续增长。
A 又生成了 4 个 Token(总共 13 Token)。Logical Block 2 被填满,需要新的 Logical Block 3。
1A: Logical Block 0 → Physical Block 0 (Token 0-3)
2 Logical Block 1 → Physical Block 1 (Token 4-7)
3 Logical Block 2 → Physical Block 4 (Token 8-11, 满)
4 Logical Block 3 → Physical Block 5 (Token 12, 3 slots 空闲)
5
6C: 不变
7
8已占用: 0, 1, 2, 3, 4, 5 空闲: 6, 7整个过程中几个值得留意的特性:
- 按需分配,不预留。 A 不需要一开始就占据能容纳最大长度的 Block 数量,从 2 个 Block 起步,随 Decode 逐步申请新 Block。
- Block 复用。 B 释放的 Block 立刻可以被 C 使用,不需要等到所有请求结束才整理碎片。
- 没有数据搬迁。 A 增长时,已经写入 Physical Block 0 和 1 的数据纹丝不动。新数据写入新分配的 Block,Block Table 加一条映射就行。
- Block Table 是唯一需要更新的元数据。 每次新分配一个 Block,只需在 Block Table 末尾追加一个条目。
这里没有展开 Continuous Batching 的调度逻辑——那涉及在 batch 中动态替换已完成的请求、调整每 step 的 batch 组成,是另一个话题。下一篇会讲。
共享前缀:分页带来的额外能力#
Block-based 的管理方式还带来了一个我一开始没想到的好处。
很多 Serving 场景下,不同请求会共享相同的 System Prompt。比如一个 chatbot 应用,所有用户的请求前面都有一段相同的系统指令。如果每个请求都独立存储这段前缀的 KV Cache,就有大量重复数据。
分页架构提供了一个自然的解决方式:让不同请求的 Logical Block 指向同一个 Physical Block。
1Request X: Logical Block 0 → Physical Block 10
2 Logical Block 1 → Physical Block 11
3 Logical Block 2 → Physical Block 6 ← X 独有
4
5Request Y: Logical Block 0 → Physical Block 10 ← 与 X 共享
6 Logical Block 1 → Physical Block 11 ← 与 X 共享
7 Logical Block 2 → Physical Block 8 ← Y 独有Physical Block 10 和 11 存放公共前缀的 KV 数据,X 和 Y 的 Block Table 都指向它们,但物理上只存一份。
这让我想到以前打 ACM 写后缀数组时用的 Longest Common Prefix(LCP)——确定两个序列的公共前缀长度,然后在分叉点之后各走各的。Prefix Caching 在概念上有类似的味道:找到 Token 序列级别的公共前缀,以完整 Block 为单位判定可复用的范围。
不过实际做 Prefix Caching 时,有几个细节不能忽略:
Block 的 hash 必须包含前缀上下文信息。 不能只看当前 Block 里的 Token ID 是否一样——同样的 Token 出现在不同的前缀上下文中,由于前面层的信息传递,它们的 KV 可能不同。正确的做法是把前缀链路上的信息也纳入 hash 计算,确保只有真正等价的 KV 才会被复用。模型配置、位置信息等也需要考虑。
Reference Count 和缓存保留是两件事。 一个 Physical Block 如果还有活跃的请求在引用它(Reference Count > 0),不能被释放——释放了正在用它的请求就出错了。但当 Reference Count 降为零时,这个 Block 未必要立刻回收。它可以作为 Prefix Cache 保留在 Pool 中,等待后续请求命中。只有当 Pool 空间紧张、需要腾地方时,才淘汰那些引用数为零的缓存 Block。
这一节不深入 Hash Chain 实现和 LRU 淘汰策略。要记住的核心点是:Block-based 的内存管理不仅支持按需分配,还为跨请求的物理 KV 共享提供了自然的表示方式。 这是连续分配方案很难优雅做到的——你很难让两段物理连续的缓冲区"共享"中间的一段。
另一条路:让 CUDA VMM 来管?#
写到这里我自然会想到一个问题:既然 PagedAttention 的核心思想是"逻辑连续、物理不连续",而 CUDA 本身就有 Virtual Memory Management API 支持虚拟地址到物理页的按需映射——为什么不直接用 CUDA VMM,非要在应用层自己管 Block Table?
vAttention(Prabhu et al., 2024)就是沿着这条路走的。它用 CUDA 的低级 VMM API(cuMemCreate、cuMemMap)为 KV Cache 预留一段连续的虚拟地址空间,然后在 Decode 过程中按需将物理页映射上去。对 Attention Kernel 来说,KV Cache 的虚拟地址始终是连续的——寻址方式和传统方案一样,不需要查 Block Table。
两种方案的核心取舍:
| PagedAttention | vAttention | |
|---|---|---|
| 逻辑到物理映射 | Serving 软件维护 Block Table | CUDA VMM 硬件/驱动层映射 |
| KV 虚拟地址布局 | 非连续 Block 寻址 | 连续虚拟地址 |
| Attention Kernel | 需要支持分页寻址的专用 Kernel | 更容易复用连续寻址的标准 Kernel |
| 关键取舍 | 灵活的软件级 Block 管理,Kernel 更复杂 | 利用 VMM 保持地址连续性,但受映射粒度和运行时约束影响 |
两者都实现了"物理内存按需分配"这个目标。区别在于把映射的复杂度放在哪一层。PagedAttention 选择在应用层管理,获得了完全的软件控制(灵活的 Block 大小、Prefix Caching、自定义淘汰策略),代价是 Kernel 需要处理间接寻址。vAttention 把映射交给 CUDA VMM,保持了虚拟地址连续性,让现有的连续寻址 Kernel 更容易复用,但需要处理 VMM API 的粒度限制和映射开销。
这不是一个谁一定更好的问题。它体现了一个在系统设计中反复出现的思想:同一个问题,可以选择在应用层解决,也可以利用更底层的抽象来解决。不同的选择会把复杂度放在不同的系统层次上。
几点体会#
回看 PagedAttention 解决的问题,有三个认知我觉得值得记住:
KV Cache 是动态增长的有状态资源。 这一点在前一篇已经推导过了,但到了内存管理这一层感受更具体。LLM Serving 管理的不是临时 Tensor,而是一堆生命周期各不相同、持续增长、需要回收复用的状态。这比传统推理的内存管理复杂一个层次。
逻辑连续不要求物理连续。 通过 Block Table 做一层间接映射,KV Cache 的逻辑 Token 序列和底层物理存储就解耦了。请求的 KV 可以持续增长,只需要分配新的 Block 并在 Block Table 末尾追加一条映射,已有的数据不需要移动。这个思想在操作系统的虚拟内存中早就存在了,PagedAttention 把它搬到了 GPU KV Cache 管理的场景。
内存管理决定 Serving 的有效容量。 PagedAttention 的主要价值不在于加速单次 Attention 计算——分页寻址的间接性甚至可能带来一些开销。它的价值在于减少显存浪费:不预留、不搬迁、按需分配、及时回收。在相同的 GPU 显存条件下,更高效的 KV Cache 管理意味着能容纳更多活跃请求,从而改善系统级的吞吐能力。
下一篇准备写 Continuous Batching 和 Chunked Prefill——有了分页管理 KV Cache 的能力之后,Serving 系统怎么在多个并发请求之间调度 GPU 的计算资源。
参考资料#
- Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023 — https://arxiv.org/abs/2309.06180
- Prabhu et al., vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention, 2024 — https://arxiv.org/abs/2405.04437
- 从一次 Forward 到持续生成:LLM Inference 为什么不一样?
- 从 MHA 到 GQA:为什么现代 LLM 不再为每个 Query Head 分配独立的 KV?