KV Cache

本文结合gemini-3.1-pro-preview生成。

简单来说，KV Cache 是一种空间换时间的推理加速技术，它只在模型生成文本（推理阶段）时使用，目的是避免重复计算。

痛点：为什么要搞个 Cache？

大语言模型（如 GPT、LLaMA）生成文本的方式叫自回归（Autoregressive），也就是“逐字生成”。 假设我们要生成一句话，已经生成了前三个词[x1, x2, x3]，现在要预测第 4 个词 x4。

在标准的 Transformer 中，Self-Attention公式： $ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V $

如果你不加干预，在第 3 步时，模型会把 x1, x2, x3 全部重新乘以权重矩阵 \(W_q\), \(W_k\), \(W_v\)，算出完整的 Q, K, V矩阵，然后再做一遍上面的公式。 但仔细一想：在生成第 4 个词时，前面的词 x1, x2, x3的特征其实在前面几步早就已经算过了啊！

这种每次把历史前缀全部重算一遍的做法，计算量会随着句子变长呈平方级爆炸\(O(N^2)\)，太浪费了。

原理：KV Cache 是怎么偷懒的？

因为大模型有因果掩码（Causal Mask），也就是“前面的词不能看后面的词”。因此，对于同一个词 x1，无论它后面接的是什么词，它自己的 Key 和 Value 在同一层里是永远不变的。

KV Cache 的核心思想就是：把过去每个词算出来的 Key 和 Value 保存下来。

我们来看看加上 KV Cache 后的数学推导：

假设现在是第 t 步（前面的 1 到 t-1 个词已经生成好了），当前我们刚刚输入了第 t 个词的特征向量 \(x_t\)（注意，当前时刻的输入只有这一个词，而不是完整的序列）。

  1. 计算当前的 Q, K, V： 我们只计算这单个词 \(x_t\) 的映射： \(q_t = x_t W_q\) , \(k_t = x_t W_k\) , \(v_t = x_t W_v\) （注意这里的 \(q_t, k_t, v_t\) 都是 \(1 \times d\) 的行向量）

  2. 更新缓存（Cache）： 我们把刚算出来的 \(k_t\) 和 \(v_t\) 拼接到之前的缓存后面：

\(K_{\text{cache}} = [k_1, k_2, \dots, k_{t-1}, k_t]\)

\(V_{\text{cache}} = [v_1, v_2, \dots, v_{t-1}, v_t]\) （此时 \(K_{\text{cache}}\) 和 \(V_{\text{cache}}\) 是 \(t \times d\) 的矩阵）

  1. 计算当前步的 Attention： 现在，我们只需要用当前时刻的 Query (\(q_t\)) 去和包含了所有历史信息的 \(K_{\text{cache}}\)计算注意力得分，再乘上 \(V_{\text{cache}}\)：\(\text{Output}_t = \text{softmax}\left(\frac{q_t K_{\text{cache}}^T}{\sqrt{d_k}}\right) V_{\text{cache}}\)

看懂和上面的Self-Attention公式的区别了吗？ 原来我们需要做的是一个巨大的矩阵相乘，现在变成了一个行向量去乘一个大矩阵。由于我们缓存了过去的 K 和 V，计算复杂度从 O(\(t^2\)) 骤降到了 O(t)。

小问题：为什么不缓存 Q（Query）？ 因为 Query 的物理意义是“我当前需要什么信息”。我们现在只关心第 t 个词需要什么信息，不需要去管过去的词需要什么信息（它们已经生成完了），所以只有当前词的 \(q_t\) 参与运算，之前的 Q 都可以丢掉。

3. 代价：现在的挑战与前沿技术

天下没有免费的午餐。KV Cache 虽然极大地节省了算力，但它非常吃显存。

我们可以简单算一笔账：对于 LLaMA-7B 这样一个小模型，如果以 fp16（每个数字 2 字节）运行，每个 token 在每一层都需要存 K 和 V。 如果batch_size=1，序列长度为 4096，那么光是存一个请求的 KV Cache 就需要大约 2GB（总显存 = 层数32 × 2(K和V) × KV头数32 × 单头维度128 × 序列长度4096 × 每个元素字节数2 × batch_size） 的显存。如果并发几十个用户，显存瞬间就会被Cache 撑爆。这就导致了现在大模型推理通常不是算力瓶颈（Compute-bound），而是显存/内存带宽瓶颈（Memory-bound）。

为了解决 KV Cache 太大的问题，近年来出现了很多基于此的改进，看论文可能会碰到：

  1. MQA (Multi-Query Attention) / GQA (Grouped-Query Attention):

  • 思路： 既然 K 和 V 这么占地方，那我让所有的注意力头（Attention Heads）共享同一组 K 和 V（MQA），或者每几个头共享一组 K 和 V（GQA）。

  • 效果： 比如 LLaMA-2 70B 和 LLaMA-3，就是用了 GQA，把 KV Cache 的大小直接缩减到了原来的 1/8。

  1. PagedAttention (vLLM 框架的核心):

  • 思路： KV Cache 随着句子变长是动态增长的，在显存里容易产生大量内存碎片。PagedAttention借用了操作系统的“虚拟内存分页”思想，把 KV Cache切成一块块的“页”（Block），按需非连续分配。

  • 效果： 显存浪费几乎降为 0，大大提升了并发能力。

  1. Kimi (Moonshot) 等超长上下文模型：
  • 它们能支持 200万 token 的上下文，主要攻克的技术难点之一就是如何存储和快速检索这庞大的、以百 GB 计的 KV Cache。
posted @ 2026-05-26 19:44  wljss  阅读(14)  评论(0)    收藏  举报