OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs

  • 文章标题: OMNIKV：用于高效长上下文大语言模型的动态上下文选择
  • 作者/机构: Jitai Hao (山东大学), Yuke Zhu (蚂蚁集团), Tian Wang (蚂蚁集团), Jun Yu (哈尔滨工业大学), Xin Xin (山东大学), Bo Zheng (蚂蚁集团), Zhaochun Ren (莱顿大学), Sheng Guo (蚂蚁集团)

A1 主要贡献

本文针对大语言模型（LLM）在长上下文推理阶段因KV缓存导致GPU显存占用巨大的问题，提出了一种名为OmniKV的推理方法。

核心问题：在长上下文场景下，LLM的KV缓存会占用大量GPU显存，且其大小随序列长度线性增长。现有方法通常通过识别并丢弃注意力分数较低的不重要token来减少显存占用。然而，作者认为这种方法存在缺陷，因为注意力分数是基于当前隐藏状态计算的，只能反映token对当前推理步骤的重要性，而无法预示其在未来生成步骤中的重要性。被丢弃的token可能在后续推理中变得关键，导致信息丢失（如图1b所示）。

研究目标：设计一种无需丢弃token、无需训练的推理方法，既能保持与原始模型相当的性能，又能显著提升长上下文场景下的推理效率并减少KV缓存的显存占用。

创新点与核心洞见：
1. 核心洞见：层间注意力相似性 (Inter-Layer Attention Similarity)。本文发现，对于一个特定的上下文，在单次生成迭代中，不同层级之间识别出的重要token集合具有高度相似性（如图1a所示）。这意味着少数几层的注意力模式可以代表大多数层的注意力模式。
2. 动态上下文选择机制：基于上述洞见，OmniKV提出了一种动态选择机制。它在解码阶段仅在少数几个“过滤器”层（filter layers）上计算完整的注意力，并用它们来选择一个重要的token子集。
3. 高效的KV缓存管理与卸载 (Offloading)：
* 在预填充（prefill）阶段，OmniKV将大部分非“过滤器”层的KV缓存卸载到CPU内存。
* 在解码（decode）阶段，其他层直接复用“过滤器”层选择的token索引，仅从CPU加载所需的一小部分KV缓存子集到GPU进行稀疏注意力计算。
* 由于多个层共享相同的token索引，加载可以打包进行（Packed Load），并利用异步传输来重叠计算和数据传输时间，从而最小化PCIe传输开销。

  1. 性能与效率：
    • 无损性能：由于不丢弃任何token的KV缓存，OmniKV能够保持与原始模型相当的性能，尤其在多步推理（如思维链）场景下表现优越。
    • 显著加速：在上下文长度超过32K时，OmniKV开始展现加速效果。在128K上下文长度下，它实现了1.68倍的加速。
    • 扩展上下文长度：OmniKV能将单张A100 GPU上Llama-3-8B模型支持的最大上下文长度从128K扩展到450K。


图1：注意力分析。(a) 层间注意力相似性。这显示了即使相隔16层，层与层之间的重要token/子上下文也具有高度相似性。(b) 一个多跳问题的例子，展示了不同生成迭代中重要token的变化。(c) 重要token多样性分析。我们保留了重要token的集合，然后计算缺失token的累积注意力分数。

A3 背景知识/关键洞见

本文的工作基于三个关键洞见，并通过使用Llama-3-8B-Instruct、Yi-9B-200k和Llama-3.1-70B-Instruct等流行模型进行了实验验证。

层内注意力稀疏性。研究一致表明，LLM层内的注意力矩阵表现出稀疏性【【索引编号:35, Spatten: Efficient sparse attention architecture with cascade token and head pruning, 2021, 2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】【索引编号:7, Attention is naturally sparse with gaussian distributed input, 2024, https://arxiv.org/abs/2404.02690】】。这一特性意味着LLM可以通过关注一个简化的token子集来生成几乎等效的输出。一些研究已经基于稀疏性提升了推理速度或降低了GPU内存需求【【索引编 号:51, H2o: Heavy-hitter oracle for efficient gen- ´ erative inference of large language models, 2024, Advances in Neural Information Processing Systems】【索引编号:19, Snapkv: Llm knows what you are looking for before generation, 2024, https://arxiv.org/abs/2404.14469】【索引编 号:33, Quest: Query-aware sparsity for efficient long-context llm inference, 2024, Forty-first International Conference on Machine Learning】】。由于稀疏性的存在，OmniKV在大多数层中仅使用一小部分token子集。通过这种方式，我们不仅减少了计算量，还降低了CPU和GPU之间的通信量。


