【笔记】KV-cache

Transformer原理讲解参考此文

KV-cache原理

解码时这三个矩阵的大小不同，事实上，通常是一个向量，而是矩阵。向量表示新的token

在注意力机制中，我们首先对查询向量（query vector）和键矩阵（key matrix）进行点积运算，然后对结果应用 $softmax$，将其作为对值矩阵（value matrix）的加权和。在自回归解码（Auto-regressive decoding）中，模型在给定所有之前的上下文的情况下逐词生成，因此键矩阵$K$和值矩阵$V$包含整个序列的信息，而查询向量（query vector）仅包含我们刚刚看到的最后一个词的信息。你可以将查询向量与键矩阵的点积理解为在当前词与所有之前的词之间同时进行的注意力操作。因为我们是逐词生成序列，所以键矩阵和值矩阵实际上并没有太多变化，这个词对应键矩阵的一列和值矩阵的一行。关键的一点是，一旦我们计算出该词的嵌入表示，它将不会再改变，无论后面生成多少词。然而，模型仍然需要在每个后续步骤中计算该词的键$k$和值向量$v$，这导致了矩阵向量乘法的数量呈二次增长，从而导致非常缓慢的计算速度。


当模型读取一个新单词时，它会像以前一样生成查询 Vector，但缓存了键和值矩阵的先前值，因此不再需要为以前的上下文计算这些向量，相反，只需要为键 Matric 计算一个新列，为值 Matrix 计算一个新行，然后像往常一样继续使用点积和 $softmax$ 来计算注意力。

KV 缓存在 Transformer 中的作用是加速自注意力层的计算。通常，自注意力层会处理整个嵌入序列，而通过 KV 缓存，只需要传入之前的 $k$ 和 $v$ 缓存以及当前词的新嵌入。自注意力层会根据当前词的嵌入计算新的键和值向量，并将它们追加到 KV 缓存中。这样，无需重复计算之前的键和值向量，从而显著加快了计算速度，同时仍然可以生成高效的自回归解码序列。

接下来，需要将这些键和值矩阵存储在 GPU 的内存中，以便在处理下一个词时可以随时检索它们。请注意，模型中当前词与前一个词交互的唯一部分是自注意力层。在其他层（例如位置嵌入、层归一化和前馈神经网络）中，当前词与之前的上下文没有交互。因此，当使用 KV 缓存时，对于每个新词只需执行恒定量的计算工作，而且即使序列长度增加，这些计算量也不会随之增加。

计算公式

$$KV-cache=2\ast Precision\ast N_{layers}\ast D_{model}\ast Seqlen\ast Batch$$
其中，
$2$=two matrices for $K$ and $V$
$Precison$=bytes per parameter(eg:4 for fp32)
$N_{layer}$ = layer in the model
$D_{model}$ = dimension of the embeddings
$Seqlen$=length of context in tokens
$Batch$ = batch sieze

举个例子

Example :OPT-30B
$2$ = two matrices for K and V
$Precison$ = 2（fp16）
$N_{layer}$ = 48
$D_{model}$ = 7168
$Seqlen$ =1024
$Batch$ = 128
所以，KV-cache = 224871681024*128 = 180G