LLM推理入门指南②：深入解析KV缓存

OneFlow深度学习框架

已于 2024-04-01 15:10:15 修改

阅读量5.1k

点赞数 15

分类专栏：前沿技术文章标签： KV缓存大模型 LLM LLM推理

于 2024-03-28 08:03:52 首次发布

本文链接：https://blog.youkuaiyun.com/OneFlow_Official/article/details/137130127

版权

本文围绕LLM推理展开，指出注意力层计算成本与总序列长度呈二次方扩展问题。介绍了KV缓存这一常用优化方式，它能使注意力计算线性扩展，以内存换计算。同时探讨了KV缓存带来的内存压力等新问题，以及相关创新解决方案。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在本系列文章《LLM推理入门指南①：文本生成的初始化与解码阶段》中，作者对Transformer解码器的文本生成算法进行了高层次概述，着重介绍了两个阶段：单步初始化阶段，即提示的处理阶段，和逐个生成补全词元的多步生成阶段。

本文进一步探讨了LLM推理的挑战 —— 第一大挑战是，注意力层（也称为自注意力层）与总序列长度（提示词元和生成补全词元）的计算成本呈二次方扩展的问题。幸运的是，生成步骤之间的许多计算是冗余的，从而能够缓存适当的结果并减少计算需求。

这种缓存方案也就是KV缓存，是LLM推理过程中的一种常用优化方式，使得（自）注意力机制的计算需求在总序列长度（提示 + 生成的完成部分）上线性扩展，而不是呈二次方扩展。更具体地说，KV缓存通过在生成过程中计算出的键和值张量存储（“缓存”）在GPU内存中，从而在每个生成步骤中节省了对过去词元键和值张量的重新计算。

KV缓存是一种折衷方案：以内存换取计算。本文后半部分将介绍KV缓存可以增长到多大，会带来哪些挑战，并介绍解决这些挑战的最常见应对策略。

（本文作者为AWS的GenAI解决方案架构师Pierre Lienhart。以下内容由OneFlow编译发布，转载请联系授权。原文：https://medium.com/@plienhar/llm-inference-series-3-kv-caching-unveiled-048152e461c8）

作者 | Pierre Lienhart

OneFlow编译

翻译｜宛子琳、杨婷

1 简要回顾Transformer注意力层

首先，我们来回顾一下vanilla Transformer（图1）中多头注意力（MHA）层的一些事实。

图1：Transformer解码器层的详细视图（上）和一个输入序列长度为3的双头（自）注意力层（下）

为简单起见，假设我们只处理长度为t的单个序列（即批处理大小为1）：

在整个过程中，输入序列（提示）中的每个词元都由一个密集向量表示（图1的浅黄色部分）。
注意力层的输入是一个密集向量序列，每个输入词元都对应一个向量，这些向量由前一个解码器块生成。
对于每个输入向量，注意力层生成一个具有相同维度的单个密集向量（图1的浅蓝色部分）。

现在，考虑单个注意力头的情况：

首先，我们通过三种不同的投影（图1最左侧的浅灰色向量）为每个输入向量生成三个较低维度的密集向量：查询（query）、键（key）和值（value）。共有t个查询向量、t个键向量和t个值向量。
对于每个查询，我们生成一个输出向量，该向量等于值的线性组合，该线性组合的系数即为注意力分数。换句话说，对于每个查询，相应的输出向量是值的注意力加权平均值。对于给定查询，注意力分数是由该查询与每个键的点积得出的。这样我们就为序列中的每个词元生成了一个表示，其中包含来自其他词元的信息，意味着我们创建了每个词元的上下文表示。
然而，在自回归解码的上下文中，我们无法使用所有可能的值来构建给定查询的输出表示。实际上，在计算与特定词元相关的查询的输出时，我们不能使用序列中后续出现的词元的值向量。这一限制通过掩码（masking）来实现，将被禁止的值向量的注意力分数设置为零，即被禁止的词元的注意力分数设置为零。

最后，将每个注意力头的输出连接，并使用最后一个线性变换进行转换，以产生最终输出。

