为什么需要KVCache?

0 阅读3分钟

为什么需要KVCache?

本blog主要介绍大模型推理中 KV Cache 的作用。先通过模拟自回归生成过程,说明模型是如何逐步生成文本的;随后进一步解释 KV Cache 为什么被提出。

1. 大模型生成内容的过程

KVCache一般只存在于含有Decoder的模型中,下面模拟会简化一些内容

1.1 规定符号

假定词表大小vocab_size=6,隐藏层维度d_model=5

词表W

!
token_id012345

embedding矩阵E

ER6×5=[x00x01x02x03x04x10x11x12x13x14x20x21x22x23x24x30x31x32x33x34x40x41x42x43x44x50x51x52x53x54]=[x0,x1,x2,x3,x4,x5]E \in \mathbb{R}^{6\times5} = \begin{bmatrix} x_{00} & x_{01} & x_{02} & x_{03} & x_{04} \\ x_{10} & x_{11} & x_{12} & x_{13} & x_{14} \\ x_{20} & x_{21} & x_{22} & x_{23} & x_{24} \\ x_{30} & x_{31} & x_{32} & x_{33} & x_{34} \\ x_{40} & x_{41} & x_{42} & x_{43} & x_{44} \\ x_{50} & x_{51} & x_{52} & x_{53} & x_{54} \end{bmatrix}=[x_0,x_1,x_2,x_3,x_4,x_5]

1.2 生成过程

输入: 我爱学

模型预期输出: 习!

下面我们来具体分析模型是怎么生成习!

  1. 模型通过词表W,将输入文字转换为token_ids=[[0, 1, 2]],token_ids.shape=(1,3)=(batch_size, seq_len)
  2. 模型通过词嵌入层embedding矩阵E,将token_ids转换为词向量表示,e=\begin{bmatrix}\begin{bmatrix}x_0 & x_1 & x_2\end{bmatrix}\end{bmatrix},e.shape=(1,3,5)=(batch_size, seq_len, d_model)
  3. 接下来需要加上位置编码,但是我们省略这个过程
  4. 经过Casual Self-Attention(Masked Self-Attention)
  5. 最后经过Linear+SoftMax进行采样
1.2.1 Casual Self-Attention

在模拟的这个注意力机制中,我们当单头注意力进行模拟

输入: 矩阵e=\begin{bmatrix}\begin{bmatrix}x_0 & x_1 & x_2\end{bmatrix}\end{bmatrix},e.shape=(1,3,5)

约定符号:

W_QW_KW_V
模拟中的形状(5,2)(5,2)(5,2)
形式化的形状(d_{model},d_q)(d_{model},d_q)(d_{model},d_v)

将矩阵e分别通过W_Q,W_K,W_V线性变换得到Q,K,V矩阵,Q.shape=(3,2),K.shape=(3,2),V.shape=(3,2)

再通过公式SoftMax(Masked(QK^T))V计算得分

image-20260205170031696.png

这里的我/我表示两个token的相关性Q[我]*K[我]

QKT=[y00y01y02y10y11y12y20y21y22]QK^T = \begin{bmatrix} y_{00} & y_{01} & y_{02} \\ y_{10} & y_{11} & y_{12} \\ y_{20} & y_{21} & y_{22} \end{bmatrix}
Masked(QKT)=[y00y10y11y20y21y22]Masked(QK^T) =\begin{bmatrix}y_{00} & -\infty & -\infty \\y_{10} & y_{11} & -\infty \\y_{20} & y_{21} & y_{22}\end{bmatrix}
P=SoftMax(Masked(QKT))=[p0000p10p110p20p21p22]=[p0p1p2]P=SoftMax(Masked(QK^T)) =\begin{bmatrix}p_{00} & 0&0 \\p_{10} & p_{11} & 0\\p_{20} & p_{21} & p_{22}\end{bmatrix}=\begin{bmatrix}p_0 \\ p1 \\ p2\end{bmatrix}

将V分块写为V=[v_0, v_1]

最终结果:

A=PV=SoftMax(Masked(QKT))V=[p0v0p0v1p1v0p1v1p2v0p2v1]A=PV=SoftMax(Masked(QK^T))V= \begin{bmatrix} p_0v_0 & p_0v_1 \\ p_1v_0 & p_1v_1 \\ p_2v_0 & p_2v_1 \end{bmatrix}
1.2.2 Final Linear + SoftMax采样

输入: 经过Casual Attention后矩阵A,A.shape=(3,2)

此时A无法与E.T进行相乘,所以需要先给A做一次线性变换W_O,W_O.shape=(2,5)

矩阵logitis=AW_OE^T,logits.shape=(6,6)

取logits最后一行进行采样,即SoftMax(logitis[-1]),假定结果为[0.1, 0.1, 0.1, 0.4, 0.1, 0.2],概率最高的为index=3,为

自回归: 将习加入输入,得到我爱学习,将这个句子重新作为输入,重复上面的推理过程,再得到输出 !

2. KVCache

我们重新回顾上面推理过程:

输入1: 我爱学

通过一系列过程,最后采样得到,将习加入输入

输入2: 我爱学习

通过一系列过程,最后采样得到 ! ,将!加入输入,直到最后采样到或者满足其它条件,才终止采样

显然两次输入存在相同的前缀我爱学,那么在我们进行第二次推理时,能否利用第一次推理过程中的部分内容呢,从上面的推理过程可以知道两次推理的相同前缀的K,V矩阵是完全相同的(甚至Q矩阵也是相同的,但是之前的Q矩阵并没有再利用的价值,所以不需要再进行缓存)

图片中灰色代表缓存的内容

image-20260205181445414.png

3. 参考资料

The Illustrated Transformer

【8】KV Cache 原理讲解