可降低GPU内存的推理框架面世：韩国团队通过卸载键值缓存节约英伟达GPU内存，实现18.95倍注意力解码加速

DeepSeek 通过工程手段实现了 GPU 的良好利用，这一做法似乎引起了全球 AI 研究人员的共鸣。

近日，来自韩国科学技术院和 DeepAuto.ai 公司的研究人员，研发出一款名为 InfiniteHiP 的大模型推理框架。

InfiniteHiP 能在推理过程中将键值（KV，Key - Value）缓存置于主机内存，从而能够显著降低 GPU 内存压力。这一框架能在单个 英伟达 L40s 48GB GPU 上处理多达 300 万个 tokens，并且不会永久丢失上下文信息。 利用这一框架，研究人员在 100 万个 tokens 上下文中实现了 18.95 倍的注意力解码加速，并且完全无需额外训练。

同时，这一框架还能高效地利用长上下文，通过模块化分层 tokens 修剪算法，能够动态地消除不相关的上下文 tokens，从而能够提高模型效率。

（来源： arXiv ）

总的来说，InfiniteHiP 让预训练大模型能够实现超长的上下文推理。在不降低质量的前提下，能让模型超越了其原始上下文长度限制，同时可以通过高效的键值缓存卸载克服 GPU 的内存限制。

作为一种无需训练的解决方案，对于任何基于 Transformer 的预训练大模型来说，InfiniteHiP 都能充当一种直接的替代方案，因为它不仅可以在模型层面和硬件层面提供更快的推理速度，还能够延长上下文的可用长度。

从赫赫有名的注意力机制说起

对于基于 Transformer 的生成式大模型来说，扩展上下文长度至关重要。但是，对于大模型来说，处理非常长的上下文窗口却是一个重大挑战。一方面这会导致推理速度减慢，另一方面这会导致内存成本增加。此外，对于大多数已有的预训练大模型来说，它们往往无法超越其原始训练序列长度。

之所以存在这一难题，得从注意力机制说起。注意力机制——是模型的基本组成部分。注意力机制能够计算每个输入 token 与所有之前 token 之间的关系，随着输入序列长度的增加，这会导致计算成本和内存成本呈显出指数级增长。

键值缓存，则是注意力机制引起的另一个问题。在模型生成内容的过程中，此前计算的注意力键和值，会缓存在 GPU 内存之中，以供在必要时进行重用。但是，键值缓存大小会随着上下文长度出现线性变化，从而会给长上下文推理带来挑战。

此前，人们已经提出各种方法来降低注意力机制的高成本。在这些方法之中，有一种名为 FlashAttention（FA2）的方法，它是一种用于加速 Transformer 模型注意力计算的技术，通过避免将整个注意力得分矩阵写入全局 GPU 内存，以此来降低内存消耗和带宽利用率。然而，FA2 并没有降低计算成本。

为了减轻键值缓存的内存负担，业内人士做出了不少努力。比如，一种名为“KV 缓存驱逐”的方法，通过选择性地“遗忘”此前上下文来节省 GPU 内存。

然而，这些方法都会永久地删除此前上下文，而这些上下文很有可能还会被用到。

此前，有研究人员提出了 HiP 注意力机制，它能将不常访问的“冷 token”卸载到更大、更便宜的主机内存中，仅在需要使用时将它们提取回 GPU，同时只将那些需要经常访问的“热 token”保留在 GPU 上。

尽管上述优化起到一定作用，但是上下文扩展的所面临的另一个问题仍然存在：即预训练大模型无法处理超过其训练上下文长度的输入。

由于注意力机制是排列不变的，于是人们开始利用旋转位置嵌入（RoPE，Rotary Positional Embeddings）等位置嵌入方法来模拟 tokens 的时间顺序。

然而，由于大模型通常在固定长度的序列上进行预训练，因此当遇到更长的上下文提示时，它们无法适应那些没有见过的位置。

为了解决这一问题，人们采用了长上下文微调的方法，即在一组较长的输入上针对模型进行微调。然而，进行微调特别是针对长序列进行微调，既需要高昂的训练成本，也需要高质量的训练数据。

正因此，超长泛化（OOL，out-of-length）能力变得越来越重要。超长泛化指的是：在没有训练的情况下，预训练模型的能力远远超出其预训练极限。而作为一款长上下文大模型框架，InfiniteHiP 综合了上述所有方法的优点。

实现 7.24 倍的端到端解码加速

具体来说，为了减轻注意力的计算负担，研究人员基于 InfiniteHiP 提出一种模块化稀疏注意力新方案，该方案能够最大限度地减少不太重要的上下文的计算量。

为了优化键值缓存卸载，InfiniteHiP 通过基于 LRU 的缓存策略，增强了 HiP 注意力的卸载策略。（注：LRU 是 Least Recently Used 的缩写，中文意思为“最近最少使用”，这是一种常用的置换算法。）

基于大模型内部的注意力模式，InfiniteHiP 能够在不同组件之内，通过仔细应用各种 RoPE 调整策略来实现超长泛化。

InfiniteHiP 的另一个独特之处，在于它对修剪模块的创新使用。对于修剪模块来说，它能够通过采用一种新颖的模块化分层修剪算法，以便有选择性地丢弃不太重要的输入 token。

每个剪枝模块都能将输入序列划分为长度为固定值的块，同时也能并行识别出每个块中注意力得分最高的近似 token。而只有最重要的前 K 个块（其中 K 为常数）会被传递到下一个模块，其余的则会被丢弃。

