什么是 KV Cache#

大模型优化中一个常用的技术就是 kv cache,可以在不影响任何计算精度的前提下通过空间换时间,提升推理性能.

背景#

生成式模型的推理过程很有特点,我们给定一个输入文本,模型会输出一个回答 (假设长度是 n),那么实际上模型进行了 n 次推理,也就是 llm 每次推理给出一个 token,输出 token 会与输入 token 拼在一起,作为下一次推理的输入,这样推理直到遇到终止符.

1
import torch
2
from transformers import GPT2LMHeadModel, GPT2Tokenizer
3

4
model = GPT2LMHeadModel.from_pretrained("/WORK/Test/gpt", torchscript=True).eval()
5

6
# tokenizer
7
tokenizer = GPT2Tokenizer.from_pretrained("/WORK/Test/gpt")
8
in_text = "Lionel Messi is a"
9
in_tokens = torch.tensor(tokenizer.encode(in_text))
10

11
# inference
12
token_eos = torch.tensor([198]) # line break symbol
13
out_token = None
14
i = 0
15
with torch.no_grad():
16
    while out_token != token_eos:
17
        logits, _ = model(in_tokens)
18
        out_token = torch.argmax(logits[-1, :], dim=0, keepdim=True)
19
        in_tokens = torch.cat((in_tokens, out_token), 0)
20
        text = tokenizer.decode(in_tokens)
21
        print(f'step {i} input: {text}', flush=True)
22
        i += 1
23

24
out_text = tokenizer.decode(in_tokens)
25
print(f' Input: {in_text}')
26
print(f'Output: {out_text}')

输出为:

1
step 0 input: Lionel Messi is a player
2
step 1 input: Lionel Messi is a player who
3
step 2 input: Lionel Messi is a player who has
4
step 3 input: Lionel Messi is a player who has been
5
step 4 input: Lionel Messi is a player who has been a
6
step 5 input: Lionel Messi is a player who has been a key
7
step 6 input: Lionel Messi is a player who has been a key part
8
step 7 input: Lionel Messi is a player who has been a key part of
9
step 8 input: Lionel Messi is a player who has been a key part of the
10
step 9 input: Lionel Messi is a player who has been a key part of the team
11
step 10 input: Lionel Messi is a player who has been a key part of the team's
12
step 11 input: Lionel Messi is a player who has been a key part of the team's success
13
step 12 input: Lionel Messi is a player who has been a key part of the team's success.
14
step 13 input: Lionel Messi is a player who has been a key part of the team's success.
15

16
 Input: Lionel Messi is a
17
Output: Lionel Messi is a player who has been a key part of the team's success.

注意这里的输出,每次推理后 token 都变长了,有没有办法让推理过程中的 flops 基本恒定不变呢?

原理#

在上述的推理中,每个 step 中输入一个 token 序列,经过 embedding 层把 token 变为三维张量[b,s,h],经过一通计算，最后经 logits 层将计算结果映射至词表空间，输出张量维度为[b, s, vocab_size]。

当前轮输出 token 与输入 tokens 拼接，并作为下一轮的输入 tokens，反复多次。可以看出第 i+1 轮输入数据只比第 i 轮输入数据新增了一个 token，其他全部相同！因此第 i+1 轮推理时必然包含了第 i 轮的部分计算。KV Cache 的出发点就在这里，缓存当前轮可重复利用的计算结果，下一轮计算时直接读取缓存结果，就是这么简单，不存在什么 Cache miss 问题。