InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

发表时间: 2024-07 · OSDI 2024

原文: https://www.usenix.org/conference/osdi24/presentation/lee

Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim, Seoul National University

速读

一句话结论 提出 InfiniGen 动态 KV Cache 管理框架，通过在 CPU 内存保留完整缓存池，并利用前一层输入预测当前层注意力分布，仅将关键 KV 预取到 GPU，在卸载架构下实现了高达 3.00 倍的推理提速且显著优于现有压缩方法的精度。

要解决什么问题 大模型生成长文本时，KV Cache 的显存占用会随序列长度和 Batch Size 线性增长，甚至超过模型权重本身。为了解决显存不足，现代推理系统（如 FlexGen）会将 KV Cache 卸载到 CPU 内存中。但这引入了新的卡点：每生成一个 Token，都需要通过带宽有限的 PCIe 将海量 KV Cache 从 CPU 搬运到 GPU，导致严重的传输延迟。现有的缓解方案（如 $\text{H}_2\text{O}$）试图通过在运行时永久丢弃“不重要”的 Token 来缩减缓存体积。然而，这种做法忽略了注意力机制的动态特性：一个在当前迭代不重要的 Token，在后续生成中完全可能变得重要。此外，不同层对上下文的依赖差异巨大（例如浅层需要广泛的注意力，深层则高度集中），且不同 Query 需要的 Key 数量也不同，固定比例的永久驱逐策略会导致长文本外推时精度严重下降。

怎么做的 核心思路是“全量保存在 CPU，按需动态预取到 GPU”。作者观察到，由于残差连接和 LayerNorm 的存在，且输入张量中少数异常值通道占据主导地位，相邻 Transformer 层的输入张量高度相似。因此，可以在第 $i-1$ 层利用其输入来“彩排”并预测第 $i$ 层的注意力分布，从而只通过 PCIe 预取真正需要的 KV 缓存。为了让预测既快又准，InfiniGen 包含三个关键设计。首先是离线权重倾斜，利用奇异值分解（SVD）找到一个正交矩阵 $A$，在离线阶段将其乘到 Query 和 Key 的权重矩阵上：

$$ \tilde{Q} = Q \times A, \quad \tilde{K} = K \times A $$


由于 $A \times A^T = I$，这在数学上等价于原注意力计算 $\tilde{Q} \times \tilde{K}^T = Q \times K^T$。这一步将权重数值的绝对大小集中到了少数几列上。其次是预填充阶段提取局部权重，在 Prefill 阶段，基于倾斜后的矩阵，按列求和并选出绝对值最大的前 30% 列，生成用于后续解码的局部 Query 权重和局部 Key 缓存。最后是解码阶段动态预取，在 Decoding 阶段执行第 $i-1$ 层时，InfiniGen 使用第 $i-1$ 层的注意力输入、第 $i$ 层的局部 Query 权重和局部 Key 缓存，计算出一个预测的注意力得分。随后设定一个动态阈值，仅将预测得分大于 $\max - \alpha$ 的 KV 缓存条目从 CPU 预取到 GPU。因为 Softmax 的特性，减去 $\alpha$ 意味着被丢弃的 Token 权重占比极小（例如 $\alpha=5$ 时占比不到 $1/e^5$）。最后，CPU 端的缓存池通过一个低开销的计数器策略来淘汰长期未使用的条目以控制内存上限。

效果如何 实验基于 OPT（最高 30B）和 Llama-2（最高 13B）模型，在配备单张 RTX A6000 GPU 和 FlexGen 卸载推理系统的硬件上进行。对比基线包括隐式内存管理路线的 UVM、代表永久驱逐路线的 $\text{H}_2\text{O}$ 以及代表数据压缩路线的 INT4 量化。在 2048 序列长度和 Batch Size 为 20 的设置下，InfiniGen 相比 FlexGen 基线实现了最高 3.00 倍的端到端加速。在精度方面，当限制 GPU 仅加载 10% 的 KV Cache 时，InfiniGen 在 lm-evaluation-harness 的 5-shot 任务中准确率比 $\text{H}_2\text{O}$ 高出最多 32.6 个百分点。此外，随着序列长度增加（如 Llama-2-7B-32K 外推测试），INT4 和 $\text{H}_2\text{O}$ 的加速比很快触顶且困惑度显著恶化，而 InfiniGen 依然能保持与全量缓存基线几乎一致的困惑度，并提供持续增长的加速收益。其代价在于，预测机制需要额外存储局部权重和局部 Key 缓存，会占用少量额外的 GPU 显存（分别约占总参数的 2.5% 和总缓存的 15%）。