图2：OmniKV在解码阶段的整体框架。OmniKV中有三种类型的层。在预填充阶段，所有层都执行全注意力计算并生成上下文的KV缓存。然后，OmniKV将绿色层生成的KV缓存卸载到CPU内存。在解码阶段，橙色层由于稀疏性较低或为了推理效率而执行全注意力计算。紫色（过滤器）层执行全注意力计算，并 a) 基于在观察窗口上计算的注意力分数，使用上下文选择器选择重要的token。绿色层 b) 将前一个过滤器层选择的KV缓存子集加载到GPU，并执行稀疏注意力计算。c) 只有GPU上的KV缓存对LLM可见，用于生成下一个token。

层间注意力相似性。我们引入了层间注意力相似性的概念，其定义为：在一个特定层中受到显著关注的固定token子集，在后续的连续层中保持其显著性。一个层的相似性值也可以被看作是该层的“过滤”能力，计算方法为该固定token子集在后续层中注意力分数的总和的均值。图1a表明，在经过某个浅层之后，一个层与后续n个层的相似性变得异常高。虽然层间的整体相似性已经很高，但某些层比其他层表现出更强的“过滤”能力。我们将这些层称为“过滤器”层。随后，这些层在OmniKV中充当上下文选择器，为每个生成迭代识别关键token，从而促进后续层的稀疏注意力计算。

Token间注意力可变性。直观上，重要的token在LLM的生成过程中应该是变化的，特别是在多任务或多步推理场景下，如思维链（CoT）【【索引编号:36, Chain-of-thought prompting elicits reasoning in large language models, 2023, https://arxiv.org/abs/2201.11903】】。如图1b所示，对于一个多跳问题，我们在CoT场景中分别高亮了两个解码步骤中注意力分数最高的12个token。我们可以观察到，除了特殊 的BOS token外，其他重要的token完全不同。同时，如图1c所示，对Multi-Hop QA任务【【索引编号:11, Constructing a multihop qa dataset for comprehensive evaluation of reasoning steps, 2020, https://arxiv. org/abs/2011.01060】】的实证研究也证实了这种可变性。对于每个生成迭代，我们计算缺失token相对于一个token集合的累积注意力分数，该集合存储了预填充阶段25%最关键的token（即H2O中的Heavy Hitters）。图中观察到的峰值表明一些缺失的token具有显著的注意力分数。这一现象证实了我们的直觉，即识别出的关键token子集在不同生成步骤中存在显著波动。受此洞见启发，OmniKV保留了所有KV缓存，以确保性能不受影响。

A2 方法细节

基于上述洞见，我们提出了OmniKV，一个无需丢弃token且无需训练的推理方法。这种设计使得OmniKV能够在多步推理设置中保持LLM的性能。如图2所示，OmniKV包含两个关键模块：上下文银行（Context Bank）和上下文选择器（Context Selector）。

自回归LLM的推理阶段。自回归LLM的推理可以分为两个阶段：1) 预填充（Prefill），该阶段将输入提示的中间计算状态编码为KV缓存K, V，以避免冗余的KV向量计算，并输出下一个token作为解码的第一个输入；2) 解码（Decode），该阶段将前一个解码迭代预测的token作为当前token，并预测下一个token。

OmniKV的工作流程。在预填充阶段，OmniKV基于层间注意力相似性初始化上下文银行，将大多数“非过滤器”层的KV缓存存储在CPU内存中。在解码阶段，OmniKV采用一个即插即用的上下文选择器，在少数几个“过滤器”层上动态识别重要的KV缓存K, V子集。然后，上下文银行将这些选择传播到“非过滤器”层，并将相应的KV缓存子集打包加载到GPU内存中，因为这些层共享相同的token索引。通过这种方式，OmniKV减少了计算和数据传输成本。

4.1 上下文银行

利用层间相似性预取重要token。所提出的上下文银行利用层间注意力相似性来预取重要token。在GPU内存不足的情况下，上下文银行还可以从CPU内存异步预加载相应的KV缓存，从而缓解内存限制。

