【大模型推理】KV Cache原理

基本原理

两句话说明 KV Cache 的原理：

• 由于Attention中的mask机制（前面token的Q，看到不到后面token的K和V），当前生成的token只依赖于 当前输入的token以及之前所有token的KV矩阵
• 因此 可以通过 缓存 注意力机制中的键值对，在自回归生成过程中减少重复计算，加速大模型的推理过程

矩阵运算

要想理解 KV Cache，需要先了解一下矩阵运算的一个基本性质：分块矩阵计算法则，也就是两个矩阵相乘可以拆分成 列向量与行向量运算的和，在张量并行中也有用到：

$$\left[\begin{array}{ll}{X}_1& X_2\end{array}\right]\times \left[\begin{array}{l}{A}_1\\ A_2\end{array}\right]=\left[X_1{A}_1+X_2{A}_2\right]=Y$$

在这个性质的基础上，此时如果矩阵 $X$ 是一个下三角矩阵，就有：

$$\begin{array}{rl}& \left[\begin{array}{cccc}{X}_{1,1}& 0& \cdots & 0\\ X_{2,1}& X_{2,2}& \cdots & 0\\ ⋮& ⋮& \ddots & ⋮\\ X_{m,1}& X_{m,2}& \cdots & X_{m,n}\end{array}\right]\cdot \left[\begin{array}{cccc}{Y}_{1,1}& Y_{1,2}& \cdots & Y_{1,p}\\ Y_{2,1}& Y_{2,2}& \cdots & Y_{2,p}\\ ⋮& ⋮& \ddots & ⋮\\ Y_{n,1}& Y_{n,2}& \cdots & Y_{n,p}\end{array}\right]\\ & =\left[\begin{array}{l}{X}_{1,1}{\vec{Y}}_1\\ X_{2,1}{\vec{Y}}_1+X_{2,2}{\vec{Y}}_2\\ ⋮\\ X_{m,1}{\vec{Y}}_1+X_{m,2}{\vec{Y}}_2+\cdots +X_{m,n}{\vec{Y}}_n\end{array}\right]\end{array}$$

可以看到，结果矩阵的第 $k$ 行只用到了矩阵 $X$ 的 第 $k$ 个行向量。所以 $X$ 不需要进行全部的矩阵乘法，每一步只取第 $k$ 个行向量即可，这就很大程度上减少了计算量，也就是 KV Cache 的数学原理。

如果把矩阵 $X$ 认为是 Attention 计算得到的 QK score，矩阵 $Y$ 认为是 $V$，上述计算流程基本上就是 KV Cache 的框架了。这里要注意一点，因为 现在大模型基本都是Decoder-only的架构 ，自回归生成的过程中，当前token在做attention计算时是看不到后面的，后面的内容都被 mask 了。所以 KV Cache 只有 Decoder-only架构才有，Encoder-only 比如 BERT模型 是没有的

KV Cache原理

下面用一个具体的例子来说明，KV Cache假设当前文本为 我要, 要 大模型 输出的内容为学习AI

第一次计算

模型第一次计算，生成学时的过程如下：

为了方便演示，忽略scale项$\sqrt{d}$，最终Attention的计算公式如下，（softmaxed 表示已经按行进行了softmax）：

$$\begin{array}{rl}At{t}_{\text{step }1}(Q,K,V)& =\mathrm{softmax}\left(\left[\begin{array}{cc}{Q}_1{K}_1^T& -\infty \\ Q_2{K}_1^T& Q_2{K}_2^T\end{array}\right]\right)\left[\begin{array}{l}\overrightarrow{V_1}\\ \overrightarrow{V_2}\end{array}\right]\\ & =\left(\left[\begin{array}{cc}\mathrm{softmaxed}\left(Q_1{K}_1^T\right)& 0\\ \mathrm{softmaxed}\left(Q_2{K}_1^T\right)& \mathrm{softmaxed}\left(Q_2{K}_2^T\right)\end{array}\right]\right)\left[\begin{array}{l}\overrightarrow{V_1}\\ \overrightarrow{V_2}\end{array}\right]\\ & =\left(\left[\begin{array}{c}\mathrm{softmaxed}\left(Q_1{K}_1^T\right)\times \overrightarrow{V_1}\\ \mathrm{softmaxed}\left(Q_2{K}_1^T\right)\times \overrightarrow{V_1}+\mathrm{softmaxed}\left(Q_2{K}_2^T\right)\times \overrightarrow{V_2}\end{array}\right]\right)\end{array}$$