通过堆叠多个剪枝模块，InfiniteHiP 能够针对块稀疏注意力掩码加以迭代优化。（注：块稀疏注意力掩码的英文是 Block Sparse Attention Mask，这是一种在处理长序列数据时用于优化注意力机制的技术，它通过将输入序列划分为较小的块（blocks），并在这些块内应用稀疏注意力，从而减少计算复杂度和内存消耗。）

（来源： arXiv ）

尽管 InfiniteHiP 是基于前人所提出的 HiP 打造而来，但它引入了几个关键性改进：

其一，与 HiP 框架基于启发式的分层修剪相比，InfiniteHiP 的分层修剪模块实现了更高的准确性。

其二，由于并行性得到增强，每个模块内的修剪算法速度明显变得更快。

其三，InfiniteHiP 采用模块化设计的方案，故能针对修剪阶段的缓存进行细粒度控制，从而能够实现更快的解码。

在 SGLang 大模型服务框架上， 研究人员证明了 InfiniteHiP 的有效性。 实验结果显示：在 3M token 上下文上，InfiniteHiP 能够实现 7.24 倍的端到端解码加速。

如下表所示，与表现最佳的基线大模型相比，基于 InfiniteHiP 这一框架，让 Llama 3 的相对得分提高 7.17%，让 Mistral 0.2 的相对得分提高 3.19%。

（来源： arXiv ）

与表现最佳的基线大模型相比，InfiniteHiP 在上述两个模型中，通过稀疏注意力所处理的密钥 tokens 减少了四分之三，从而能够产生更好的解码延迟。

为了进一步展示 InfiniteHiP 的超长泛化能力，研究人员将 Bench 的英语多项选择得分，在不同上下文长度下与 Llama 3.1 8B 模型进行对比。

（来源： arXiv ）

结果发现：随着上下文长度的增加，InfiniteHiP 的性能持续提升。而那些不具备超长泛化能力的基线模型，在超出预训练上下文长度（128K）后，性能会显著下降。

同时，研究人员还针对其他短上下文大型语言模型进行了实验。在这些短上下文模型中，研究人员在扩展上下文环境下观察到了十分显著的性能提升。

例如：与 FA2 方法相比，当在 Gemma2 模型中使用 InfiniteHiP 时，其在英语多项选择任务中得到 24.45% 以上的性能提升，在英语问答任务中得到 22.03% 以上的性能提升。

（来源： arXiv ）

解决长上下文推理问题，能通过键值缓存卸载节约 GPU 内存

与此同时，研究人员分析了 InfiniteHiP 中各个组件的延迟情况以及它们所产生的影响。

具体来说，研究人员分析了 InfiniteHiP 在 100 万 tokens 上下文中的延迟情况，并将其与能够产生相似基准测试分数的基线模型进行比较。

研究人员先是测量了注意力方法的延迟，在进行 100 万 tokens 的预填充时，其发现 InfiniteHiP 分别比 FA2 快 20.29 倍，比 InfLLM 出快 6%，并且与 HiP 的延迟相近。

在进行 100 万 tokens 上下文的解码时，InfiniteHiP 分别比 FA2 方法快 19.85 倍、比 InfLLM 快 4.98 倍、比 HiP 快 92%。

当启用上下文扩展时，针对预计算的余弦向量和正弦向量进行额外内存读取会产生开销，而 InfiniteHiP 在预填充时的速度会减慢约 1.6 倍，在解码时会减慢 5%。

在启用上下文扩展的预填充过程中，InfiniteHiP 比 InfLLM 慢 50%，但是在解码过程中却明显更快，这是因为解码过程会受到内存的限制。在 3k tokens 的上下文窗口之下，InfiniteHiP 读取的上下文 tokens 数量，比 InfLLM 在 12k tokens 的上下文窗口下读取的要少。

与此同时，研究人员在一个密码密钥检索任务样本上，测量了启用键值缓存卸载时的解码延迟。

在 256K 上下文解码中，InfiniteHiP 比 InfLLM 快 3.64 倍。InfLLM 在执行其注意力核时，采用的是不访问 CPU 内存的方法，因此不得不牺牲前 k 个估计算法的精度，这就需要更大的块大小和上下文窗口大小，来维持模型在下游任务上的性能。

相比之下，InfiniteHiP 在注意力核执行期间访问 CPU 内存，这为算法设计提供了更大的灵活性，在下游自然语言理解任务表现更佳。

在下图中，研究人员展示了使用 英伟达 RTX 4090（24GB）和英伟达 L40S（48GB）GPU 的方法的解码吞吐量。在英伟达 RTX 4090（24GB）上，与 SRT（SGlang Runtime with FlashInfer）的预估解码吞吐量相比，研究人员的方法在 1M 上下文长度下实现了 3.20 倍的吞吐量。在英伟达 L40S（48GB）GPU 上，研究人员的方法在 3M 上下文长度下比 SRT 高出 7.25 倍。

（来源： arXiv ）

在下图中，研究人员证明 InfiniteHiP 对于重要 tokens 具有更好的覆盖范围，这意味着关键 tokens 的注意力概率具有更高的召回率（recall）。更好的召回率表明，InfiniteHiP 比基线方法能够更紧密地遵循预训练的注意力模式。

（来源： arXiv ）

而在下图中，研究人员通过进行消融研究，找到了每种设置的延迟性能最优修剪模块组合。

（来源： arXiv ）