Prefill阶段的KV缓存创建。为简化分析，我们此处忽略批量大小。在预填充阶段，一个L层的LLM通过将注意力投影矩阵$W_i^k$和$W_i^v$应用于隐藏状态$h_i^p$，创建KV缓存$\{K_i, V_i\}_{i=1}^L$，从而为键（key）和值（value）生成维度为$R^{H \times N \times d}$的张量。这里，N表示预填充上下文中token化提示p的长度，H表示注意力头的数量，d表示每个注意力头的隐藏大小。

选择“过滤器”层。首先，我们需要确定哪些层在识别重要token方面更有效。这些选定的层被称为“过滤器”层。如图1a所示，尽管相邻16个层之间的相似性很高，但与16个相邻层相比，8个相邻层的平均相似性之间观察到显著差距。为了通过减小“过滤器”层之间的距离来提升性能，OmniKV使用了一组超参数L，其大小为m（m ≤ 3）。与使用单个“过滤器”层相比，使用多个“过滤器”层时，“非过滤器”层使用的子上下文理论上表现出更高程度的相似性。

上下文选择器与稀疏注意力。因此，OmniKV在L中的层上执行全注意力计算，以识别一个重要的token小子集$T_i, (i \in L)$。由于浅层网络的稀疏性较低，OmniKV还在$L_0$层之前的层（$l < L_0$）上执行全注意力计算而不进行选择。然后，OmniKV仅使用这些重要token $T_i$作为稀疏注意力层$l$（其中$L_i < l < L_{i+1}$）的子上下文。这里，$h_i^w$表示第i层观察窗口的隐藏状态。上下文选择器识别出在观察窗口上具有显著注意力分数的重要token。

避免GPU等待并制定注意力机制。为了在加载KV缓存时避免不必要的GPU等待，OmniKV还在L的相邻层（表示为$\{l+1\}_{l \in L}$）上执行全注意力计算。这将数据传输与计算交错进行。这里，$h_l$表示最后一个token的隐藏状态。最后，整个注意力机制可以公式化如下：

从CPU检索子集。OmniKV在稀疏注意力层中将序列长度显著减少到不足10%，这导致了时间复杂度的降低。在$L_t$（其中$L_t=i$）中识别出关键token $T_i$后，OmniKV从CPU内存中检索相应的KV缓存子集$K_j[T_i], V_j[T_i]$（其中$L_t+1 < j < L_{t+1}$）作为稀疏层的子上下文。

打包加载。由于“过滤器”层之间的各层共享相同的子上下文token索引T，一系列连续稀疏注意力层的KV缓存可以被打包，并在最近的前一个“过滤器”层从CPU加载到GPU。因此，OmniKV只进行m次（m ≤ 3）加载，与在每一层都进行加载相比，显著减少了缓慢的PCIe传输开销。

4.2 上下文选择器

统一的token选择框架。如第4.1节所述，OmniKV在“过滤器”层L中选择重要的token $T_i$。受先前工作的启发【【索引编号:19, Snapkv: Llm knows what you are looking for before generation, 2024, https://arxiv.org/abs/2404.14469】【索引编 号:38, Infllm: Training-free long-context extrapolation for llms with an efficient context memory, 2024, https://arxiv.org/abs/2402.04617】】，我们提出了一个统一的token选择框架。OmniKV基于一个得分向 量$S_i \in R^N$来选择重要token。得分$S_i$是使用一个观察窗口$h_i^w$计算的。我们以局部窗口作为查询状态，以完整上下文$h_i^c$作为键状态，计算注意力得分$A_i$：

计算得分并使用top-k选择。接下来，为了得到得分$S_i$，我们首先应用reduce-max操作来获得跨注意力头的最大得分。随后，利用一个加权向量$\alpha$对注意力得分进行加权求和。最后，我们利用topk在得分$S_i$上识别出重要的token $T_i$。

