【深度解析】为什么有KV Cache而没有Q Cache？揭秘缓存机制的背后！

offer捷报

KV Cache 在已经成为了 LLM 推理系统中必须要用到的一个优化技术。

KV Cache 的想法很简单，即通过把中间结果缓存下来，避免重复计算，但是 KV Cache 有时候又是比较令人困惑的。

学习 KV Cache，搞清楚下面的三个问题就算是完全弄清楚了：

• 缓存的是什么？

• 为什么是 KV Cache，而不是 QKV Cache？

• KV Cache 对模型的输出有任何的影响吗？

如果我们从矩阵乘法的角度观察 attention 运算的过程，那么上面的三个问题很容易理解。

01

自回归模型的特点

LLM 的推理被称为是自回归的过程，也就是说模型上一步的输出会被当作是下一步的输出。

如下面的动图所示，用户输入的 prompt 是"recite the first law $"。

模型产生的第一个 token 是"A"，然后输出的"A"会添加到用户的 prompt 后面，再次输入到模型中，模型输出"robot"这个 token。如此反复，直到模型输出结束词。

我们把运行一次模型成为一个 step，每一个 step 生成一个新的 token。

02

KV Cache 的探索过程

下面，我们从最朴素的推理方法出发，一步一步推理出 KV Cache 的优化方法。假设，我们的模型有两个 attention 层。

不考虑任何 Cache 的推理：

Step 0

我们先不考虑任何 Cache，观察我们的推理情况。 如图 1 所示，我们以“你好”作为 prompt 输入到模型中。

图1

模型的 embedding 层会将 token 转化为对应的 embedding(由于篇幅原因，embedding 层就没有展示在图 1 中) ，得到一个 2x4 的矩阵 embedding1，其中第一行是“你”的 embedding 表示，第二行是“好”的 embedding 表示。

之后，我们进入到第一个 attention 层。embedding1 会与 attention 层中的三个权重矩阵相乘，分别得到 Q、K 和 V 矩阵，如图 2 所示。

图2

在 attention layer 中，我们根据 attention 层的计算公式对 Q、K 和 V 矩阵进行计算。

首先，Q 和 K 矩阵相乘，得到一个 attention 矩阵，其中包括了每两个词之间的注意力值。比如，第二行第一列表示“好”这个 token 对“你”这个 token 的注意力值。

需要额外注意的是，“你”这个 token 对“好”这个 token 的注意力值被 mask 了，即每一个 token 只计算与之前的 token 的注意力值，而不计算之后的 token 的注意力值。

然后，attention 矩阵与 V 矩阵相乘，再把相乘结果输入到 FFN 层中(这里的 FFN 对我们要讨论的 KV Cache 没有什么影响，所以可以忽略)，就得到了新的 embedding2。

第二个 attenion 层也是如此类推，并输出 embedding3。

最后，embedding3 的最后一行，也就是“好”这个 token 对应的 embedding 被输入到一个分类其中，预测下一个 token。得到的预测结果是“啊“。到此，我们就生成第一个 token，即完成了 step0。

Step 1

在预测了 token“啊”之后，我们将“你好啊”再次输入到模型中，如图 3 所示。

图3

除了我们的 token 数量从 2 变成了 3 之外，需要执行的运算与 Step 0 完全相同，输入的 embedding 经过两层 attention 层之后，得到 embedding3。embedding3 的最后一行被输入到分类器中，预测得到 token“!”。

一个细节是，Step 1 和 Step 0 的红色的部分是相同的。比如说，Step 1 的 K1 矩阵前两行与 Step 0 的 K1 矩阵的前两行是相同的。

这是因为 K1 矩阵由 embedding1 矩阵和模型的权重矩阵相乘得到，而两个 Step 的权重矩阵是相同的，唯一不同的地方在于 Step 1 的 embedding1 比 Step 0 多了最后一行，所以两个 Step 的 QKV 矩阵的前两行是相同的，如下图所示。

图4

删除无效计算

观察图 1 和图 3 可以发现，模型之后的 Classifier 只会使用 embedding3 矩阵的最后一行作为输入，并输出预测结果。

由此我们可以发现，实际上，每一个 embedding 矩阵只有最新输入的 token 的行是有效的。

比如，在图 3 的 embedding3 中，只有 step 0 产生的新 token“啊”对应的第三行是有用的，前两行都是无效的。因此，我们可以删除这些无用的行，以节约计算和存储开销。

图 5 以 Step1 的第一个 attention layer 作为演示。我们首先删除三个 embedding 矩阵的红色的部分。

因为 embedding2 和 embedding3 矩阵是通过 attention 和 V1 矩阵相乘得到的，所以我们也需要删除 attention 矩阵的重复部分。

而 attention 矩阵是由 Q 矩阵和 K 矩阵相乘得到，所以我们还需要删除 Q 矩阵的重复部分。

图5

所以，通过删除无效计算，我们的模型推理过程可以用图 6 表示。

图6

缓存重复计算

经过仔细观察我们可以发现，在图 6 中，Step 1 和 Step 2 的红色部分是完全相同的，只有白色的部分是不相同的。

