现代 AI 系统的“内存墙”：KV Cache 的压力与优化路径

在 LLM（大语言模型）的推理过程中，最昂贵的资源往往不是计算量（FLOPs），而是内存带宽。当我们谈论 AI 系统的性能瓶颈时，一个核心概念是 KV Cache（Key-Value Cache）。本文将深入探讨 KV Cache 的本质、它如何造成“内存墙”，以及工业界目前主流的优化方案。

什么是 KV Cache？

Transformer 架构的核心是注意力机制（Attention）。在生成文本时，模型采用自回归方式：每产生一个新 Token，都需要回顾之前所有已生成的 Token。

如果没有缓存，每生成第 $N$ 个 Token，模型都要重新计算前 $N-1$ 个 Token 的 Key 和 Value 向量。这意味着计算量随序列长度呈平方级增长。为了避免重复计算，我们将之前步骤产生的 Key 和 Value 向量存储在显存中，这就是 KV Cache。

“内存墙”的成因

虽然 KV Cache 解决了计算冗余，但它引入了巨大的内存压力：

1. 空间占用惊人：KV Cache 的大小与模型层数、隐藏层维度、序列长度以及 Batch Size 成正比。对于一个 Llama-3-70B 模型，在 FP16 精度下，单个请求的 KV Cache 在长文本场景下可达数 GB。
2. 带宽受限 (Memory Bound)：推理过程是典型的访存密集型任务。GPU 计算核心的速度远快于显存读取速度。在解码阶段（Decoding phase），模型需要将巨大的 KV Cache 从 HBM（高带宽内存）搬运到 SRAM 中进行计算，导致 GPU 大部分时间在“等待数据”，而非执行计算。
3. 碎片化问题：传统的显存分配方式会导致大量碎片（External Fragmentation），使得实际可用显存低于物理上限，限制了最大并发数（Throughput）。

工业界的优化路径

为了打破这堵“墙”，目前的优化方向主要集中在三个维度：减少存储量、提高访存效率、优化内存管理。

1. 减少存储量：MQA 与 GQA

早期的 Multi-Head Attention (MHA) 为每个 Query Head 分配一对 KV Head。
- MQA (Multi-Query Attention)：所有 Query Head 共用一组 KV Head。极大地降低了缓存大小，但可能损失部分精度。
- GQA (Grouped-Query Attention)：折中方案。将 Query 分组，每组共用一组 KV Head。Llama-3 等主流模型均采用 GQA，在性能和精度之间取得了平衡。

2. 提高访存效率：PagedAttention (vLLM)

借鉴操作系统虚拟内存的思想，vLLM 提出了 PagedAttention。它不再为每个请求分配连续的显存块，而是将 KV Cache 分页存储在不连续的物理块中。
- 消除碎片：几乎消除了内部和外部碎片。
- 动态共享：允许不同的请求共享相同的前缀缓存（如相同的 System Prompt），极大提升了吞吐量。

3. 量化与压缩

将 KV Cache 从 FP16/BF16 量化为 INT8 或 FP8，甚至更低位宽（如 4-bit）。这直接将内存占用减半或更多，并减轻了带宽压力，而对模型生成质量的影响在可控范围内。

总结

KV Cache 是 LLM 推理的必然产物，也是性能优化的主战场。从算法层面的 GQA 到系统层面的 PagedAttention，再到硬件友好的量化技术，其核心目标只有一个：让数据搬运的速度跟上计算的速度。对于开发者而言，理解这一机制有助于在选择部署框架和配置并发参数时做出更科学的决策。