现代 AI 的“记忆管理”:KV Cache 与 PagedAttention 的工程真相

在 LLM 的推理过程中,最昂贵的资源不是计算量(FLOPs),而是内存带宽。当你与 AI 对话时,模型需要回顾之前所有的对话历史。如果每次生成新 Token 都要重新计算一遍之前的所有 Key 和 Value 向量,推理速度将呈平方级下降。

专属插画
现代 AI 的“记忆管理”:KV Cache 与 PagedAttention 的工程真相

现代 AI 的“记忆管理”:KV Cache 与 PagedAttention 的工程真相

在 LLM 的推理过程中,最昂贵的资源不是计算量(FLOPs),而是内存带宽。当你与 AI 对话时,模型需要回顾之前所有的对话历史。如果每次生成新 Token 都要重新计算一遍之前的所有 Key 和 Value 向量,推理速度将呈平方级下降。

为了解决这个问题,工业界引入了 KV Cache(键值缓存)。但 KV Cache 带来了另一个噩梦:内存碎片化。

KV Cache:用空间换时间的权衡

在 Transformer 的解码阶段,每个 Token 的生成依赖于之前所有 Token 的注意力权重。这意味着对于第 $n$ 个 Token,我们需要它与前 $n-1$ 个 Token 的 $\text{Key}$ 和 $\text{Value}$ 向量进行点积运算。

由于前 $n-1$ 个 Token 在之前的步骤中已经计算过了,且其结果在推理过程中保持不变,我们可以将这些 $\text{K}$ 和 $\text{V}$ 向量存储在显存中。这样,每一步只需要计算当前 Token 的 $\text{K}$ 和 $\text{V}$ 并将其追加到缓存中即可。

代价是显存的剧增。对于一个 Llama-3-70B 模型,在 FP16 精度下,单请求的 KV Cache 占用量随序列长度线性增长。当并发请求增加时,显存会被迅速填满,导致 OOM(Out of Memory)。

静态分配的困境:内存碎片化

早期的推理框架(如 HuggingFace Transformers)采用的是连续内存分配策略。为了保证性能,系统会为每个请求预先分配一块足够大的连续显存空间(例如最大支持 4096 个 Token)。

这种做法导致了严重的内存浪费:
1. 内部碎片:如果用户只输入了 10 个 Token,剩下的 4086 个位置依然被占用且无法给他人使用。
2. 外部碎片:随着请求的动态增加和释放,显存中会出现大量不连续的小空隙,无法容纳新的大请求。
3. 预留浪费:为了防止序列增长导致崩溃,开发者不得不保守地设置最大长度,进一步降低了吞吐量。

PagedAttention:借鉴操作系统的虚拟内存

vLLM 团队提出的 PagedAttention 从根本上改变了这一现状。其核心思想是将 LLM 的 KV Cache 类比为操作系统的虚拟内存页(Virtual Memory Pages)。

核心机制

PagedAttention 不再要求 KV Cache 在物理显存中连续存储,而是将其划分为固定大小的块(Blocks)
- 逻辑块 $\rightarrow$ 物理块:模型在逻辑上看到的是连续的序列,但底层通过一个“块表”(Block Table)将逻辑索引映射到不连续的物理显存块中。
- 动态按需分配:只有当当前的物理块填满时,系统才会为该请求分配一个新的物理块。
- 高效共享:在并行采样(Parallel Sampling)或 Beam Search 时,多个输出序列可以共享同一个前缀(Prompt)的物理块,无需重复存储相同的 KV 数据。

工程效果

通过 PagedAttention, KV Cache 的内存利用率从之前的 $60\% \sim 80\%$ 提升到了接近 $96\%$ 以上。这意味着在同样的硬件条件下,单机能够承载的并发请求数(Batch Size)提升了数倍,直接降低了推理成本并提高了响应速度。

总结:从算法到工程的闭环

KV Cache 是对 Transformer 计算冗余的优化,而 PagedAttention 是对 KV Cache 存储冗余的优化。这揭示了现代 AI 系统的一个关键趋势:模型能力的上限由算法决定,但商业落地的下限由工程实现决定。当我们讨论“推理加速”时,关注点已从单纯的算子优化转向了更深层的资源调度与内存管理方案。

留言区

欢迎分享你的想法!

发表留言

0/500

加载留言中…