不同的选择器方法。不同的$\alpha$值会为局部窗口的token分配不同的权重，并产生不同的选择模式。为进一步研究观察窗口中的哪些token具有更强的“过滤”能力来识别重要token T，本研究探讨了三种方法：1) Uniform：$\alpha = \{1\}_{i=0}^{|h_i^w|}$。这种方法意味着窗口中的每个token在加权求和注意力分数时产生等效的贡献。2) Exponential：$\alpha = \{2^{i-|h_i^w|}\}_{i=0}^{|h_i^w|}$。这种方法意味着越靠近窗口末尾的token产生越高的贡献。3) LastToken：concat(α = {0}|hwi |−1i=0 , {1})。这种方法意味着我们只考虑最后一个token。

A4 实验环境

  • 模型架构:

    • Llama-3-8B-262K
    • Yi-9B-200K
    • Llama-3.1-70B-Instruct
    • Llama-3.1-405B (附录中测试)

  • 数据集:

    • InfiniteBench: 平均长度145.1K，覆盖多项任务，测试时统一使用128K上下文。
    • LongBench: 包含18个任务，平均长度在5K到15K之间。
    • 多步推理数据集:
      • 2WikiMQA 和 HotpotQA (来自LongBench)
      • 2StageRetr: 作者提出的新基准，用于评估多步推理能力，包含一个字典和一个加法方程，需要先计算再检索。

  • 硬件配置:

    • GPU: Nvidia A100 GPU
    • CPU: Intel Xeon Platinum 8369B CPU (12核, 2.90GHz)

  • 软件配置与实现细节:

    • 解码策略: 大多数任务使用贪心解码；InfiniteBench中的摘要任务使用top-p解码（p=0.95, temperature=0.8）。
    • 模型精度: Llama-3.1-70B使用bitsandbytes进行4位权重 quantization；其他模型使用float16。
    • 代码库: 基于Huggingface的transformers进行少量修改。

    • OmniKV超参数:

      • 局部窗口大小（uniform和exponential选择器）: 16
      • “过滤器”层 L 的设置:
        • Llama-3-8B-262K: {2, 8, 18}
        • Yi-9B-200K: {6, 11, 30}
        • Llama-3.1-70B-Instruct: {4, 19, 41}

    • 基线方法 (Baselines):

      • H2O: 基于注意力分数丢弃token。
      • InfLLM: 不丢弃token，将序列分块并选择代表性token进行检索。
      • StreamingLLM: 保留初始窗口和滑动窗口作为上下文。
      • Full Attention: 原始模型，作为性能上限。

A4 实验结果

5.1 性能评估

  • 单步推理性能:

    • 实验内容: 在LongBench和InfiniteBench两个标准基准上评估单步推理性能，模型直接回答问题。
    • 实验结果 (Table 1, 2): OmniKV在所有任务上均取得了最佳性能，且结果非常接近全注意力基线，展现了其稳定性。尤其在Llama-3.1-70B模型上，OmniKV显著优于其他基线方法。与InfLLM相比，OmniKV在不同任务上的表现更稳定，例如在KV Retrieval任务上远超InfLLM。

    表1：在LongBench上的单步推理性能。(Bai et al., 2023)。斜体表示模型使用全注意力基线。粗体表示同一模型下的最佳性能。~ 指的是InfLLM的KV缓存内存预算由于其高度集成的实现而大致设置为特定值。+ 表示由于与flash attention不兼容，在H2O的每个块内丢弃token。每个子任务的详细结果见表4和表5。

    表2：在InfiniteBench上的单步推理性能。(Zhang et al., 2024a)。

  • 多步推理性能:

    • 实验内容: 在2WikiMQA、HotpotQA和2StageRetr三个多步推理任务上进行测试，采用思维链（CoT）输出格式，并使用精确匹配（Exact Match）作为评估指标。
    • 实验结果 (Figure 3): OmniKV在所有预算设置下都取得了最佳性能，证明了其动态上下文选择的有效性和不丢弃token的必要性。在2StageRetr任务中，H2O的准确率无法超过其token预算比例，说明它只能随机保留字典中的键值对。


    图3：在三个多跳任务上，不同KV缓存token预算比例下的多步推理性能。

  • 上下文选择器选择:

    • 分析结论: 在单步推理中，“exponential”和“last”选择器表现更优；在多步推理中，“last”选择器性能最佳。从工程实现和推理效率角度看，“last”选择器最简单且延迟最低，因此后续分析主要采用该选择器。

