现代 AI 系统中的 KV Cache 优化：从内存墙到 PagedAttention

在 LLM 推理过程中，最核心的性能瓶颈之一并非计算量（Compute-bound），而是内存带宽（Memory-bound）。当我们讨论生成式 AI 的推理速度时，实际上是在讨论如何高效地管理 KV Cache（Key-Value Cache）。

什么是 KV Cache？

在 Transformer 的解码阶段，模型每生成一个新 Token，都需要与之前所有已生成的 Token 进行 Attention 计算。如果每次都重新计算所有历史 Token 的 Key 和 Value 向量，计算复杂度将随序列长度呈平方级增长。

为了避免重复计算，我们将之前步骤产生的 K 和 V 向量存储在内存中，这就是 KV Cache。这意味着在生成第 $n$ 个 Token 时，我们只需要计算当前 Token 的 K 和 V，并将其追加到缓存中。

内存墙与碎片化问题

尽管 KV Cache 解决了计算冗余，但它引入了巨大的内存压力。一个典型的 Llama-3-70B 模型，其 KV Cache 在 FP16 精度下，每个 Token 消耗的内存量惊人。随着并发请求（Batch Size）和上下文长度（Context Length）的增加，显存会被迅速填满。

传统的 KV Cache 管理方式是预先分配一块连续的内存空间。这种做法存在两个致命缺陷：
1. 内部碎片化：为了应对最大可能长度（Max Sequence Length），系统必须为每个请求预留最大空间。如果实际生成的长度远低于最大值，大量显存被浪费。
2. 外部碎片化：由于请求的生命周期不同且长度不一，内存空间会变得支离破碎，导致无法分配大块连续空间给新请求。

PagedAttention：借鉴虚拟内存机制

vLLM 提出的 PagedAttention 彻底改变了这一现状。其核心思想是将 KV Cache 的存储方式从“连续数组”改为“分页存储”，类似于操作系统的虚拟内存管理。

工作原理

PagedAttention 将 KV Cache 分割为固定大小的“块”（Blocks）。每个块可以存储固定数量的 Token（例如 16 个）。
- 逻辑块 $\rightarrow$ 物理块：模型在逻辑上认为 KV Cache 是连续的，但物理上它们分布在不连续的显存块中。
- 块表（Block Table）：系统维护一张映射表，记录每个逻辑块对应的物理地址。
- 动态分配：只有当当前的物理块被填满时，系统才会为该请求分配一个新的物理块。

带来的提升

1. 近乎零浪费：除了最后一个块可能存在少量空隙外，所有显存都被高效利用。
2. 高效共享（Copy-on-Write）：在并行采样（Parallel Sampling）或多轮对话中，多个请求可以共享同一组前缀（Prefix）的物理块。只有当某个请求需要修改内容时，才触发 Copy-on-Write 操作复制该块。

工程实践中的权衡

虽然 PagedAttention 大幅提升了吞吐量，但在实际部署中仍需关注以下维度：
- Block Size 的选择：过小的 Block 会增加映射表的开销；过大的 Block 则会重新引入内部碎片化。通常 16 或 32 是平衡点。
- 量化压缩：为了进一步降低压力，业界采用了 FP8 或 INT8 量化 KV Cache $\text{KV_quant}$，甚至使用 GQA（Grouped Query Attention）来减少 K 和 V 头数。

总结

KV Cache 的演进路径是从“简单缓存” $\rightarrow$ “静态预分配” $\rightarrow$ “动态分页管理”。PagedAttention 不仅是算法的优化，更是对底层硬件资源管理的一次深刻重构。对于构建高性能 AI 系统而言，理解并优化内存布局比单纯追求算力堆砌更为关键。