KV Cache

KV Cache是LLM中加速运算的非常重要的一步，能够保证模型在对话过程中，输入文本越来越长却不影响推理速度。

如图所示，在LLM推理过程中，KV Cache占据了很大一部分的显存。

注意力机制

LLM的流程本质是输入n个token, 输出第n+1个token，当得到第n +1个token，此时再根据0~n+1个token预测第n+2个token,以此反复。

LLM中预测下一个token中，最关键的一步是Self Attention的计算，假设我们输入的N个token得到了N个q,k,v向量，用数学表示就是$q_i, k_i, v_i (i \in (0, n))$

Self Attention的过程是对N个token进行qkv计算，具体的公式为$\sum_{j = 0}^{n} Softmax( \frac{1}{\sqrt{d_k}} q_i^T k_j) v_j$, 去除固定系数$\frac{1}{\sqrt{d_k}}$, 得到$\sum_{j = 0}^n Softmax( q_i^T k_j ) v_j$

这样就得到了第$i$个token的注意力向量为$\sum_{j = 0}^n Softmax( q_i^T k_j) v_j$

因果掩码

在预测第 $i + 1$ 个token的过程中，llm只能看到 $0$ 到 $i$ 个位置的token，而不能看到 $i+1$ 以及之后的token。从数学上描述就是$\sum_{j = 0}^i Softmax( q_i^T k_j) v_j$ , 注意 $j$ 的范围变化。

KV Cache

可以注意到，计算$i + 1$个token的时候，需要对于每个 $q_i$都要计算$\sum_{j = 0}^i Softmax( q_i^T k_j) v_j$ ， 其中$k_j v_j$ 是共有的。

其次当计算$i + 2$个token的时候，需要对于每个$q_i$都要计算$\sum_{j = 0}^{i + 1} Softmax( q_i^T k_j) v_j = \sum_{j = 0}^i Softmax( q_i^T k_j ) v_j + q_i^T k_{j + 1} v_{j + 1}$ ， 又出现$k_j v_j$

因此kv可以进行缓存，当计算$i + 2$个token的时候，只需要再计算一下$k_{j + 1} v_{j + 1}$就行

从矩阵的角度更容易理解，如图计算出了过往所有Token的KV，公式为$QK^T V = Q (K^TV)$ 。在通过当前Token来计算下一个Token时，公式为$q_i K_{i * i}^TV_{i * i} = q_i k_i^T v_i + q_i K^T_{(i - 1) * (i - 1)} V_{(i - 1)* (i - 1)}$ , 因此缓存$K^TV$就行了

为什么没有Q Cache

从上面的推导可以知道，预测下一个token并不需要和以前token的Q向量进行计算，只需要使用当前token对应的Q向量，因此没必要缓存