5.2 延迟与权衡

  • 端到端延迟评估:

    • 实验内容: 使用单张NVIDIA A100 80GB GPU，在Llama-3-8B-262K模型上评估端到端推理延迟，批大小为1，稀疏注意力层的token预算为2048。
    • 实验结果 (Figure 4):
      • 解码阶段: 当上下文长度超过32K时，OmniKV的延迟低于全注意力。在128K上下文下，不进行卸载的OmniKV（OmniKV w/o offload）实现了1.68倍的加速（21.0 tokens/s）。在GPU显存不足时，通过卸载机制，OmniKV能在单A100上以7.5 tokens/s的速度处理450K上下文，比在3张A100上运行的原始模型快1.87倍。
      • 预填充阶段: OmniKV的延迟与全注意力相同，因为卸载过程被计算时间所覆盖。


    图4：端到端延迟结果。左图和右图分别显示解码和预填充阶段的延迟。OOM表示某些方法在单张A100 GPU上支持的最大上下文长度。

  • 性能与效率的权衡:

    • 实验内容: 在Llama-3-8B-262K上，使用不同token预算在InfiniteBench（128K上下文）上评估性能与效率的权衡。
    • 实验结果 (Figure 5): 即使每个稀疏层只检索128个相关token，OmniKV的平均得分（35.9）也高于H2O+。选择1024个token可以在性能和延迟之间取得良好的平衡。


    图5：稀疏层token预算的权衡。左轴显示在InfiniteBench上的平均得分，右轴显示解码延迟。

5.3 分析

  • “过滤器”层的任务无关性:

    • 分析: 对比了Llama-3-8B和Yi-9B在不同任务上的“过滤”能力（层间注意力相似性）。
    • 结论 (Figure 6a): 不同任务的相似性曲线趋势基本一致，表明“过滤”能力是模型层的内在特性，而非任务依赖。

  • 上下文选择的准确性:

    • 分析: 使用CLongEval数据集测试不同层定位答案所在上下文块的命中率。
    • 结论 (Figure 6b): “过滤”能力更强的层（如Llama-3-8B的第8层）在识别包含答案的重要token方面表现出更高的命中率。

  • 性能与“过滤”能力的关系:

    • 分析: 在LongBench上测试了使用不同“过滤器”层的性能。
    • 结论 (Figure 6c): 模型的性能与所选“过滤器”层的“过滤”能力有很强的相关性。选择“过滤”能力强的层（如8、10、11、13层）能带来更好的性能。


    图6：对层“过滤”能力的分析。对于(a)和(b)，左：Llama-3-8B-262K，右：Yi-9B-200K。(a) 层间注意力相似性，也称为层的“过滤”能力，在不同任务中显示出相似的趋势。(b) 捕捉重要token的能力不同。(c) 在LongBench上使用Llama-3-8B-262K的不同性能表现。

A5 结论

本文提出了OmniKV，一个无需丢弃token且无需训练的推理方法。它在长文本场景下，能够在不牺牲性能的前提下，将推理效率提升1.7倍。此外，OmniKV与卸载技术高度兼容，显著减少了KV缓存的内存消耗。该方法实现简单，具有广阔的实际应用前景。

未来工作: 计划探索将OmniKV与KV缓存量化技术相结合，以进一步减少KV缓存的使用并提升效率。

A6 附录

A 2STAGERETR

基准测试设计。为了减轻先验知识对结果的影响，我们提出了2StageRetr基准。其核心思想是构建一个两步推理问题，利用第一步的结果来完成第二步。任务描述被放置在提示的末尾，以防止模型“心算”出第一步的结果，从而避免LLM基于心算结果为答案token分配高注意力分数。具体来说，2StageRetr由一个包含多对数字和颜色组成的字典以及一个加法方程构成。我们确保加法的结果总是字典中的一个键。然后模型需要根据加法结果在字典中找到对应的颜色并输出。以下是2StageRetr的一个例子。

2StageRetr 示例。

<blockquote>

我们来玩个游戏。你有一个字典和一个数学加法方程。字典的键可以是任何数字。你需要在执行加法后，在字典中找到对应的键值，并输出该键对应的值。
字典是 {0: lime, 1: yellow, 2: red, 3: black, 4: brown, ...... 17: brown, 18: maroon, 19: teal, 20: red, ..... 28: brown, 29: violet}
方程是 8 + 10 = ?
根据加法结果回答相应的颜色。

