@misc{kv-cache,
  author = {Jingxi Qiu},
  title = {大模型 KV Cache},
  howpublished = {Jingxi Qiu},
  year = {2024},
  url = {https://jingxiqiu.com/blog/kv-cache/}
}

Qiu, J. (2024). 大模型 KV Cache. In Jingxi Qiu.

Qiu, Jingxi. “大模型 KV Cache.” Jingxi Qiu, 2024.

Qiu, Jingxi. 2024. “大模型 KV Cache.” In Jingxi Qiu. Preprint.

KV-Cache#

KV-Cache是目前Transformer常用的功能，应用于Decoder架构的推理加速。

Self-Attention 不带KV-Cache流程#

以下流程来自知乎看图学[1]

对于一段文字输入，有以下计算过程：

1. 首先输入一个token的embedding，大小为[1, emb_size]

经过公式计算得出attention向量，这里为方便表示$\sqrt{d}$被省略

$$Att_1(Q,K,V)=softmax(\frac{Q_1,K^T_1}{\sqrt{d}})\overset{\rightarrow}{V}$$

1. 输入第二个token

在这里可以发现，attention的计算公式为：

$$\begin{align}Att_2(Q,K,V)&=softmax\left(\left[\begin{matrix}Q_1K_1^T -\infty \\Q_2K_1^T Q_2K_2^T\end{matrix}\right]\right)\left[\begin{matrix}\overset{\rightarrow}{V_1} \\ \overset{\rightarrow}{V_2}\end{matrix}\right] \\ &=\left(\left[\begin{matrix}softmax(Q_1K_1^T)\overset{\rightarrow}{V_1}+0 \\ softmax(Q_2K_1^T)\overset{\rightarrow}{V_1} + softmax(Q_2K_2^T)\overset{\rightarrow}{V_2} \end{matrix}\right]\right)\end{align}$$

因此可以发现：

$$\begin{align}Att_1(Q,K,V)&=softmax(Q_1K_1^T)\overset{\rightarrow}{V_1} \\ Att_2(Q,K,V)&=softmax(Q_2K_1^T)\overset{\rightarrow}{V_1} + softmax(Q_2K_2^T)\overset{\rightarrow}{V_2} \end{align}$$

随后的attention也同理可得

$$Att_i(Q,K,V)=\sum^i_{j=1}softmax(Q_iK^T_j)\overset{\rightarrow}{V_j}$$

因此可以得出结论：

1. 当前计算方式存在大量冗余，计算量可以优化
2. Attention只与$Q_k$相关
3. 推理第i个字符只需要i-1即可

推理加速示意图#

引用：#

[1] https://zhuanlan.zhihu.com/p/662498827