比如说，Step 1 的 Q1 矩阵、K1 矩阵和 V1 矩阵和 Step 0 的 Q1 矩阵、K1 矩阵和 V1 矩阵的前两行是完全相同的。

那么我们其实完全可以在 Step 0 的时候把这些矩阵 Cache 下来，然后在 Step 1 的时候 load 进来，从而节省大量的计算。

下面是进行缓存之后的计算，如图 7 所示。

图7

在 Step 0 中，我们把红色部分的数据 Cache 下来，然后在 Step 1 的 load 进行计算。

在图 7 的 Step 1 中，红色部分的数据都是历史缓存，只有白色部分是新产生的数据。

我们将缓存数据和新产生的数据拼接起来，即可得到一个新的矩阵。当然，在 Step 1 中，我们也需要缓存新产生的数据，以供后面的 Step 使用。

03

总结

在 KV Cache 中，我们把第一个 Step，即 Step 0，称为是 Prefill 阶段，因为这个阶段是在“填充”KV Cache。而之后的所有 Step 称为是 Decode 阶段。

通过观察我们可以发现，Prefill 阶段是是 Compute bound，即计算密集型的，因为这个阶段我们的 embedding 矩阵都特别大，需要进行大量的矩阵乘法计算；

而 Decode 阶段是 Memory bound，即访存密集型的，因为这个阶段需要从显存乃至内存中加载 KV Cache。

大模型推理系统基本上要解决下面的问题：

如何加速 Prefill 阶段的计算，比如，现在很多工作使用预查找表加速矩阵乘法计算。

如何应对 Decode 阶段的访存开销，比如一些工作尝试使用 CSD(Computational Storage Device)缓解推理过程的 Memory bound，提高推理的吞吐率。

推理任务的调度问题。因为 LLM 的是自回归的，所以我们不知道模型究竟需要执行多久。

并且，模型的输出长度越长，KV Cache 占据的存储空间就越多。所以，如何在多个服务器实例之间调度推理任务是非常 tricky 的。之后阿里今年新出的论文 Llumnix 要解决的就是 LLM 的推理任务调度的问题。

由此，我们可以回答开始提到的三个问题：

（1）缓存的是什么？

缓存的是 K 和 V 矩阵。

（2）为什么是 KV Cache，而不是 QKV Cache？

如图 5 所示，我们只需要 embedding 的最后一行，所以我们也只需要 attention 的最后一行，因此只需要 Q 矩阵的最后一行。

所以，不缓存 Q 矩阵的原因是没必要缓存，我们只需要 Q 矩阵的最后一行即可。

（3）KV Cache 对模型的输出有任何的影响吗？

因为输入到 Classifier 的 embedding 是没有改变的，所以模型的输出也不会有任何变化。

END

AI大模型学习福利

作为一名热心肠的互联网老兵，我决定把宝贵的AI知识分享给大家。 至于能学习到多少就看你的学习毅力和能力了 。我已将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

大模型&AI产品经理如何学习

求大家的点赞和收藏，我花2万买的大模型学习资料免费共享给你们，来看看有哪些东西。

1.学习路线图

第一阶段： 从大模型系统设计入手，讲解大模型的主要方法；

第二阶段： 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用；

第三阶段： 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统；

第四阶段： 大模型知识库应用开发以LangChain框架为例，构建物流行业咨询智能问答系统；

第五阶段： 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型；

第六阶段： 以SD多模态大模型为主，搭建了文生图小程序案例；

第七阶段： 以大模型平台应用与开发为主，通过星火大模型，文心大模型等成熟大模型构建大模型行业应用。

2.视频教程

网上虽然也有很多的学习资源，但基本上都残缺不全的，这是我自己整理的大模型视频教程，上面路线图的每一个知识点，我都有配套的视频讲解。

（都打包成一块的了，不能一一展开，总共300多集）

因篇幅有限，仅展示部分资料，需要点击下方图片前往获取

3.技术文档和电子书

这里主要整理了大模型相关PDF书籍、行业报告、文档，有几百本，都是目前行业最新的。

4.LLM面试题和面经合集

这里主要整理了行业目前最新的大模型面试题和各种大厂offer面经合集。

👉学会后的收获：👈
• 基于大模型全栈工程实现（前端、后端、产品经理、设计、数据分析等），通过这门课可获得不同能力；

• 能够利用大模型解决相关实际项目需求： 大数据时代，越来越多的企业和机构需要处理海量数据，利用大模型技术可以更好地处理这些数据，提高数据分析和决策的准确性。因此，掌握大模型应用开发技能，可以让程序员更好地应对实际项目需求；

• 基于大模型和企业数据AI应用开发，实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能， 学会Fine-tuning垂直训练大模型（数据准备、数据蒸馏、大模型部署）一站式掌握；

• 能够完成时下热门大模型垂直领域模型训练能力，提高程序员的编码能力： 大模型应用开发需要掌握机器学习算法、深度学习框架等技术，这些技术的掌握可以提高程序员的编码能力和分析能力，让程序员更加熟练地编写高质量的代码。

1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集

👉获取方式：
😝有需要的小伙伴，可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