大家都知道大模型是通过语言序列预测下一个词的概率。假定 x 1 , x 2 , x 3 , . . . , x n − 1 { x_1, x_2, x_3, ..., x_{n-1} } x1,x2,x3,...,xn−1 为已知序列,其中 x 1 x_1 x1, x 2 x_2 x2, x 3 x_3 x3, …, x n − 1 x_{n-1} xn−1 均为维度是 d m o d e l d_{model} dmodel 的向量, q n q_n qn, k n k_n kn, v n v_n vn 同为向量。当输入 x n x_n xn 时,需要预测 x n + 1 x_{n+1} xn+1 的概率分布。
KV Cache 干了什么?
Attention 机制的目标是输入 x n x_n xn,输出 z n z_n zn。在具体实现过程中,输入 x n x_n xn,生成 q n q_n qn、 k n k_n kn 和 v n v_n vn,并在实际计算中不再需要重复计算 k 1 k_1 k1, k 2 k_2 k2, …, k n − 1 k_{n-1} kn−1 和 v 1 v_1 v1, v 2 v_2 v2, …, v n − 1 v_{n-1} vn−1,直接从缓存中取即可。
具体 Attention 机制计算流程如下图所示。
观察注意力矩阵最下面一行(放大图我放下面了)。新输入的 x n x_n xn 通过矩阵 W q W_q Wq 生成 q n q_n qn,其中 q n q_n qn 与 k 1 k_1 k1, k 2 k_2 k2, …, k n k_n kn 均有运算关系。所以可以通过缓存 k 1 k_1 k1, k 2 k_2 k2, …, k n − 1 k_{n-1} kn−1 向量加速推理。这是 K 矩阵需要缓存的原因。
不过很意外的发现最右边一列 q 1 q_1 q1, q 2 q_2 q2, …, q n − 1 q_{n-1} qn−1 与 k n k_n kn 之间存在计算。
不是说好的只有 KV 缓存,没有 Q 矩阵缓存?如果推导成立,新输入 x n x_n xn 是否会改变 x 1 x_1 x1, x 2 x_2 x2, …, x n − 1 x_{n-1} xn−1 的注意力分布?
推导没有错,也没有 Q 矩阵缓存。因为在推理阶段,Attention 机制有一个非常重要的细节:mask 掩码
注意力矩阵在训练推理过程中,为了模拟真实推理场景,当前位置 token 是看不到下一位置的,且只能看到上一位置以及前面序列的信息,所以在训练推理的时候加了 attention mask
。具体实现如下图所示:
将上图灰色区域全部重置为 − i n f -inf −inf(负无穷大),这样方便 softmax 的时候置为 0。当新输入 x n x_n xn,注意力的计算(见注意力矩阵最下面一行)与 q 1 q_1 q1, q 2 q_2 q2, …, q n − 1 q_{n-1} qn−1 无关,因此无需缓存 Q 矩阵。
另外,还有个 V 矩阵,参照图1就干了一件事。
z n = a 1 ∗ v 1 + a 2 ∗ v 2 + . . . + a n ∗ v n z_n = a_1 * v_1 + a_2 * v_2 + ... + a_n * v_n zn=a1∗v1+a2∗v2+...+an∗vn
我可以提前缓存 v 1 v_1 v1, v 2 v_2 v2, …, v n − 1 v_{n-1} vn−1,计算的时候从缓存中取即可,这是 V 矩阵需要缓存的原因。