主要贡献

基于Transformer的大型语言模型（LLMs）在各种自然语言处理任务中表现出色。然而，在生成长文本时，LLM推理面临着巨大挑战，这是因为被称为键值（KV）缓存的瞬态状态会占用巨大的内存空间，且该空间随序列长度和批处理大小的增加而扩展。现代LLM服务系统支持将数据卸载到CPU内存以在硬件预算内高效服务，但将海量KV缓存从CPU内存传输到GPU成为了LLM推理中新的性能瓶颈。

为了解决这一问题，本文提出了InfiniGen，这是一个专为长文本生成量身定制的新型KV缓存管理框架，能够与现代基于卸载的推理系统协同工作。InfiniGen的创新点在于：
1. 提出了一个动态KV缓存管理框架，通过在CPU内存中智能管理KV缓存池，并结合短暂修剪（ephemeral pruning）的新型KV缓存预取技术，保留大部分KV缓存于CPU中，仅将必不可少的部分带入GPU。
2. 利用一个关键洞察：Transformer中用于计算后续注意力层的少数重要token，可以通过使用当前层的输入以及后续层的部分查询权重和键缓存执行最小化演练来进行推测。
3. 提出在离线阶段倾斜模型权重（Skewing），使得查询和键矩阵的少数列被放大，从而使注意力分数的推测更加高效和精确。
4. 在现代基于卸载的系统上实现了InfiniGen，评估表明，与现有的KV缓存管理方法相比，InfiniGen将整体性能提高了最高 $3.00\times$，同时提供了显著更好的模型精度，并在更大的模型、更长的序列和更大的批处理大小下持续提供性能提升。

背景知识与设计原则

大型语言模型的注意力机制与KV缓存：
LLM由堆叠的Transformer块组成，每个块包含一个注意力层和一个前馈层（FFN）。输入张量被层归一化后进入注意力层，与权重矩阵相乘生成查询（$Q$）、键（$K$）和值（$V$）矩阵。注意力计算公式为：$\text{softmax}(QK^T)V$。在生成式推理中，过程分为预填充（Prefill）阶段和解码（Decoding）阶段。为了避免在解码阶段重新计算所有先前token的键和值，系统通常将它们记忆在内存中，即KV缓存。KV缓存的大小随迭代次数和批处理大小线性增长。

LLM中的异常值与奇异值分解（SVD）：
LLM在Transformer块输入张量中存在异常值，这些异常值出现在跨层的少数固定通道（即2D矩阵中的列）中，源于模型的内在属性（例如层归一化权重中的大幅值）。为了更好地预测重要token，可以通过奇异值分解（SVD）对查询和键矩阵进行偏置转换。对于实矩阵 $Q$，其SVD分解为 $Q = \mathbf{U}\Sigma\mathbf{V}^T$。通过乘以正交矩阵，可以旋转和拉伸向量，使得少数通道的幅度远大于其他通道。

图1：利用SVD从矩阵 $\mathbf{V}^T$ 到矩阵 $Q$ 的变换。正交矩阵 $A$ 最大化了 $\tilde{Q}$ 列向量之间幅度的差异。
图1：利用SVD从矩阵 $\mathbf{V}^T$ 到矩阵 $Q$ 的变换。正交矩阵 $A$ 最大化了 $\tilde{Q}$ 列向量之间幅度的差异。

KV缓存管理面临的挑战：
在基于卸载的系统中，将KV缓存移动到CPU可以解决GPU内存限制，但低PCIe带宽导致传输延迟成为瓶颈。传统的预取技术只能隐藏部分延迟。

图2：OPT-30B在不同序列长度和批处理大小下KV缓存和模型权重的总大小。虚线表示模型权重的大小。
图2：OPT-30B在不同序列长度和批处理大小下KV缓存和模型权重的总大小。虚线表示模型权重的大小。
图3：不同Transformer块执行方式之间的时序图比较。
图3：不同Transformer块执行方式之间的时序图比较。

