理解KV Cache

你有没有想过为什么当和ChatGPT等大模型对话时，第一个词的生成需要等待一段时间，但后续的生成都非常顺利，一气呵成？这实际上可以归功于KV Cache的设计。

KV Cache其实是一个非常简单的设计，旨在解决大模型生成文本时的重复计算问题。

标准的推理过程

当大模型生成文本时，它需要看所有先前的tokens（在上下文长度之内的）以预测下一个token。一般地，它对每一个新token会重复相同的计算，可能会很慢。具体而言，在自注意力机制计算时，例如，对于第一个token，其输入嵌入要首先要与Wq、Wk、Wv相乘，得到q、k、v向量，然后（为简单起见，以不带掩码的注意力机制为例），每个q与其他所有的k相乘（在带掩码的注意力机制里，是每个q与自己的k和先前的k相乘，例如，q1只与k1相乘，q2与k1和k2相乘），得到注意力权重，所有的注意力权重soft经Softmax后再与v向量相乘，得到最终的上下文向量，如图所示：

可以发现，之后的每个q都要与其他的k、v（在因果注意力中是与先前的k、v）进行乘法计算。

在实际的大模型中，一般都使用因果注意力或者叫带掩码的注意力机制，例如，q3要与k1、k2、k3和v1、v2、v3相乘，得到$\omega$3，既然如此，可以先把所有的k和v向量提前缓存好，然后在后续的计算中直接使用，而无需重复计算，这样可以大幅度节约推理的时间和成本。

KV Cache是怎么工作的？

1. 第一次生成：当模型看见第一个输入时，它计算并保存它的k和v向量到cache中。
2. 下一个词：对于每一个新词，模型检索存储的k和v向量，并把当前的k和v添加到cache中，而不是从头计算每一个k、v。
3. 高效的Attention计算：对于新的Q，使用cached的K和V来计算输出。
4. 更新输入：把生成的token加入到输入中，回到第二步，直到结束生成。

可以诠释为如下过程：

1
2
3
4

Token 1: [K1, V1] ➔ Cache: [K1, V1]
Token 2: [K2, V2] ➔ Cache: [K1, K2], [V1, V2]
...
Token n: [Kn, Vn] ➔ Cache: [K1, K2, ..., Kn], [V1, V2, ..., Vn]

对比KV Cache和标准推理