假设$At{t}_1$表示Attention结果的第一行，$At{t}_2$表示Attention结果的第二行，那么有下面的表示：

$$\begin{array}{rl}& At{t}_1(Q,K,V)=\mathrm{softmaxed}\left(Q_1{K}_1^T\right){\vec{V}}_1\\ & At{t}_2(Q,K,V)=\mathrm{softmaxed}\left(Q_2{K}_1^T\right){\vec{V}}_1+\mathrm{softmaxed}\left(Q_2{K}_2^T\right){\vec{V}}_2\end{array}$$

可以看到：

• 在计算$At{t}_1$时，$Q_1{K}_2^T$这个score会被mask掉，也就是当前token计算Attention时看不到后面的token信息
• 在计算$At{t}_2$时，仅仅依赖于$Q_2$以及$K_1^T,K_2^T$和${\vec{V}}_1,{\vec{V}}_2$，与$Q_1$无关

第二次计算

如果没有KV Cache，模型进行第二次前向推理计算，生成习时，计算过程如下：

此时，Attention的计算公式为：

$$\begin{array}{rl}& {\mathrm{Att}}_1(Q,K,V)=\mathrm{softmaxed}\left(Q_1{K}_1^T\right){\vec{V}}_1\\ & {\mathrm{Att}}_2(Q,K,V)=\mathrm{softmaxed}\left(Q_2{K}_1^T\right){\vec{V}}_1+\mathrm{softmaxed}\left(Q_2{K}_2^T\right){\vec{V}}_2\\ & {\mathrm{Att}}_3(Q,K,V)=\mathrm{softmaxed}\left(Q_3{K}_1^T\right){\vec{V}}_1+\mathrm{softmaxed}\left(Q_3{K}_2^T\right){\vec{V}}_2+\mathrm{softmaxed}\left(Q_3{K}_3^T\right){\vec{V}}_3\end{array}$$

同样的，$At{t}_3$的计算结果，之和$Q_3$有关，与$Q_1,Q_2$无关，因此，可以模型的推理可以简化为：

这就是KV Cache的计算过程

第三次计算

同理，生成最后AI时，使用KV Cache计算过程如下：

$At{t}_4$的计算公式为：

$$\begin{array}{rl}{\mathrm{Att}}_4(Q,K,V)& =\mathrm{softmaxed}\left(Q_4{K}_1^T\right)\overrightarrow{V_1}+\mathrm{softmaxed}\left(Q_4{K}_2^T\right)\overrightarrow{V_2}\\ & +\mathrm{softmaxed}\left(Q_4{K}_3^T\right)\overrightarrow{V_3}+\mathrm{softmaxed}\left(Q_4{K}_4^T\right)\overrightarrow{V_4}\end{array}$$

因此，可以看出：

• 不使用KV Cache的方法，存在大量冗余的计算，也就是要生成$At{t}_k$时，还需要重复计算$At{t}_1,..,At{t}_{k-1}$
• 计算$At{t}_k$时，之和$Q_k$有关，与之前的$Q_1,...,Q_{k-1}$都没关系
• 生成第$x_k$个token时，只需要输入上一轮生成的$x_{k-1}$即可

所以每一步其实只需要根据$Q_k$ 计算 $At{t}_k$ 就可以，但是 $K$ 和 $V$ 是全程参与计算的。从优化推理速度角度来看，只需要把每一步的 $K,V$ 缓存起来就可以， 所以叫 KV Cache。

最后需要注意当 sequence 比较长，或者 batch 特别大的时候，KV Cache 其实还是个Memory刺客，所以如何减少 KV 的内存变得尤为重要。

目前各种框架，针对 KV Cache 做了优化，比如 vLLM 的 Page Attention, Prefix Caching，Token 的稀疏化，KV 共享或者压缩（MQA、GQA 和 MLA），LayerSkip，Mooncake 等等，可以说 KV Cache 目前是推理的基石，各种基于 KV Cache 的优化方法撑起了大模型推理加速的半壁江山。

参考资料

• [1] https://note.mowen.cn/note/detail?noteUuid=Uwudr7Hu_STXSicuxHei2