示例输出：因为 8 + 10 = 18，所以对应的颜色是“maroon”。

</blockquote>

设计原理。这种设计本质上确保了LLM在没有问题先验知识的情况下，使用预训练的先验知识为字典中的某些键值对分配更高的注意力分数。因此，H2O方法基本上等同于从字典中随机丢弃键值对。因此，其性能不能超过其token预算的比例。相比之下，我们的方法每次都根据观察窗口内的信息选择最相关的上下文，从而实现更优的性能。

数据集细节。我们设置字典大小不超过200，数字从0到最大值顺序排列。这种算术序列对于当前的LLM来说应该相对简单，从而主要评估模型的检索能力。该数据集的平均长度仅为739个token，最大长度为1382个token。然而，当前基于缓存丢弃的方法在该任务上表现不佳。

B 详细实现

伪代码。实际代码实现仅需对Huggingface的Transformers库【【索引编号:37, Huggingface’s transformers: State-ofthe-art natural language processing, 2020, https://arxiv.org/abs/1910.03771】】的源代码进行少量修改。这里我们展示了OmniKV在算法1中的核心伪代码。该算法演示了单层的注意力前向传播过程 。

算法1：OmniKV的注意力前向传播
输入：观察窗口的隐藏状态hw，过滤器层L，上下文缓存K, V，可见KV缓存Kv, Vv，当前层i，注意力查询权重Wq，token预算k，窗口权重α
if i ∈ L then
    Qwi ← Wqi hwi ▷ 获取查询状态
    Ai ← Softmax(Qwi K⊤i / √d) ▷ 获取注意力分数
    Si ← Σ|hwi|−1j=0 αj max0≤h<H Ai[h, j] ▷ 获取上下文分数
    Ti ← arg top k(Si) ▷ 选择重要token 0≤t<|hci|
    t ← 获取i在L中的索引
    for j = Lt + 2 → Lt+1 − 1 do
        Kvj , Vvj ← LoadToGPU(Kj [Ti], Vj [Ti]) ▷ 加载KV缓存子集到GPU
    end
end
if i ∈ L or i < L0 or i − 1 ∈ L then
    Kvi , Vvi ← Ki, Vi ▷ 对执行全注意力的层使用原始缓存
end
最后，使用可见KV缓存Kvi , Vvi 执行正常的注意力计算。

C 基线设置

详细设置说明。此处我们提供不同基线设置的详细描述。

  1. H2O。由于H2O需要注意力分数的输出，这与flash attention不兼容，处理长序列时中间激活值会直接导致内存溢出。因此，当上下文长度超过60K或使用大于30B的模型时，我们必须分块计算注意力，并根据每个块内的注意力分数移除token。我们在结果中用+表示这种方法。为确保比较的公平性，我们严格设置H2O保留Mem%的KV缓存。我们还修改了代码以支持Llama-3和Yi中的GQA【【索引编号:2, Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023, https://arxiv.org/abs/2305.13245】】 。
  2. InfLLM。由于InfLLM使用基于LRU的块缓存且有许多超参数，按百分比限制KV缓存会导致效率和性能下降。因此，InfLLM被配置为平均使用Mem%的KV缓存。
  3. StreamingLLM。我们严格设置StreamingLLM保留Mem%的KV缓存，其中1%用于初始sink token，Mem-1%用于局部窗口token。
  4. Full Attention (original model)。为避免MLP中的中间激活值消耗过多GPU内存，我们将隐藏状态沿序列长度维度拆分为大小为4000的块。这在不影响效率的情况下降低了峰值内存使用。

D 消融研究和详细实验

D.1 “大海捞针”测试

实验设置与结果。由于OmniKV使用8B模型在单个A100 GPU上能够支持450K上下文，我们进行了最大上下文长度为450K的“大海捞针”测试，并取得了完全准确的性能。我们继续使用L = {2, 8, 18}，稀疏层的token预算为1024。我们还在单个NVIDIA H20上用512K的输入长度进行了测试，同样取得了完全准确的结果。结果如图7所示。

D.2 详细的权衡分析

效率与性能权衡。这里，我们展示了效率与性能之间详细的权衡。我们在InfiniteBench上的测试结果如表3所示。此处的延迟指的是解码阶段每个token的端到端时间。

算法1：OmniKV的注意力前向传播
表3：性能与效率权衡的详细结果。


