一文读懂KVCache

本文转载自知乎：https://zhuanlan.zhihu.com/p/...

在LLM的推理过程中，KVCache已经属于必备的技术了。然而，网上现有的解读文章并不够清晰。看了一些文章之后，我决定还是自己写一篇新的，以飨读者。

Motivation: 为什么要KVCache

为了回答KVCache的必要性，我们首先需要对transformer类大语言模型的计算过程有个整体的理解。

我们可以用简单的代码来验证这一结论：

import torch  
from transformers import GPT2Tokenizer, GPT2Model  
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')  
model = GPT2Model.from_pretrained('gpt2')  
  
# text: "The quick brown fox jumps over the lazy"  
tokens = [[464, 2068, 7586, 21831, 18045, 625, 262, 16931]]  
input_n = torch.tensor(tokens)  
output_n = model(input_ids=input_n, output_hidden_states=True)  
  
# text: " dog"  
tokens[0].append(3290)  
input_n_plus_1 = torch.tensor(tokens)  
output_n_plus_1 = model(input_ids=input_n_plus_1, output_hidden_states=True)  
  
for i, (hidden_n, hidden_n_plus_1) in enumerate(zip(output_n.hidden_states, output_n_plus_1.hidden_states)):  
    print(f"layer {i}, max difference {(hidden_n - hidden_n_plus_1[:, :-1, :]).abs().max().item()}")  
    assert torch.allclose(hidden_n, hidden_n_plus_1[:, :-1, :], atol=1e-4)

输出结果：

layer 0, max difference 0.0  
layer 1, max difference 3.814697265625e-06  
layer 2, max difference 5.7220458984375e-06  
layer 3, max difference 5.7220458984375e-06  
layer 4, max difference 1.52587890625e-05  
layer 5, max difference 9.5367431640625e-06  
layer 6, max difference 7.62939453125e-06  
layer 7, max difference 1.1444091796875e-05  
layer 8, max difference 1.9073486328125e-05  
layer 9, max difference 1.52587890625e-05  
layer 10, max difference 2.288818359375e-05  
layer 11, max difference 3.0517578125e-05  
layer 12, max difference 3.0517578125e-05

更具体的计算流程如下图所示：

我们使用不同的颜色来区分初始输入（黑色）、生成第一个token时计算得出的中间结果（红色）、后续生成每个token需要计算并保存的中间结果（蓝色）。

严格来说，本文中的表示输入prompt的token数目，随着generation进行，我们还需要引入另一个变量来表示generate的token个数。为了避免引入太多符号，下文中我们用表示prompt的token数目加上已经generate的token个数，也就是当前拥有的token数。这样我们就能聚焦在next token prediction本身了。

小插曲：神经网络计算的确定性

有些强迫症患者可能看到上述GPT2的代码输出非常不爽，毕竟还有大约的误差。这里建议强迫症患者深呼吸一下，因为接下来的代码可能让他们更加震惊：

x = torch.randn(9990).abs()  
z1 = x.sum().item()  
z2 = x[:5000].sum().item() + x[5000:].sum().item()  
print((z1 - z2 == 0)) # False

一个数组直接求和，与分段求和再整体求和的结果不一样（当然，相差不多）。这是因为，加法在数学上具有结合律，然而浮点数的加法并没有结合律。

更有甚至，同一段代码执行多次的结果也可能不一样：

x = torch.randn(9990).abs()  
indices = torch.randint(0, 9990, (9000,))  
  
def f(x, indices):  
    y = x.clone()  
    y[indices] += x[:9000]  
    return y  
  
y1 = f(x, indices)  
y2 = f(x, indices)  
print((y1 - y2).abs().max().item() == 0) # False

这是因为indices里面有重复的元素，所以y[indices] += x[:9000]涉及并行竞争。而并行的顺序一般是不确定的，所以这段代码的执行本身具有随机性。

一般来说，神经网络的计算并不会要求特别强的确定性，甚至不会要求特别高的精度，这也正是这几年算力突飞猛进的原因（使用低精度计算）。

KVCache的原理及设计细节

考虑到一层里面有个cache，三种方式的总FLOPS的主导项都是，在计算量上面是差不多的。因此我们选择存储占用最少的方案，也就是常用的KVCache（一般远大于）。如果哪一天某个奇怪的模型，或者存在非常明显的KVCache复用（本次的KVCache还会在下一轮对话中用到，比如针对长文本的多次summarization），那可能可以考虑一下s cache。

至此，我们不仅了解了KVCache的原理，而且还分析了它为什么设计成这样而不是别的样子。

KVCache的存储及实现细节

KVCache的总大小为，正比于当前token数量、向量维度、层数。这里面，最令人头疼的是，它是在推理过程中不断变大的一个量。变长数据的存储总是很烦人的，具体解决起来无外乎三种方法：

分配一个最大容量的缓冲区，要求提前预知最大的token数量。而且现在大模型卷得厉害，动不动支持上百万长度，而大部分的用户请求都很短。因此，按照最大容量来分配是非常浪费的。
动态分配缓冲区大小，类似经典的vector append的处理方式，超过容量了就扩增一倍。这也是一种可行的解决方案，但是（在GPU设备上）频繁申请、释放内存的开销很大，效率不高。
把数据拆散，按最小单元存储，用一份元数据记录每一块数据的位置。

最后一种方案，就是目前采用最多的方案，也叫PageAttention。程序在初始化时申请一整块显存（例如4GB），按照KVCache的大小划分成一个一个的小块，并记录每个token在推理时要用到第几个小块。小块显存的申请、释放、管理，类似操作系统对物理内存的虚拟化过程，这就是大名鼎鼎的vLLM的思路（具体参见论文Efficient Memory Management for Large Language Model Serving with PagedAttention）。

KVCache成立条件

总结

本文分析了LLM推理中常用的KVCache技术的原理、设计细节以及成立条件，顺带提了神经网络计算（GPU计算）中的误差与不可避免的随机性。感谢@朱力耕在本文撰写过程中提供了很多建议。

作者：游凯超
来源：机智流

推荐阅读

欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

Motivation: 为什么要KVCache

小插曲：神经网络计算的确定性

KVCache的原理及设计细节

KVCache的存储及实现细节

KVCache成立条件

总结

推荐阅读

目录