0x06 KV-Cache
LLM 生成过程做的是 Next Token Prediction，也就是输入前 \(n\) 个 token，预测第 \(n+1\) 个 token。首先考察做 Self-Attention 时，按照正常计算方式是如何运作的。下面举例输入是三个 token 的情况，也就是根据前三个 token 预测第四个 token：
\[
O = \begin{bmatrix}
O_0\\
O_1\\
O_2
\end{bmatrix}
=
\mathrm{softmax}\!\left(
\frac{
\begin{bmatrix}
Q_0K_0^\top & Q_0K_1^\top & Q_0K_2^\top\\
Q_1K_0^\top & Q_1K_1^\top & Q_1K_2^\top\\
Q_2K_0^\top & Q_2K_1^\top & Q_2K_2^\top
\end{bmatrix}
}{\sqrt{d_k}}
\right)
\begin{bmatrix}
V_0\\
V_1\\
V_2
\end{bmatrix}.
\]
注意，在计算 \(QK^\top\) 之后，还会对 Attention Score 矩阵进行下三角的 Mask，也就是实际有效的注意力分数为：
\[
\begin{bmatrix}
Q_0K_0^\top & 0 & 0\\
Q_1K_0^\top & Q_1K_1^\top & 0\\
Q_2K_0^\top & Q_2K_1^\top & Q_2K_2^\top
\end{bmatrix}.
\]
从计算中可以看到，实际参与到第三个 token 预测 \(O_2\) 中的只有 Attention Score 计算的最后一行，和最后的 \(V\) ，也就是：
\[
\begin{aligned}
\mathrm{softmax}\!\left(\frac{[\,Q_3K_1^\top,\;Q_3K_2^\top,\;Q_3K_3^\top\,]}{\sqrt{d_k}}\right)
\begin{bmatrix} V_1\\ V_2\\ V_3 \end{bmatrix}.
\end{aligned}
\]
在这个计算中包含之前历史的 \(K、V\) 以及当前输入的第三个 token 对应的 \(Q_3\)。
简单推广一下
- 在输入 \(n\) 个 token 预测第 \(n+1\) 个 token 的时候，会用到前 \(n\) 个 token 的 \(K,V\) 以及第 \(n\) 个 token 的 \(Q\)
- 在输入 \(n+1\) 个 token 预测 \(n+2\) 个 token 时，会用到前 \(n+1\) 个 token 的 \(K,V\) 以及第 \(n + 2\) 个 token 的 \(Q\)。
不难发现，这其中包括对历史 KV 的重复使用，并且这些 KV 在后续逐个 token 的生成过程中，并不会发生改变，因此避免对这些 KV 的重复计算，很直接的方式就是将他们缓存起来，达到减少计算，提高升成速度的效果。
- 为什么不把 Attention Score 矩阵（\(QK^\top\)）整个缓存？
- 为什么不缓存历史的 \(Q\)