韩国通过卸载键值缓存节约英伟达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，从而能够提高模型效率。

总的来说，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），并在这些块内应用稀疏注意力，从而减少计算复杂度和内存消耗。）

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

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

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

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

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

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

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

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

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

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

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

解决长上下文推理问题，能通过键值缓存卸载节约 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 倍。

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

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

这些实验均说明：InfiniteHiP 是一个无需训练的大模型推理框架，它能够被用于高效的长上下文推理，支持超线性泛化和动态键值缓存卸载。

同时，它也有效解决了长上下文大模型推理中的三个主要挑战：不仅解决了长上下文的高效推理的问题，也解决了超长泛化的问题，并能在不“遗忘”此前上下文的同时，通过键值缓存卸载来节约 GPU 内存。

参考资料：

https://arxiv.org/pdf/2502.08910

运营/排版：何晨龙