结合图例，理解transformer推理加速技术之KV Cache

先要明确：KV cache是优化transformer的decoder端；decoder的输入有两部分，encoder的输出和decoder先前预测的结果，我们重点关注后者。

一、回顾传统transformer解码过程

对于一个翻译任务，输入[<我>，<爱>，<你>]，期待输出[< I >，< love >，< you >]。

假设：模型的decoder只有一层，且目前已经执行了两轮，预测出了[< I >, < love >]，现在正在预测第3个word，那么decoder的输入端要包括前面的输出，即[< start >，< I >，< love >]（当然还需要包含起始符）。接下来注意力机制的相关计算如下：

transformer模型不管是decoder还是encoder输入输出都是一致的size，但是decoder端我们只取最后一个维度作为本轮输出，也就是< you >作为本次decoder的预测输出，再与之前的输出[< I >, < love >]拼接得到结果[< I >, < love >，< you >]（之后还要预测到终止符或直接预测到固定长度，这里暂不考虑）。

二、传统注意力模块冗余分析

分析一下，在decoder端做自回归预测时，我们是否需要每次都计算输入decoder的所有词的q、k、v特征？

上图从右向左看；

我们本层最终得到的结果< you >，事实上在上一步只与softmax后的score的最后一行和当前输入的全部词的特征v相关，而softmax操作是按行计算的，那么再反推到上一轮，改变它的是输入特征q的最后一维（也就是当前输入最后一个word的q特征，即仅从< love >提取的特征q）和全部词的特征k；

显然，上图白色部分数据的加载或计算都没有意义，因为后续没有用到（attn操作后还有一个全连接层，是对v的行做操作，不会影响上图颜色的分布）。

因此，当预测< you >时[< start >, < I >]的q不需要计算，事实上[< start >, < I >]的k和v我们也不用再计算，因为decoder是反复执行的，预测< love >的时候已经计算过[< start >, < I >]的k\v，而< start >的k\v也在预测< I >的时候算过了，那么，我们可以提前“缓存”之前k\v的值再和当前k\v拼接，那么当前就只需要计算第三个输入的q\k\v特征（直接看下一张图更直观）。

三、加入KV Cache的attention

现在每轮只需要提取上一步推理出的最后1个字符的q\k\v，而不是计算前面所有字符的；

计算出当前的字符的qkv特征，k和v拼接上之前缓存的k’和v‘，同时将当前字符的的k和v缓存；

图中仅白色的部分是需要计算的，和二中的图相比，我们可以看到计算量大幅减少。

四、扩展——多头注意力

2-head attension示例

五、关于mask

事实上在decoder端还有一个mask机制：

如图，因为softmax是按行处理，因此对前面的分析并不会有什么影响。

六、结论

kv-cache主要用于优化大模型推理过程；
kv-cache主要用于优化要自回归操作的decoder端；
kv-cache通过避免反复加载计算attention操作中不需要计算的数据或是需要多次加载的数据减少了计算负载。

本篇都是基于本人个人的理解，回答的内容可能理解有偏差甚至是本质的错误还望各位大佬指出！