现有的通过在有限预算内驱逐KV缓存来减小体积的方法（如 $\text{H}_2\text{O}$）存在以下挑战：
1. 迭代间注意力模式的动态特性：先前的驱逐策略假设当前不重要的token未来也不重要。然而，实际中当前不重要的token在后续迭代中可能变得重要。永久驱逐会导致序列变长时模型精度下降。
图4：具有完整缓存的基线模型与 (a) $\text{H}_2\text{O}$ 或 (b) Optimal 之间的注意力权重余弦相似度。
2. 跨层调整KV条目数量的必要性：不同层对KV缓存的需求不同。例如，第0层具有广泛的注意力模式，需要大量键token才能达到累计权重0.9；而第18层分布高度倾斜，大多数查询token仅需少数键token。
图5：直方图显示了OPT-6.7B模型 (a) 第0层和 (b) 第18层达到0.9总注意力权重所需的键token数量。分布跨层动态变化。
3. 跨查询调整KV条目数量的必要性：即使在同一层内，相邻的查询token所需的键token数量也存在巨大差异。固定比例的KV缓存预算无法适应这种方差，导致管理效率低下。因此，必须动态调整要选择的键/值token数量。

方法细节

利用CPU内存扩大评估窗口并动态预取：InfiniGen利用充足的CPU内存容量来增加识别KV缓存中重要token时的窗口大小。因此，在生成新token时，KV缓存的大部分token都保留在CPU内存中，而不是像先前的工作【37，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time+2023+NeurIPS】【78，$\text{H}_2\text{O}$: Heavyhitter oracle for efficient generative inference of large language models+2023+NeurIPS】那样完全丢弃它们。然而，该方法并不将整个KV缓存带到GPU进行注意力计算，而是仅加载和计算少数重要token的键和值，动态地丢弃其他不重要的token。为此，InfiniGen在CPU内存中维护KV缓存池，并选择性地、推测性地加载少量token。具体而言，使用前一个Transformer层的注意力输入来推测并预取当前层重要token的键和值。这种推测是通过在前一层执行当前层注意力计算的最小化演练来完成的。这允许仅传输对注意力计算至关重要的键和值，从而减少PCIe带宽的浪费，同时保持模型精度。此外，尽管KV缓存被卸载到比GPU内存更便宜且更大的CPU内存中，但仍需管理KV缓存池的大小，以免对CPU内存造成过大压力。

图6：InfiniGen设计概述。
图6：InfiniGen设计概述。

层间注意力输入的高相似性：预取模块建立在一个关键观察之上，即在LLM中连续注意力层的注意力输入高度相似。这主要有两个原因：一是LLM中存在异常值，二是由于层归一化（LayerNorm）。首先，Transformer块 $i$ 的输入（$\text{Tblock\_in}_i$）可以通过一系列带有残差连接和层归一化的操作从第 $i-1$ 层的输入（$\text{Tblock\_in}_{i-1}$）推导出来。因为注意力输出（$\text{Attn\_out}_{i-1}$）和FFN输出（$\text{FFN\_out}_{i-1}$）的输入都经过了层归一化，降低了每个值的幅度，所以它们的输出值相对于 $\text{Tblock\_in}_{i-1}$ 较小。因此，$\text{Tblock\_in}_i$ 高度受 $\text{Tblock\_in}_{i-1}$ 的影响。这种连续Transformer块之间高度相似的输入导致跨注意力层的输入相似。实验数据证实，$\text{Tblock\_in}_i$ 与 $\text{Tblock\_in}_{i-1}$ 之间的余弦相似度极高。InfiniGen利用这一关键观察，使用第 $i-1$ 层的注意力输入来推测第 $i$ 层的注意力模式。

