KV Cache in LLMs 阅读：https://outcomeschool.com/blog/kv-cache-in-llms…

大语言模型中的 KV Cache

来源：https://outcomeschool.com/blog/kv-cache-in-llms 大语言模型中的 KV Cache

在这篇博客中，我们将了解 KV Cache（其中 K 代表 Key，V 代表 Value），以及为什么在大语言模型（LLMs）中使用它来加速文本生成。

我们将首先了解 LLM 如何逐个 Token 生成文本，理解模型内部 Key、Value 和 Query 的作用，通过一个示例看看重复计算的问题，然后逐步了解 KV Cache 如何通过存储和重用过去的结果来解决这个问题。

我是 Amit Shekhar，Outcome School (https://outcomeschool.com/) 的创始人。我曾教导和辅导过许多开发者，他们的努力使他们获得了高薪的科技工作，帮助许多科技公司解决了他们独特的问题，并创建了许多被顶级公司使用的开源库。我热衷于通过开源、博客和视频分享知识。

我在 Outcome School 教授人工智能和机器学习 (https://outcomeschool.com/program/ai-and-machine-learning)。

让我们开始吧。

LLM 如何生成文本

在了解 KV Cache 之前，我们首先需要理解 LLM 如何生成文本。

LLM（大语言模型）是一种在海量文本数据上训练过的模型。它能够理解和生成人类语言。当我们给它一个句子时，它会预测下一个是什么。

LLM 生成文本是一次一个 Token 的。Token 是文本的一小部分——它可以是一个词、词的一部分，甚至是一个单个字符。为了简单起见，让我们把每个词当作一个 Token。

假设我们给模型这个输入：

模型查看 “I” 和 “love”，并预测下一个 Token：“teaching”。

现在完整的序列变为：

模型查看 “I”、“love” 和 “teaching”，并预测下一个 Token：“AI”。

现在完整的序列变为：

这个过程一次一个 Token 地继续，直到模型决定停止。

这里值得注意的是：每次模型预测新 Token 时，它都需要查看所有之前的 Token 来决定接下来是什么。

模型内部发生了什么

现在，让我们了解模型在生成每个 Token 时内部发生了什么。

模型有一个叫做 attention layer（注意力层） 的组件，这是 Transformer 架构 (https://outcomeschool.com/blog/decoding-transformer-architecture) 的核心。注意力层的工作是帮助模型确定哪些之前的 Token 对于预测下一个 Token 是重要的。并非所有之前的 Token 都同样有用。有些比其他更重要。

在注意力层内部，每个 Token 都被转换为三个东西。让我们通过一个简单的类比来理解它们。

想象一个教室，有一个新学生加入，想知道谁能帮助某个特定主题：

• Query (Q)：新学生的问题——“这里谁懂这个主题？”每个被预测的 Token 都有一个 Query。它表示该 Token 在寻找什么。
• Key (K)：每个现有学生佩戴的名牌，上面写着他们知道什么——“我懂数学”或“我懂科学”。每个之前的 Token 都有一个 Key。它描述该 Token 包含什么信息。
• Value (V)：每个学生实际拥有的笔记。一旦新学生通过名牌（Key）找到合适的人，笔记（Value）就是他们实际使用的东西。每个之前的 Token 都有一个 Value。它承载实际的信息。

所以，当前 Token 使用其 Query 与所有之前 Token 的 Keys 进行比较。这种比较产生 attention scores（注意力分数）——数字告诉模型对每个之前 Token 给予多少关注。较高的分数意味着该 Token 更相关。较低的分数意味着相关性较低。然后模型使用这些分数从相关 Token 中收集 Values。如果我们想更深入地了解 Q、K 和 V 是如何计算的，我们有一篇关于注意力背后的数学：Q、K、V (https://outcomeschool.com/blog/math-behind-attention-qkv) 的详细博客。

这就是注意力层的工作方式。

现在，让我们看看问题。

要想动手学习注意力层、Q/K/V 内部机制和 Transformer 架构，并参与真实项目，请查看 Outcome School 的人工智能和机器学习课程 (https://outcomeschool.com/program/ai-and-machine-learning)。

问题：重复计算

每次模型预测下一个 Token 时，它都会为序列中的所有 Token 计算 Key、Value 和 Query——而不仅仅是新的 Token。

让我们通过我们的示例逐步看看发生了什么：

步骤 1： 输入是 “I love”

模型计算以下内容的 Key、Value 和 Query：

• “I”
• “love”

它使用这些来预测下一个 Token：“teaching”。

步骤 2： 输入是 “I love teaching”

模型计算以下内容的 Key、Value 和 Query：

• “I”（已在步骤 1 中计算，但再次计算）
• “love”（已在步骤 1 中计算，但再次计算）
• “teaching”（新）

它使用这些来预测下一个 Token：“AI”。

步骤 3： 输入是 “I love teaching AI”

模型计算以下内容的 Key、Value 和 Query：

• “I”（已在步骤 1 和步骤 2 中计算，但再次计算）
• “love”（已在步骤 1 和步骤 2 中计算，但再次计算）
• “teaching”（已在步骤 2 中计算，但再次计算）
• “AI”（新）

你看到问题了吗？

“I” 的 Key 和 Value 在步骤 1 中已经计算。但模型在步骤 2 中再次计算，在步骤 3 中又再次计算。“love” 和 “teaching” 也是如此。

模型正在为已经见过的 Token 重复相同的工作。 这是浪费的计算。

随着序列变长，这个问题会变得更糟。如果模型已经生成了 100 个 Token，那么在下一步，它将重新计算所有 100 个之前 Token 的 Key 和 Value，只是为了预测一个新 Token。这使得文本生成非常慢。

解决方案：KV Cache

KV Cache 背后的想法很简单：为每个 Token 计算一次 Key 和 Value，存储它们，并在每一步中重用它们。

想象一下做笔记。假设你正在开会。每次有新的人发言时，你不是要求所有之前的发言人重复他们说的话，而是简单地从你的笔记中读取，只听新发言人的话。笔记就是你的缓存。

类似地，KV Cache 是一个内存，我们保存每个已处理的 Token 的 Key 和 Value。这样，模型就不需要再次计算它们。

让我们看看带有 KV Cache 的相同示例如何工作：

步骤 1： 输入是 “I love”

模型计算以下内容的 Key、Value 和 Query：

• “I”
• “love”

它将 “I” 和 “love” 的 Key 和 Value 保存在 KV Cache 中。

它使用这些来预测下一个 Token：“teaching”。

KV Cache 现在包含： “I” 和 “love” 的 Key 和 Value

步骤 2： 输入仅为 “teaching”（只有新的 Token）

模型从 KV Cache 中检索 “I” 和 “love” 的 Key 和 Value。无需重新计算。

它仅为新 Token 计算 Key、Value 和 Query：“teaching”。

它将 “teaching” 的 Key 和 Value 保存在 KV Cache 中。

它使用所有这些一起预测下一个 Token：“AI”。

KV Cache 现在包含： “I”、“love”、“teaching” 的 Key 和 Value

步骤 3： 输入仅为 “AI”（只有新的 Token）

模型从 KV Cache 中检索 “I”、“love” 和 “teaching” 的 Key 和 Value。无需重新计算。

它仅为新 Token 计算 Key、Value 和 Query：“AI”。

它将 “AI” 的 Key 和 Value 保存在 KV Cache 中。

KV Cache 现在包含： “I”、“love”、“teaching”、“AI” 的 Key 和 Value

所以，模型不是在每一步都为每个 Token 重新计算 Key 和 Value，而是仅为新 Token 计算，并重用所有之前 Token 的缓存值。

这就是 KV Cache 避免重复计算的方式。

为什么只缓存 Key 和 Value，而不缓存 Query

你自然会问：为什么我们只缓存 Key 和 Value，而不缓存 Query？

Query 仅用于当前 Token——即当前正在生成的那个。当前 Token 使用其 Query 与所有之前 Token 的 Keys 进行比较，以找出哪些是相关的。一旦预测完成，该 Query 就不再需要。

但每个过去 Token 的 Key 和 Value 在每个未来步骤中都需要，因为每个新 Token 必须查看所有之前的 Token 来进行预测。

所以，我们只需要存储 Keys 和 Values。这就是为什么它被称为 KV Cache——它缓存 Keys 和 Values。

速度能快多少

让我们并排比较这两种方法：

没有 KV Cache：

步骤 1：为 2 个 Token 计算 K、V、Q 步骤 2：为 3 个 Token 计算 K、V、Q（2 个重新计算） 步骤 3：为 4 个 Token 计算 K、V、Q（3 个重新计算） 步骤 4：为 5 个 Token 计算 K、V、Q（4 个重新计算） ... 步骤 N：为 (N+1) 个 Token 计算 K、V、Q（N 个重新计算）

计算量在每一步都不断增加。

使用 KV Cache：

步骤 1：为 2 个 Token 计算 K、V、Q → 将 2 个 Token 的 K、V 保存在缓存中 步骤 2：为 1 个 Token 计算 K、V、Q → 从缓存中重用 2 个 Token 的 K、V 步骤 3：为 1 个 Token 计算 K、V、Q → 从缓存中重用 3 个 Token 的 K、V 步骤 4：为 1 个 Token 计算 K、V、Q → 从缓存中重用 4 个 Token 的 K、V ... 步骤 N：为 1 个 Token 计算 K、V、Q → 从缓存中重用 N 个 Token 的 K、V

在第一步之后，模型在每一步中仅为一个新 Token 计算，而不是整个序列。

为了直观理解这一点：如果模型正在生成一个包含 100 个 Token 的序列，没有 KV Cache，所有步骤中 Key 和 Value 的总计算次数将是 2 + 3 + 4 + … + 100 = 5,049 次计算。使用 KV Cache，它将是 2 + 1 + 1 + … + 1 = 101 次计算。这大约是减少了 50 倍的计算量。序列越长，节省越多。

这就是 KV Cache 在加速文本生成方面如此有效的原因。

权衡：速度与内存

KV Cache 使生成更快，但它有一个权衡：它使用额外的内存来存储到目前为止生成的每个 Token 的所有 Key 和 Value 信息。

随着序列变长，缓存会变大。对于包含数千个 Token 的非常长的序列，缓存会消耗大量内存。

所以，KV Cache 是一个权衡：我们使用更多内存来节省计算时间。对于大多数用例来说，这种权衡是非常值得的，因为速度提升显著。

现在，我们已经理解了 LLM 中的 KV Cache。

在下一篇博客中，我们将学习 Paged Attention (https://outcomeschool.com/blog/paged-attention-in-llms)，它解决了 KV Cache 的内存问题。

为 AI 工程面试做好准备：AI Engineering Interview Questions (https://github.com/amitshekhariitbhu/ai-engineering-interview-questions)

目前就到这里。

谢谢

Amit Shekhar Outcome School (https://outcomeschool.com/) 创始人

你可以在以下平台联系我：

• X (https://x.com/amitiitbhu)
• LinkedIn (https://www.linkedin.com/in/amit-shekhar-iitbhu)
• YouTube (https://www.youtube.com/@amitshekhar)
• GitHub (https://github.com/amitshekhariitbhu)

关注 Outcome School：

• X (https://x.com/outcome_school)
• LinkedIn (https://www.linkedin.com/company/outcomeschool)
• YouTube (https://youtube.com/@OutcomeSchool)
• GitHub (http://github.com/OutcomeSchool)

在这里阅读我们所有高质量的博客。 (https://outcomeschool.com/blog)