一、為什么要研究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 generationsstart = 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架構