通过正交矩阵变换实现权重矩阵倾斜：观察到注意力分数高度依赖于查询和键矩阵中的少数几列。大幅度列对注意力模式有巨大影响，因为查询和键之间的点积受这几列的影响很大。如果使查询和键矩阵中的少数几列具有比其他列大得多的幅度，那么少得多的列将显著影响注意力模式。这可以通过将查询和键权重矩阵乘以相同的正交矩阵 $A$ 来实现。由于正交矩阵的转置是其自身的逆矩阵，所提出的操作等价于 $Q \times K^T$，不会改变最终计算结果。具体来说，首先使用SVD分解查询矩阵 $Q = \mathbf{U}\Sigma\mathbf{V}^T$。然后将 $A$ 设置为正交矩阵 $\mathbf{V}$，以使列向量与标准单位向量对齐。倾斜的查询矩阵公式化为 $\tilde{Q} = Q \times A = \mathbf{U}\Sigma\mathbf{V}^T \times \mathbf{V}$。通过这种方式，可以在不改变计算结果的情况下，在 $\tilde{Q}$ 中制造出少数具有大幅度的列。

图7：(a) 连续Transformer块之间输入相似性的可视化。(b) OPT-13B模型第18层的查询矩阵。
图7：(a) 连续Transformer块之间输入相似性的可视化。(b) OPT-13B模型第18层的查询矩阵。

离线阶段的权重倾斜操作：在离线阶段，InfiniGen修改权重矩阵以生成倾斜的查询和键矩阵。为实现这一目标，InfiniGen首先使用样本输入运行模型的一次前向传递。在此过程中，InfiniGen从每一层收集查询矩阵，并对每个查询矩阵执行奇异值分解（SVD）。每层的倾斜矩阵使用查询矩阵的分解矩阵获得。然后将该矩阵与相应层中的每个查询和键权重矩阵相乘。相乘后，权重矩阵的维度保持不变。倾斜是一次性的离线过程，不会产生任何运行时开销。由于利用了源于模型固有属性的逐列模式，倾斜后的查询和键值会表现出高度的倾斜性，从而提高了预取模块的有效性。

图8：InfiniGen预取模块的操作流程。
图8：InfiniGen预取模块的操作流程。

Prefill阶段的局部权重生成：在预填充阶段，InfiniGen从查询权重矩阵和键缓存中选择几个重要的列来推测注意力模式，并生成用于解码阶段的局部查询权重和键缓存矩阵。因为需要将查询矩阵中的每一列与转置键矩阵中的相应行相乘，所以必须在查询权重矩阵和键缓存中选择相同的列索引。为了获得捕获异常值的局部矩阵，首先取倾斜查询和键矩阵的逐元素绝对值，然后将这两个矩阵加在一起。这有助于计算每列的总和并仅执行一次 top-$k$ 操作。接着对每列中的元素求和并选择矩阵中的 top-$k$ 列（本研究中选择了 $30\%$ 的列）。使用列值的总和捕获了每列的全局趋势，由于使用了倾斜的查询和键矩阵，选定的列更好地近似了注意力模式。

图9：Prefill阶段的局部权重生成。
图9：Prefill阶段的局部权重生成。

Decoding阶段的注意力分数推测与动态KV选择：在解码阶段，InfiniGen推测下一层的注意力模式并决定要预取的关键键和值。在第 $i-1$ 层，使用在Prefill阶段识别的第 $i$ 层的局部查询权重矩阵和键缓存，以及第 $i-1$ 层的注意力输入。在将局部查询和局部键缓存相乘后，InfiniGen选择具有高注意力分数的token。为了动态调整，InfiniGen考虑到推测注意力分数的最大值来设置阈值。仅选择注意力分数大于“最大分数减去 alpha”的token。由于从注意力分数中减去常数等价于在softmax之后进行除法（例如减去5导致权重被除以 $e^5 \approx 148.4$），排除这些token不会明显损害模型的准确性。由于多个注意力头并行计算，通过在各头之间平均最大分数和阈值之间的token数量，确保同一层中的每个头获取相同数量的token。InfiniGen从第1层开始推测和预取，因为利用输入相似性所需的异常值在第0层的计算期间出现。

图10：Decoding阶段的注意力分数推测。
图10：Decoding阶段的注意力分数推测。

