为什么需要KVCache?
本blog主要介绍大模型推理中 KV Cache 的作用。先通过模拟自回归生成过程,说明模型是如何逐步生成文本的;随后进一步解释 KV Cache 为什么被提出。
1. 大模型生成内容的过程
KVCache一般只存在于含有Decoder的模型中,下面模拟会简化一些内容
1.1 规定符号
假定词表大小vocab_size=6,隐藏层维度d_model=5
词表W:
| 我 | 爱 | 学 | 习 | ! | ||
|---|---|---|---|---|---|---|
| token_id | 0 | 1 | 2 | 3 | 4 | 5 |
embedding矩阵E:
1.2 生成过程
输入: 我爱学
模型预期输出: 习!
下面我们来具体分析模型是怎么生成习!
- 模型通过词表
W,将输入文字转换为token_ids=[[0, 1, 2]],token_ids.shape=(1,3)=(batch_size, seq_len) - 模型通过词嵌入层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) - 接下来需要加上位置编码,但是我们省略这个过程
- 经过Casual Self-Attention(Masked Self-Attention)
- 最后经过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_Q | W_K | W_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计算得分
这里的我/我表示两个token的相关性Q[我]*K[我]
将V分块写为V=[v_0, v_1]
最终结果:
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矩阵并没有再利用的价值,所以不需要再进行缓存)
图片中灰色代表缓存的内容