为什么KV Cache只有KV，没有Q

大家都知道大模型是通过语言序列预测下一个词的概率。假定 $x_1,x_2,x_3,...,x_{n-1}$ 为已知序列，其中 $x_1$, $x_2$, $x_3$, …, $x_{n-1}$ 均为维度是 $d_{model}$ 的向量，$q_n$, $k_n$, $v_n$ 同为向量。当输入 $x_n$ 时，需要预测 $x_{n+1}$ 的概率分布。

KV Cache 干了什么？

Attention 机制的目标是输入 $x_n$，输出 $z_n$。在具体实现过程中，输入 $x_n$，生成 $q_n$、$k_n$ 和 $v_n$，并在实际计算中不再需要重复计算 $k_1$, $k_2$, …, $k_{n-1}$ 和 $v_1$, $v_2$, …, $v_{n-1}$，直接从缓存中取即可。

具体 Attention 机制计算流程如下图所示。

观察注意力矩阵最下面一行（放大图我放下面了）。新输入的 $x_n$ 通过矩阵 $W_q$ 生成 $q_n$，其中 $q_n$ 与 $k_1$, $k_2$, …, $k_n$ 均有运算关系。所以可以通过缓存 $k_1$, $k_2$, …, $k_{n-1}$ 向量加速推理。这是 K 矩阵需要缓存的原因。

不过很意外的发现最右边一列 $q_1$, $q_2$, …, $q_{n-1}$ 与 $k_n$ 之间存在计算。

不是说好的只有 KV 缓存，没有 Q 矩阵缓存？如果推导成立，新输入 $x_n$ 是否会改变 $x_1$, $x_2$, …, $x_{n-1}$ 的注意力分布？

推导没有错，也没有 Q 矩阵缓存。因为在推理阶段，Attention 机制有一个非常重要的细节：mask 掩码

注意力矩阵在训练推理过程中，为了模拟真实推理场景，当前位置 token 是看不到下一位置的，且只能看到上一位置以及前面序列的信息，所以在训练推理的时候加了 attention mask。具体实现如下图所示：

将上图灰色区域全部重置为 $-inf$（负无穷大），这样方便 softmax 的时候置为 0。当新输入 $x_n$，注意力的计算（见注意力矩阵最下面一行）与 $q_1$, $q_2$, …, $q_{n-1}$ 无关，因此无需缓存 Q 矩阵。

另外，还有个 V 矩阵，参照图1就干了一件事。

$$z_n=a_1\ast v_1+a_2\ast v_2+...+a_n\ast v_n$$

我可以提前缓存 $v_1$, $v_2$, …, $v_{n-1}$，计算的时候从缓存中取即可，这是 V 矩阵需要缓存的原因。