基于计数器的KV缓存池内存淘汰机制：将KV缓存作为一个池进行管理并卸载到CPU内存中。虽然CPU内存较大，但仍有容量限制。因此，扩展了设计以纳入用户定义的内存大小限制。在运行时，当CPU内存的大小达到限制时，KV缓存池管理器选择一个受害者KV条目进行逐出。随后，管理器用新生成的键和值覆盖所选的受害者，并更新驻留在GPU中的相应局部键缓存。在受害者选择策略上，考虑了基于计数器的策略以及FIFO【7，The CacheLib caching engine: Design and experiences at scale+2020+OSDI】【69，Fifo queues are all you need for cache eviction+2023+SOSP】【70，CacheSack: Admission optimization for google datacenter flash caches+2022+USENIX ATC】和LRU【2，Memcached+URL】策略。基于FIFO的策略会导致较大的精度下降；基于LRU的策略表现出较小的精度下降但需要较高的运行时开销（需维护带锁的双向链表）。基于计数器的策略为每个预取的KV条目增加一个计数器，并选择计数最小的受害者（若计数器饱和则全部减半）。观察到基于计数器的策略和LRU策略显示出可比的模型精度，因此选择了基于计数器的方法，以简化设计并避免原子内存更新，从而获得更好的并行性。

方法细节参考文献汇总：
- 【37，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time+2023+NeurIPS】：在4.1节中被引用，作为先前完全丢弃KV缓存的代表性工作。
- 【78，$\text{H}_2\text{O}$: Heavyhitter oracle for efficient generative inference of large language models+2023+NeurIPS】：在4.1节中被引用，同样作为先前通过评估重要性来永久移除KV缓存条目的方法。
- 【7，The CacheLib caching engine: Design and experiences at scale+2020+OSDI】：在4.4节中被引用，作为广泛使用的FIFO软件缓存逐出策略的参考。
- 【69，Fifo queues are all you need for cache eviction+2023+SOSP】：在4.4节中被引用，作为FIFO策略的参考。
- 【70，CacheSack: Admission optimization for google datacenter flash caches+2022+USENIX ATC】：在4.4节中被引用，作为FIFO策略的参考。
- 【2，Memcached+URL】：在4.4节中被引用，作为LRU缓存逐出策略的参考。

实验环境

实验结果

  1. lm-evaluation-harness上的准确性实验：
    - 实验内容：在不同模型和5-shot任务上，比较InfiniGen与全缓存基线、$\text{H}_2\text{O}$、量化方法在不同相对KV缓存大小下的准确性。
    - 实验结果：当相对KV缓存大小小于 $10\%$ 时，InfiniGen始终表现出更好的准确性，而其他方法出现明显的准确性下降。当相对大小大于 $10\%$ 时，InfiniGen的准确性与全缓存基线紧密匹配，甚至在某些情况下略好。
    - 分析结论：InfiniGen能够有效减少KV缓存传输开销同时保留模型精度。减少参与注意力计算的KV缓存量有时能帮助模型更好地聚焦于关键token。
    - 图表引用：Fig 11。
    图11：LLM在lm-evaluation-harness中5-shot任务上的准确性。

  2. 序列长度对语言建模困惑度的影响：
    - 实验内容：在WikiText-2上，测量OPT-13B（2048长度）和Llama-2-13B（4096长度）随着生成块ID增加时的困惑度。
    - 实验结果：随着序列变长，InfiniGen的困惑度始终与全缓存基线相当，而 $\text{H}_2\text{O}$ 与基线的偏差越来越大。
    - 分析结论：$\text{H}_2\text{O}$ 受到永久KV缓存消除的困扰，而InfiniGen动态计算注意力，仅使用必要的KV缓存量，更适应长序列处理。
    - 图表引用：Fig 12。
    图12：OPT-13B和Llama-2-13B在WikiText-2数据集上的困惑度。越低越好。

  3. 倾斜（Skewing）效果实验：
    - 实验内容：在OPT-6.7B上，使用固定的 $20\%$ KV缓存预算，比较有无键/查询倾斜时的准确性。
    - 实验结果：如果没有倾斜，OPT-6.7B的准确性大幅下降。应用倾斜后，实现了与全缓存基线相似的准确性。
    - 分析结论：倾斜方法有效地调整了键和查询矩阵，使得少数列能更好地表示原始矩阵，对于某些模型（如OPT）至关重要。
    - 图表引用：Fig 13。
    图13：在OPT-6.7B上有无倾斜的lm-evaluation-harness基准测试准确性。

  4. 推理延迟与批处理大小可扩展性实验：
    - 实验内容：在OPT-13B（序列长度2048）上，测量批处理大小为20时的推理延迟，以及不同批处理大小（4到20）下的延迟。
    - 实验结果：InfiniGen比基线实现了 $1.63\times - 32.93\times$ 的加速。随着批处理大小增加，InfiniGen的吞吐量显著增加，而INT4和 $\text{H}_2\text{O}$ 的加速比饱和。
    - 分析结论：性能优势主要来自通过动态方法显著减少了从CPU内存加载的KV缓存量，实现了跨批处理大小的可扩展性能。
    - 图表引用：Fig 14, Fig 15。
    图14：OPT-13B在序列长度2048和批处理大小20下的推理延迟。
    图15：OPT-13B在序列长度2048下5种不同批处理大小的推理延迟。

  5. 序列长度与模型大小可扩展性实验：
    - 实验内容：在OPT-13B上测量不同序列长度（512到2048）的加速比；在不同模型大小（6.7B, 13B, 30B）上测量加速比。
    - 实验结果：InfiniGen的加速比随序列长度持续增加（最高 $5.28\times$），而其他方法饱和。在30B模型（需卸载 $30\%$ 模型参数）上，InfiniGen仍显示出 $1.34\times$ 的加速。
    - 分析结论：InfiniGen通过动态观察推测的注意力分数，自然捕获了重要token数量的非线性增长趋势，在更长序列和更大模型上具备更好的可扩展性。
    - 图表引用：Fig 16。
    图16：相对于FlexGen基线在 (a) 序列长度和 (b) 模型大小上的加速比。

