现代 AI 的“KV Cache”：为什么 LLM 越聊越慢，且显存会被吃光？

当你与大语言模型（LLM）进行长对话时，可能会发现两个现象：一是随着对话长度增加，生成速度逐渐变慢；二是显存占用在不断攀升，即使模型本身的大小没有变化。这背后的核心机制就是 KV Cache（Key-Value Cache）。

什么是 KV Cache？

要理解 KV Cache，首先要理解 Transformer 的注意力机制（Attention）。在生成每一个新 Token 时，模型需要回顾之前所有已经生成的 Token。

在计算注意力时，每个 Token 都会被转化为三个向量：Query (Q), Key (K), 和 Value (V)。
- Query：当前 Token “想要寻找什么”。
- Key：之前的 Token “能提供什么信息”。
- Value：之前的 Token “实际包含的内容”。

计算过程是：用当前的 $Q$ 与之前所有 Token 的 $K$ 做点积，计算权重，然后根据权重对 $V$ 进行加权求和。

关键点在于： 对于已经生成的 Token，$K$ 和 $V$ 向量在后续的生成步骤中是完全不变的。如果每次生成新 Token 都要重新计算一遍之前所有 Token 的 $K$ 和 $V$，那么计算量将随长度呈平方级增长（$O(n^2)$），这会导致极大的资源浪费。

为了优化，工程师将这些已经计算过的 $K$ 和 $V$ 存储在显存中，这就是 KV Cache。下次生成时，直接从内存读取旧的 $K, V$，只计算当前新 Token 的 $K, V$ 即可。

KV Cache 的代价：显存的“隐形杀手”

虽然 KV Cache 大幅提升了推理速度（将复杂度降为线性），但它引入了一个巨大的空间成本。

1. 显存占用公式

KV Cache 的大小取决于：
- 模型层数 (Layers)
- 隐藏层维度 (Hidden Dimension)
- 注意力头数 (Heads)
- 序列长度 (Sequence Length)
- 精度 (Precision, 如 FP16/BF16)

一个简单的估算公式（对于标准 Transformer）：
$\text{Memory} = 2 \times \text{Layers} \times \text{Heads} \times \text{Dim_per_head} \times \text{Seq_len} \times \text{Bytes_per_param}$

例如，一个 Llama-3-8B 模型（32层，32头，128维/头），使用 FP16（2字节）：
每个 Token 需要的缓存 $\approx 2 \times 32 \times 32 \times 128 \times 2 = 524,288$ 字节 $\approx 0.5\text{ MB}$。
如果上下文达到 32k tokens，仅 KV Cache 就需要 $\approx 16\text{ GB}$ 的显存！这还不包括模型本身的权重。

2. “碎片化”问题

传统的 KV Cache 分配像是在内存中申请连续的大块空间。但由于对话长度不可预测，系统往往会预先分配一个最大长度的缓冲区。这导致了严重的内存碎片化——很多预留空间被浪费了，而真正需要时却可能因为碎片过多而 OOM（Out of Memory）。

如何优化？（工业界的解决方案）

为了解决上述问题，现代 AI 系统采用了多种黑科技：

PagedAttention (vLLM)

这是目前最主流的方案。借鉴了操作系统的“虚拟内存”概念，将 KV Cache 分成固定大小的“页”（Pages）。不再要求连续存储，而是通过页表映射。这样可以几乎消除内存碎片，允许更高的并发量（Throughput）。

MQA 与 GQA (Multi-Query / Grouped-Query Attention)

既然 $K, V$ 太占空间，能不能少存点？
- MQA: 所有 Query 头共享同一组 $K, V$ 头。显存占用直接降低到原来的 $1/\text{Heads}$。但效果会有所下降。
- GQA: 将 Query 头分组，每组共享一组 $K, V$ 头（如 Llama-3 使用 GQA）。在性能和显存之间取得了极佳的平衡。

量化 (Quantization)

将 KV Cache 从 FP16 量化到 INT8 或 FP8。这能直接将缓存空间减半，且对模型精度的影响微乎其微。

总结

KV Cache 是 LLM 推理中的一种“空间换时间”策略。它让实时对话成为可能，但也成为了限制上下文长度和并发数的瓶颈。从 PagedAttention 到 GQA 再到量化，AI 工程化的核心目标之一就是如何更高效地管理这块昂贵的显存缓存。