图7：“大海捞针”测试。结果表明，应用于Llama-3-8B-1M的OmniKV实现了完美的检索结果。

D.3 LongBench详细结果

各子任务得分。我们展示了在LongBench上的详细结果，列出了每个子任务的得分，如表4和表5所示。

表4：LongBench上的详细结果（第一部分）。

D.4 Yi-9B-200K的多步推理结果

进一步验证。为了进一步验证OmniKV在多步推理中的有效性，我们继续使用Yi-9B-200K进行实验。2StageRetr的平均长度为739，因此token预算为739×0.067 = 49。为避免预算过低，同时确保比较的公平性，我们设置“过滤器”层L = {3, 11, 30}，为稀疏层分配更多的token预算。

结果分析。如表6所示，OmniKV在所有三个数据集上均取得了最佳结果，特别是在受限的token预算下。这进一步验证了我们方法的有效性。Yi-9B-200K在2StageRetr中没有遵循指令直接回答。

D.5 消融研究

模块重要性。OmniKV由上下文选择器和上下文银行组成，这两个模块是高度耦合的。如果移除上下文选择器，意味着需要计算整个上下文，那么在GPU内存充足的情况下，该方法等同于全注意力。否则，我们必须卸载某些层的KV缓存，这时该方法就变得与“暴力卸载”（Brutal Offload）相同。尽管“暴力卸载”的性能完全等同于原始模型，但在解码阶段频繁的加载和卸载会带来巨大的开销。

表5：LongBench上的详细结果（第二部分）。

表6：Yi-9B-200K的多步推理结果。

上下文银行的作用。如果没有上下文银行，意味着我们不应用层间注意力相似性，那么我们只能在每一层使用上下文选择器来选择token集T。然而，选择重要token/上下文本身需要在当前层进行全注意力计算，这意味着效率将显著降低。唯一潜在的好处可能是使用更短的上下文来避免无关信息的干扰，从而可能带来更好的性能。

表7：性能与效率权衡的详细结果。

D.6 过滤器层的性能

理论与实践。理论上，如果我们选择具有更强“过滤”能力的“过滤器”层，可以预期获得更好的性能。然而，在测试不同选定层时，要确保完全公平的性能比较是具有挑战性的。这是由于层的顺序性；即使我们确保执行全注意力的总层数保持一致，仍然可能存在“过滤器”层之间间距不均的问题。尝试均匀分布“过滤器”层集合会引入额外的变量。尽管如此，实验结果表明，具有增强“过滤”能力的层在某种程度上倾向于表现出更优的性能。我们在LongBench和InfiniteBench上进行了实验。结果如表8所示。

表8：在LongBench上不同过滤器层设置的性能。

D.7 层间注意力相似性图示例

可视化分析。为了直接观察层的“过滤”能力，我们还展示了任意两层之间的相似性。我们可视化了Llama-3-8B-262K在HotpotQA上top-2048 token集的累积注意力分数作为相似性的度量，与之前保持一致。如图8所示，对于第8层，尽管与第20层相距12层，由第8层选择的token集T8在第20层仍然达到了0.87的累积注意力分数。这充分证明了第8层选择重要token的有效能力。

D.8 OmniKV在更大模型上的有效性

405B模型分析。我们首先可视化了Llama-3.1 405B的层间注意力相似性，如图9所示。405B模型仍然表现出非常高的层间相似性，表明OmniKV可以有效地应用于该模型。此外，我们在LongBench的两个任务Qasper和Qmsum上进行了进一步评估。结果呈现在表9中。

表9：OmniKV在405B模型上的表现。

D.9 OmniKV的兼容性

与推理引擎集成。本文中的大多数实验是使用Huggingface Transformers进行的。然而，这个框架通常不作为推理引擎使用。当前用于大模型推理的引擎，如vLLM【【索引编号:16, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】】和LightLLM【【索引编号:26, Lightllm: An efficient lightweight llm serving framework, 2023, https://github.com/ ModelTC/lightllm】】，比基于Hugging Face Transformers的引擎快得多。因此，我们将OmniKV适配到其中之一，即LightLLM，以证明OmniKV在真实世界场景中具有出色的可用性。结果如表10所示。