补充细节

参数敏感性分析：

图17：在不同 (a) alpha 值和 (b) 局部权重比例下的准确性和推理延迟。
图17：在不同 (a) alpha 值和 (b) 局部权重比例下的准确性和推理延迟。

开销分析：
- 预取延迟分解：图18显示，FlexGen和 $\text{H}_2\text{O}$ 的主要瓶颈是数据传输（占 $>90\%$ 时间）。InfiniGen通过大幅减少数据传输量，仅比理想情况（无数据传输）慢 $1.52\times$。
- 内存消耗：局部查询权重和键缓存分别仅占总模型参数和总KV缓存的 $2.5\%$ 和 $15\%$。可以通过仅存储列索引或将局部键缓存放置在CPU中来进一步优化GPU存储开销。

图18：OPT-13B在序列长度2048和批处理大小8下Transformer块的延迟分解。
图18：OPT-13B在序列长度2048和批处理大小8下Transformer块的延迟分解。

超长上下文窗口的潜力：
- 在支持32K token的Llama-2-7B-32K模型上，图19表明随着相对KV缓存大小减小或序列变长，InfiniGen能保持接近全缓存的困惑度，而 $\text{H}_2\text{O}$ 差距显著拉大。
- 针对百万级token的分析（图20）显示，序列越长，关注极少数键的查询比例越高。且键token的注意力权重在迭代中会发生剧烈变化。InfiniGen保留暂时不重要的KV条目，能有效防止关键上下文丢失。

图19：Llama-2-7B-32K 在不同 (a) 相对KV缓存大小和 (b) 序列长度下的困惑度。
图19：Llama-2-7B-32K 在不同 (a) 相对KV缓存大小和 (b) 序列长度下的困惑度。
图20：使用Llama-3-8B-1048K对100万个token的分析。
图20：使用Llama-3-8B-1048K对100万个token的分析。

结论

KV缓存的庞大体积在基于卸载的高吞吐量推理系统中引发了严重的可扩展性问题，其大小甚至超过了模型参数。现有的KV缓存驱逐策略在应用于卸载系统时，不仅会导致模型准确性大幅下降，且未能有效利用互连带宽。本文提出了InfiniGen，这是一个基于卸载的动态KV缓存管理框架，能够高效执行大型语言模型的推理。InfiniGen利用前一层的注意力输入，推测性地预取重要token的KV缓存，并通过操纵查询和键权重使这种推测更加高效。实验证明，InfiniGen在大幅缩短推理延迟的同时保持了语言模型的性能，并且与现有解决方案相比，在批处理大小、序列长度和模型大小方面展现出了卓越的可扩展性。