注意力计算的二次方扩展

我们看一下计算注意力分数所需的浮点运算数（FLOPs）。对于一个给定的注意力头，对于批处理大小为b，总长度为t的每个序列（包括提示和生成的补全部分），注意力分数矩阵是通过将形状为(t, d_head)的查询张量与形状为(d_head, t)的转置键张量相乘而创建的。

一个矩阵乘法所需的FLOPs是多少？将形状为(n, p)的矩阵与另一个形状为(n, m)的矩阵相乘大致需要 2.m.n.p 次操作。在这种情况下，单头单序列的注意力分数计算大约需要 2.d_head.t^2 FLOPs。因此，总体上注意力分数计算需要 2.b.n_layers.n_head.d_head.t^2=2.b.n_layers.d_model.t^2 FLOPs。现在，t 的二次方扩展变得清晰可见。

从具体数字来看，以Meta的Llama2–7B为例，n_layers=32 且 d_model=4096。

注意：使用值张量乘以掩码注意力分数矩阵所需的FLOPs与上述计算量相同。

那么涉及模型权重的矩阵乘法呢？通过类似的分析，可以证明其计算复杂度为 O(b.n_layers.d_model^2.t)，即计算需求与总序列长度 t 呈线性增长。

下面我们通过一个例子来理解二次方扩展的严重性。如果要生成第1001个词元，模型必须执行比生成第101个词元多100倍的FLOPs。这种计算量的指数级增长显然很快会变得难以承受。幸运的是，由于掩码的使用，实际上在生成步骤之间可以节省大量计算。

3 掩码在生成阶段引发了冗余计算

现在来到问题的关键。由于掩码技术的使用，对于给定词元，输出表示仅使用先前词元的表示生成。由于先前的词元在迭代过程中保持不变，因此对于该特定词元的输出表示对于所有后续迭代也将是相同的，因此意味着存在冗余计算。

我们以上一篇博文中的词元序列为例（其特性为每个单词即为一个词元）。假设我们刚刚从输入序列“ What color is the sky? The sky ”中生成了“ is ”。在过去的迭代中，“sky ”是输入序列的最后一个词元，因此与该词元相关联的输出表示是使用序列中所有词元的表示生成的，即“ What”、“ color”、“ is”、“ the”、“ sky”、“?”、“ The ”和“sky ”的值向量。

下一次迭代的输入序列将是“ What color is the sky? The sky is ”，但由于使用了掩码，从“sky”的视角来看，它似乎仍然是“ What color is the sky? The sky ”的输入序列。因此，“sky”生成的输出表示将与上一次迭代相同。

现在看一个图示（见图2），使用图1中的图表。启动步骤应该处理长度为1的输入序列。冗余计算的部分在浅红色和浅紫色中突出显示。浅紫色元素对应冗余计算的键和值。

图2 — 生成阶段注意力层中的冗余计算

回到我们的例子，对于使用“ What color is the sky? The sky is ”作为输入的新迭代，我们在先前生成步骤中尚未计算的唯一表示是输入序列中的最后一个词元“ is ”。

更具体地说，我们需要什么来完成这一过程？

“is”的查询向量。
“What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”、“sky ”和“ is ”的键向量，用于计算注意力分数。
“What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”、“sky ”和“ is ”的值向量，用于计算输出。

关于键和值向量，它们已经在先前的迭代中为所有词元计算过，除了“ is ”。因此，我们可以保存（即缓存）并重复使用先前迭代中的键和值向量。这种优化简称为KV缓存。之后，计算“ is ”的输出表示就会变得非常简单：

计算“is”的查询、键和值。
从缓存中获取“ What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”和“sky ”的键和值向量，并将它们与我们刚刚为“is”计算的键和值向量连接起来。
使用“is”的查询和所有键来计算注意力分数。
使用注意力分数和所有值来计算“is”的输出向量。

从我们的输入来看，只要可以使用它们的键和值向量，实际上就不再需要先前的词元。当我们进行KV缓存时，模型的实际输入是最后生成的词元（而不是整