transformer之KV Cache

一、为什么要研究KV Cache

非常有效的加速推理速度，效果如下所示：

import numpy as np
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
NAME_OR_PATH = r'***************'
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(NAME_OR_PATH)
model = AutoModelForCausalLM.from_pretrained(NAME_OR_PATH).to(device)
model.config.pad_token_id = tokenizer.eos_token_id
for use_cache in (True, False):
  times = []
  for _ in range(10):  # measuring 10 generations
    start = time.time()
    model.generate(**tokenizer("What is KV caching?", return_tensors="pt").to(device), use_cache=use_cache, max_new_tokens=10)
    times.append(time.time() - start)
  print(f"{'with' if use_cache else 'without'} KV caching: {round(np.mean(times), 3)} +- {round(np.std(times), 3)} seconds")
  #===================================max_new_tokens=1=======================================================
with KV caching: 0.072 +- 0.008 seconds
without KV caching: 0.081 +- 0.02 seconds
#===================================max_new_tokens=10=======================================================
with KV caching: 0.428 +- 0.021 seconds
without KV caching: 0.751 +- 0.04 seconds
#===================================max_new_tokens=100=======================================================
with KV caching: 4.606 +- 0.072 seconds
without KV caching: 19.257 +- 1.663 seconds
#===================================max_new_tokens=1000=======================================================
with KV caching: 42.941 +- 0.711 seconds
without KV caching: 349.954 +- 1.523 seconds

使用KV Cache的推理速度是明显优于没有使用KV Cache的，而且生成的token越长速度提升就越明显，当最大生成token数为1000时，近10倍的加速，一次推理近6分钟。

二、为什么KV Cache能加速

2.1 原理是什么

最本质的原理是避免重复计算，将需要重复计算的结果进行缓存，需要缓存的值为历史token对应的KV值，所以叫KV Cache。

2.2 为什么只需要KV

预测新的token只与输入的最后一个token相关，输入的最后一个token因为只需要计算注意力值，而注意力的值需要将输入token的V值进行加权即得到结果，进行加权就需要将当前的Q与与所有的K进行计算得到权重，所以只需要缓存历史token的KV值。

2.2 为什么会存在重复计算

首先，生成式模型每生成一个新token都需要调用整个模型进行一次推理，历史token计算得到的中间激活值在Decoder架构的模型中每次推理时都是一样的，所以可以进行缓存。
这是因为Decoder架构中，当前token只用之前的token计算得到注意力值，通过Causal Mask实现，换句话说，在推理的时候前面已经生成的字符不需要与后面的字符产生attention，从而使得前面已经计算的K和V可以缓存起来。

2.3 预测新token的计算步骤

由于Causal Mask矩阵的存在，预测下一个token只与输入的最后一个token的QKV和历史token的KV有关；
如果没有Causal Mask，比如说是encoder架构，每次推理时每个token需要考虑所有输入的token，所以得到的注意力值都会变化，就不存在重复计算的情况。

三、结论

1、KV Cache是通过空间换时间，避免重复计算进而提升推理速度
2、预测新token只与输入的最后一个token的注意力值相关，而注意力值与最后一个token的Q和所有输入token的KV相关，每一层的注意力不变，进而每一层的KV都是不变的
3、只适用于Decoder架构，因为只与之前的token进行计算，得到的注意力值不会变化，第一层、第二层、第三层到第$l$层； 如果只有一层，那就不仅仅适用于Decoder架构