连续批处理 (Continuous Batching)。连续批处理是一种在LLM训练和推理中用于优化计算效率的技术，特别是在模型处理不同长度序列的场景中。其核心思想是动态地将输入序列分组为可以并行处理的批次，即使序列长度不同。OmniKV可以自然地与连续批处理集成。基于LightLLM，我们发现连续批处理技术将多个64K长度请求的平均延迟从41.7秒降低到37.8秒。


图8：层间注意力相似性图示例。

表10：每个请求的平均每token延迟（ms/token, tp=4）。

张量并行 (Tensor Parallelism)。张量并行(TP)是一种在LLM和其他深度学习模型中使用的技术，用于将张量（多维数组）的计算分布到多个设备（如GPU）上。这种方法对于训练和推理那些无法装入单个设备内存的非常大的模型特别有用。使用LightLLM，我们实现了带张量并行的OmniKV。如表11所示的实验结果表明，启用TP后，与原始模型相比，OmniKV仍然可以实现相当大的加速。另一个观察是，随着序列长度的增加，Lightllm+OmniKV的解码延迟降低。这是因为“非过滤器层”的序列长度固定为2048。随着批大小的减小，“非过滤器层”的计算量也随之减少。

TP的更多技术细节。在张量并行中，注意力头被分布在不同的GPU上，输出通过all-reduce聚合，使得GPU间通信成为效率的关键因素。对于tp=n，OmniKV为每个卡维护n个重要token的索引。每个重要token的索引仅由该卡自己的注意力头计算，从而消除了GPU间的通信需求。由于不同的注意力头可能有不同的索引，系统直观上表现出更大的灵活性。我们也通过实验证实了这一点（见表12）。


图9：Llama-3.1 405B的层间注意力相似性。

表11：解码吞吐量结果（tp=4），使用Llama-3-8B-Instruct-262K。

表12：跨数据集的性能比较。

流水线并行 (Pipeline Parallelism)。流水线并行(PP)是一种用于在多个设备或处理器上并行化LLM执行的技术。这种方法对于深度神经网络特别有用，其中模型可以被划分为较小的、顺序的阶段，每个阶段可以在不同的设备上处理。OmniKV也自然兼容流水线并行。表1中显示的Llama-3.1-70B的实验就是使用流水线并行进行的。

上下文并行 (Context Parallelism)。上下文并行(CP)是一种在LLM中用于通过将输入序列的处理分布到多个设备或处理器上来提高训练和推理效率的技术。这种方法利用了许多LLM（如transformer）处理输入序列的方式允许并行计算的特点。实际上，当使用OmniKV时，没有必要启用上下文并行。这是因为OmniKV仅需约2048个token就能取得优异的结果。尽管如此，OmniKV仍然可以与上下文并行集成。在这种情况下，多GPU并行中的通信延迟通常超过单个GPU的计算延迟。因此，我们可以仅在密集（全）注意力层中使用上下文并行，而在稀疏注意力层中使用单个GPU。这意味着我们只需配置哪些层将实现上下文并行。当批大小大于1时，我们可以将样本分配到不同的GPU以防止资源浪费。

D.10 预填充加速

设计与实现。OmniKV旨在提升解码速度。然而，考虑到预填充加速的重要性，我们修改了OmniKV以优化预填充阶段。实验结果表明，与输入长度为256K的原始模型相比，OmniKV-prefill在不损失性能的情况下实现了1.90倍的延迟降低。此处的代码是使用Transformers库实现的，并用Flash Attention进行了增强。具体来说，在预填充阶段，我们将序列维度划分为多个块进行计算。对于每个块，我们也使用“过滤器层”来选择重要token。由于查询块的序列长度不再是1，我们采用第4.2节中介绍的uniform计算方法来在“过滤器层”中选择重要token。

结果分析。表13和表14显示了延迟和质量保持结果。这里，“r”表示所选token的大小，token大小与块大小的比率为“r”。因此，r的值越大，信息损失越少。结果表明，r=2在效率和性能之间取得了良好的平衡。MInference【【索引编号:14, Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention, 2024, arXiv preprint arXiv:2407.02490】】专为预填充加速设计，并优化了多个CUDA核。但OmniKV仅依赖于Huggingface-transformers库，因此效率较低。

表13：在单个NVIDIA-H20上的预填充延迟（秒）。

表14：在LongBench上的质量保持结果。