RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

作者/机构: Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Danning Ke, Shikuan Hong, Yiwu Yao, and Gongyi Wang (Huawei Technologies Co., Ltd)

A1 主要贡献

长上下文大型语言模型（LLM）在多种自然语言处理任务中取得了显著进展，但随着输入长度增加，键值（KV）缓存的增长已成为部署的主要瓶颈。现有的KV缓存压缩方法，如量化、token丢弃和局部注意力，试图缓解此问题。其中，一个主要方向是直接丢弃被认为不重要的token，但这种方法存在一个固有假设：即当前不重要的token在未来的查询中也不会被需要。这个假设在实际场景中并不成立，例如当用户查询与文本主旨不直接相关的信息或进行多轮对话时，可能会导致所需信息被丢弃，从而严重影响性能（如图2所示）。这引出了一个关键问题：“我们能否在不丢失语义信息的情况下减小KV缓存的大小？”

本文从一个新的视角解决了这个问题。研究发现，LLM在处理长上下文时存在一种“检索与处理”机制。具体来说：
1. 检索头（retrieval heads）：LLM通过一部分特定的注意力头（我们称之为“检索头”）从长输入中精确地回忆起被查询的信息。这些头能够将其大部分注意力权重集中在与查询相关的信息上。
2. 非检索头（non-retrieval heads）：其余的非检索头主要关注局部上下文或注意力池（attention sink），无法有效利用所有输入中的语义信息。

基于这些发现，本文提出一个假设：LLM的推理过程是基于“检索与处理”的。模型首先使用检索头收集相关信息，然后非检索头处理这些检索到的信息并生成最终响应。这启发我们为不同类型的注意力头设计独立的缓存策略：
* 对于检索头，我们保持KV缓存的完整性。
* 对于非检索头，我们只缓存最近的token和注意力池。

为了进一步减小直接丢弃非检索头中远程token所带来的精度损失，本文设计了一种“补偿token”机制，将丢弃的缓存压缩成一个token，从而进一步恢复信息。结合检索头和补偿token，本文提出的算法RazorAttention能够成功将KV缓存压缩70%，而性能没有明显下降（如图1所示）。RazorAttention甚至可以支持对1024K序列的压缩，在将KV缓存压缩50%后，在“大海捞针”任务中实现了几乎无损的精度（如图5所示）。

此外，以往基于重要性的token丢弃方法依赖注意力权重来计算重要性得分，这使得它们无法与作为长上下文推理关键组件的FlashAttention兼容。RazorAttention解决了这个问题，因为它不使用注意力图作为度量标准。其逐头的剪枝标准与FlashAttention完全兼容，并且补偿token的计算开销可以忽略不计。因此，RazorAttention在推理加速方面比以往方法有显著提升。

本文的主要贡献总结如下：
* 系统性分析：我们系统地分析了Transformer在长输入下的注意力动态，揭示了只有少数“检索头”能从整个输入中回忆信息，而其余的头主要关注局部上下文。
* 新颖算法RazorAttention：我们提出了一种名为RazorAttention的新算法，能够在对8K到100K token上下文的性能影响极小的情况下，将KV缓存大小减少70%。我们设计了一种精确且无需数据的度量标准来分配所有检索头，并结合了一种错误补偿策略来弥补因截断KV缓存造成的信息损失。
* 高效与易用性：RazorAttention在压缩过程中引入的开销可忽略不计，并且与FlashAttention兼容，使其成为一种高效、即插即用的解决方案，无需训练或显著开销即可提升LLM的推理效率。大量实验证明，RazorAttention可以有效地应用于各种模型和任务。

图 1: RazorAttention在压缩70%的KV缓存后，仍能实现与原始模型相当的性能。为了证明这一点，我们在大海捞针基准测试【gkamradt, 2023】上测试了Llama2-13B-64K【Fu et al., 2024】。

图 2: 基于重要性的token丢弃方法在查询与主旨不太相关的信息时会失效。这里，我们使用一个来自LongBench【Bai et al., 2023b】的8K文档，并添加了两句与主旨无关的句子。在这种情况下，H2O丢弃了与主旨不太相关的token，导致在Q1和Q2中都失败了。SnapKV根据第一个查询丢弃token，对Q1有效，但在后续查询（如Q2）中失败。只有RazorAttention即使在我们压缩了70%的KV缓存后，也能成功地从长输入中输出确切的信息。

A2 方法细节

3.1 针对ALiBi模型的RAZORATTENTION

ALiBi模型的注意力分数计算。对于使用ALiBi位置嵌入的模型，其第h个注意力头的注意力分数根据以下公式计算：

$$S_{m \rightarrow n}(\boldsymbol{q}; \boldsymbol{k})=\boldsymbol{q}_{m} \boldsymbol{k}_{n}^{\top}-l_{h}(m-n),$$

其中，$q_m$是位置m的查询张量，$k_n$是位置n的键张量，$l_h$是该头特定的斜率，$S_{m \to n}(q; k)$是注意力分数。需要注意的是，由于注意力的因果性，$(m \ge n)$是保证的。

基于注意力衰减的token丢弃。在$l_h(m-n)$显著大于$q_m k_n^\top$的情况下，查询$q_m$和键$k_n$之间的注意力将衰减至零，这意味着对于位置m的输出，任何位置比n更远的token的贡献都变得可以忽略不计。下面的定理将这一观察形式化。

定理1：注意力权重上界。给定一个根据公式1计算注意力分数的注意力头，对于任何$\epsilon \in (0, 1)$，从$q_m$到$k_n$的注意力权重可以被上界约束为：

$$\begin{aligned} \begin{aligned} & Attn_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})=\frac{\exp \left(S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)}{\sum_{n=0}^m \exp \left(S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)} \leq \epsilon, \quad \forall n<m-C_0, \\ & L_h:=\frac{2\left\|W_{Q_h} W_{K_h}\right\|_2\left(\|\boldsymbol{\gamma}\|^2+\|\boldsymbol{b}\|^2\right)-\log (\epsilon)}{l_h} . \end{aligned} \end{aligned}$$ <p>这里，$W_{Q_h}$和$W_{K_h}$是第h个注意力头的查询和键矩阵，$\gamma$和$b$是注意力层之前的LayerNorm层的权重和偏置（对于RMSNorm【Zhang & Sennrich, 2019】，b=0），$||\cdot||_2$表示矩阵的l2范数。$L_h$可以被看作是该头的“视野范围”。详细证明见附录A.1。

动态调整KV缓存。定理1表明，当$q_m$和$k_n$之间的距离超过$C_0$时，这两个token之间的注意力权重将低于$\epsilon$。当$\epsilon$足够小（例如，0.1%）时，远程token对最终输出的影响极小，因此可以被丢弃。基于这一原则，ALiBi模型可以为每个头动态调整KV缓存的大小。我们首先计算有效注意力范围$L_h$，然后只在KV缓存中保留最近的$L_h$个token，因为任何距离超过$L_h$的token所施加的注意力权重不超过$\epsilon$，我们可以安全地丢弃它们以进行压缩。因此，对于ALiBi模型，检索头是那些具有较大$L_h$的头，而非检索头则具有较小的注意力视野$L_h$。

3.2 针对RoPE模型的RAZORATTENTION

RoPE模型的注意力分数计算。对于RoPE模型，每个注意力头根据以下公式计算注意力分数：

$$S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})=\boldsymbol{q}_{m} \boldsymbol{k}_{n}^{\top}, \quad \boldsymbol{q}_{m}=\mathcal{R}_{m} \boldsymbol{q}, \quad \boldsymbol{k}_{n}=\mathcal{R}_{n} \boldsymbol{k}$$

其中，$q_m$和$k_n$是经过旋转变换后的查询和键状态，$R_m$和$R_n$是位置m和n的旋转矩阵（详见【Su et al., 2023, Roformer: Enhanced Transformer with Rotary Position Embedding, 2023】）。尽管RoPE嵌入本身并不意味着长距离衰减的注意力，但我们的经验发现表明，只有大约15%的头（我们称之为“检索头”）能够有效利用长距离信息来保证模型的推理性能，而其余的头主要关注局部上下文。

检索头的重要性。如表1所示，当所有头都设置为只关注局部上下文时（保护0%的头），准确率平均显著下降7.3%；当随机选择15%的头关注长距离信息（随机保护15%的头）而其余头关注局部上下文时，准确率平均显著下降7.1%；然而，当选择检索头关注长距离信息（保护15%的检索头）而其余头关注局部上下文时，准确率与基线相当，仅下降0.3%。

补偿token机制。基于上述发现，我们直接减少了所有非检索头的KV缓存。如表1所示，模型的性能大部分得以保留。然而，仍然存在一个显著的准确率差距，表明一些信息仍在丢失。此外，即使我们保护了检索头的KV缓存，在“大海捞针”测试中的结果也显示出明显的性能下降（见图8中的消融结果）。为了进一步提升性能，我们设计了一种轻量级且有效的方法，将被丢弃token中的信息压缩成一个“补偿token”。该补偿token定义如下：

$$\hat{\boldsymbol{k}}=\frac{1}{N_{d}} \sum_{m \in\{\mathcal{D}\}} \boldsymbol{k}_{m}, \quad \hat{\boldsymbol{v}}=\frac{1}{N_{d}} \sum_{m \in\{\mathcal{D}\}} \boldsymbol{v}_{m} .$$

这里，$\hat{k}, \hat{v}$是丢弃的KV缓存的补偿token，$\{D\}$包含被丢弃token的索引，$N_d$是被丢弃token的数量。之后，我们丢弃这些token，并用补偿token $\hat{k}$和$\hat{v}$来增强KV缓存，其中$\{K, V\}$是经过旋转变换后剩余token的KV缓存。记压缩后的KV缓存为$\{K, \hat{k}\}$和$\{V, \hat{v}\}$，当前token的注意力输出遵循以下公式：

$$ \operatorname{Attn}(\boldsymbol{q}_m, \{K, \hat{\boldsymbol{k}}\}, \{V, \hat{\boldsymbol{v}}\}) = \frac{N_d \exp\left(\boldsymbol{q}_m \hat{\boldsymbol{k}}^{\intercal}\right)\hat{\boldsymbol{v}} + \sum_{n \notin \{\mathcal{D}\}} \exp\left(\boldsymbol{q}_m \boldsymbol{k}_n^{\intercal}\right)\boldsymbol{v}_n}{N_d \exp\left(\boldsymbol{q}_m \hat{\boldsymbol{k}}^{\intercal}\right) + \sum_{n \notin \{\mathcal{D}\}} \exp\left(\boldsymbol{q}_m \boldsymbol{k}_n^{\intercal}\right)} $$

RazorAttention图示。图3(a)中我们提供了RoPE模型上RazorAttention的一个说明性示例。通过补偿token，准确率得到进一步提升，使得RazorAttention即使在非检索头中丢弃70%的KV缓存也几乎无损。下面我们介绍如何确定检索头组。

图 3: 图3(a)展示了RazorAttention如何压缩KV缓存。对于检索头，我们维护一个完整的缓存以保留所有token的信息。对于非检索头，我们直接丢弃远程token，并将丢弃的token压缩成一个补偿token，其KV缓存表示为$\{\hat{k}, \hat{v}\}$。图3(b)提供了一个回声头（echo head）和归纳头（induction head）的说明性示例。当前token是“B”，生成的token是“C”。在这种情况下，回声头将主要关注token“B”，而归纳头主要关注之前上下文中的token“C”。

3.3 检索头的识别

RoPE模型的检索头识别。对于ALiBi模型，注意力范围可以直接通过公式2确定，并相应地丢弃KV缓存。然而，对于RoPE模型，检索头的识别需要一种更复杂的方式。我们的研究揭示，有两组头在处理长上下文中至关重要，因此它们都应被包含为检索头，如下所述：
* 回声头（Echo head）：倾向于关注与当前token相同的先前token（称为回声token）。
* 归纳头（Induction head）：倾向于关注紧随当前token之后的先前token（即归纳token）。基本上，它关注的是之前上下文中也存在的下一个token。

识别过程。图3(b)中我们展示了一个解释回声头和归纳头的说明性示例。为了识别检索头，我们生成K个（例如，K = 2500）随机token，将这些token重复4次，然后将其用作模型的输入。这种设计最小化了token之间的语义依赖，从而可以更清晰地观察回声头和归纳头的行为。

检索头选择标准。随后，我们计算了所有头中所有词的回声分数（对回声token的注意力权重）和归纳分数（对归纳token的注意力权重）。检索头的选择包括归纳分数最高的top-14%的注意力头和回声分数最高的top-1%的注意力头（见表2）。需要注意的是，尽管我们使用的回声头远少于归纳头，但我们的研究表明，这两种头对于LLM的检索性能都至关重要（见第4.2节的消融结果）。

RazorAttention算法。在识别出检索头后，我们在此介绍用于RoPE模型的RazorAttention算法1。

算法1：针对RoPE模型的RazorAttention

输入: 非检索头集合$\{H\}$，原始KV缓存（旋转变换后）$\{K, V\}$，压缩比$C$，压缩阈值$S_0$，初始token（sink token）数量$N_0$。
1. 对于非检索头 $h \in \{H\}$ 执行:
2.     计算缓冲区长度 $L_h = \max(S_0, N/C)$，其中$N$是该头中的token数量。
3.     仅保留靠近输出的最近$L_h$个token和最初的$N_0$个初始token，丢弃其余token并根据公式4将它们压缩成一个补偿token。
4. 结束循环
5. 非检索头根据公式5计算注意力，而检索头则遵循原始的注意力计算方式。

A4 实验环境

• 模型:

  • Qwen系列: Qwen1.5-7B, Qwen1.5-72B【Bai et al., 2023a】
  • Llama系列: Llama2-7B-80K, Llama2-13B-64K【Touvron et al., 2023; Fu et al., 2024】, Llama3-8B【AI@Meta, 2024】, Llama3.1-8B-Instruct, Llama 3.1-70B
  • Baichuan系列: Baichuan2-13B【Baichuan, 2023】
  • GLM系列: GLM-9B-1M【GLM et al., 2024】

• 数据集与基准测试:

  • LongBench【Bai et al., 2023b】: 用于评估LLM长上下文相关能力的综合基准。
  • Needle In A Haystack【gkamradt, 2023】: 用于测试模型在长文本中精确回忆信息的能力。
  • RULER【Hsieh et al., 2024】和InfiniteBench【Zhang et al., 2024】: 用于额外的准确性评估。

• 硬件配置:

  • 主要实验平台: NVIDIA GeForce RTX 4090 (24GB)。
  • 效率评估平台: 8颗昇腾（Ascend） 910B NPU。

• 软件与基线:

  • 对比方法: H2O【Zhang et al., 2023】, StreamingLLM【Xiao et al., 2024】, SnapKV【Li et al., 2024】。
  • 实验设置: 除非特别说明，RazorAttention使用表2中的超参数配置，该配置在长上下文输入下可实现3.125倍的KV缓存压缩。
    • 缓冲区长度: $\max(4000, N/5)$
    • 归纳头保护: top 14%
    • 回声头保护: top 1%
    • 初始token数: 4

A4 实验结果

4.1 准确性评估

4.1.1 LONGBENCH评估

实验内容：在LongBench基准上，我们评估了RazorAttention在不同模型上的性能，包括Qwen1.5-7B/72B（RoPE模型）、Llama3-8B（GQA模型）和Baichuan2-13B（ALiBi模型）。
实验结果：如表3所示，与StreamingLLM和H2O相比，RazorAttention在所有模型上均取得了优越的性能。即使在3倍压缩比下，其性能也与未压缩的基线模型相当。对于GQA模型（Llama3-8B），我们将一个组内的所有注意力头都视为检索头（只要组内有一个头满足归纳或回声属性），结果表明RazorAttention同样有效。

4.1.2 “大海捞针”评估

实验内容：我们使用上下文长度为80K的Llama2-7B-80K模型在“大海捞针”任务上进行评估。
实验结果：如图4所示，与在长输入下性能严重下降的H2O不同，RazorAttention能够准确地回忆起被查询的信息，这有力地证明了该方法保留了原始上下文中的所有语义信息。
超长序列测试：如图5所示，我们使用GLM-9B-1M模型在1024K序列长度的“大海捞针”任务上验证了RazorAttention对超长序列的支持能力。通过保留更多的检索头（28%的归纳头 + 2%的回声头），RazorAttention实现了几乎无损的准确率。
与SnapKV对比：如图6所示，在Llama 3.1-70B模型上，我们对比了RazorAttention和SnapKV。SnapKV在查询已知（query-aware）的情况下表现出色，但在查询未知（query-agnostic）的场景下性能显著下降。相比之下，RazorAttention凭借其免训练设计和逐头稀疏模式，在各种查询类型下都表现出卓越的鲁棒性和一致的高准确率。

图 4: RazorAttention与其他压缩算法在Llama2-7b-80K上的“大海捞针”性能比较。请注意，H2O与FlashAttention不兼容，因此在较长序列上测试时出现OOM错误，且其性能在这种情况下已变得不可用。

表 3: RazorAttention与其他压缩算法在不同LLM上于LongBench的性能比较。请注意，Llama3-8B-Instruct在TREC和LSHT上的性能不适用（接近0），因此我们未包含其在Llama3-8B上的结果。

图 5: GLM-9B-1M在1024K“大海捞针”任务上的表现。

图 6: RazorAttention和SnapKV在Llama 3.1-70B模型上于“大海捞针”任务准确率的比较。SnapKV在查询感知和查询未知设置下显示出显著的准确率差异，如(a)和(b)所示，而RazorAttention则始终保持较高准确率，如(c)所示。

4.2 消融研究

4.2.1 回声头的重要性

实验内容：我们评估了仅保留1%的回声头对模型性能的影响。
实验结果：如图7所示，尽管数量很少，但这组头在长上下文信息检索中至关重要。移除回声头会显著损害检索性能。一个可能的解释是归纳头的运作依赖于回声头的存在，正如【Olsson et al., 2022】所讨论的。

4.2.2 归纳头的数量

实验内容：我们测试了不同数量的归纳头对RazorAttention准确率的影响。
实验结果：如表4所示，随着归纳头数量的增加，准确率持续提升。我们选择保留14%的归纳头，以在压缩率和模型性能之间达到最佳平衡。

4.2.3 补偿token的重要性

实验内容：我们对比了使用和不使用补偿token的RazorAttention性能。
实验结果：如图8所示，补偿token对于RazorAttention的性能至关重要。它成功地压缩了被丢弃token的大部分信息，从而即使在显著的KV缓存压缩下也能保持高准确率。

图 7: 增加1%的回声头可以显著增强RazorAttention在Llama2-7B-80k上的检索性能。

图 8: 补偿token对于恢复由截断KV缓存引入的信息损失至关重要。

补充实验结果

RULER和InfiniteBench评估：我们在RULER和InfiniteBench数据集上使用LLaMA3.1-8B-Instruct模型进行了额外评估。结果（见附录表5和表6）显示，RazorAttention在这些基准测试中也表现出色，性能与原始模型相当。
效率评估：我们在8个Ascend 910B NPU上评估了GLM-9B-1M模型的解码延迟和吞吐量。如附录表7所示，由于KV缓存压缩后可以使用更大的批处理大小，在输入长度为128k和256k时，最大吞吐量分别提升了64.2%和71.81%。

A5 结论

本文提出了一种新颖的KV缓存压缩算法——RazorAttention，它成功地为使用RoPE或ALiBi嵌入的模型实现了3倍的压缩率。与以往基于重要性的token丢弃方法不可避免地丢弃语义信息不同，RazorAttention在检索头中保留了所有语义信息。我们证明了在非检索头中，远程token可以被有效地压缩成补偿token。此外，我们的逐头剪枝标准与FlashAttention完全兼容，使得RazorAttention成为一种即插即用的压缩方法，能够在扩展上下文下加速LLM的推理。我们的实验表明，RazorAttention可以达到与原始模型相当的性能，并在准确性和效率上超越了以往的方法。

A7 补充细节

6 局限性

未来工作的挑战。尽管我们的工作取得了成功，但仍存在一些局限性。第一个问题是，为什么LLM中的注意力头表现如此不同，以及检索头在长输入下是如何运作的。第二个挑战在于实现更高的压缩率。虽然我们已成功将KV缓存减少了70%，但我们相信这个数字可以进一步提高。此外，尽管我们已在多个模型上测试了我们的算法，但其他模型上的最佳配置可能会有所不同，这意味着在不同情况下我们可能需要更多或更少的检索头。这些主题都非常重要，我们将在未来的工作中继续研究。

A6 附录

A.1 定理1的证明

证明思路概述。我们首先给出查询和键乘积的一个上界，然后证明当位置偏差显著大于该上界时，注意力权重会衰减到零。
第一步：qk乘积上界。由于我们有$q = W_{Q_h}x$和$k = W_{K_h}x$，其中x是Attention块的输入，这导致：

$$\boldsymbol{q}\boldsymbol{k}^{\mathsf{T}}=\boldsymbol{x}W_{Q_{h}}W_{K_{h}}\boldsymbol{x}^{\mathsf{T}}\le\|W_{Q_{h}}W_{K_{h}}\|_{2}\|\boldsymbol{x}\|^{2}.$$

第二步：LayerNorm输出范数上界。由于x是在LayerNorm之后获得的，这意味着：

$$\begin{aligned} \begin{aligned} \boldsymbol{x} & = \boldsymbol{\gamma} \odot \frac{\hat{x} - \mu}{\sigma} + \boldsymbol{b}, \\ \mu & = \frac{1}{d} \sum_{i=1}^d \hat{x}_i, \quad \sigma = \frac{1}{d} \sum_{i=1}^d (\hat{x}_i - \mu)^2. \end{aligned} \end{aligned}$$

这里$\hat{x}$是LayerNorm的输入，d是其维度，$\hat{x}_i$是$\hat{x}$的第i维。上述方程导致：

$$\begin{aligned} \begin{aligned} \|\boldsymbol{x}\|^{2} & =\left\|\boldsymbol{\gamma} \odot \frac{\hat{x}-\mu}{\sigma}+\boldsymbol{b}\right\|^{2} \\ & \leq 2\left\|\boldsymbol{\gamma} \odot \frac{\hat{x}-\mu}{\sigma}\right\|^{2}+2\|\boldsymbol{b}\|^{2} \\ & \leq 2\|\boldsymbol{\gamma}\|^{2}+2\|\boldsymbol{b}\|^{2} . \end{aligned} \end{aligned}$$

第三步：组合上界。结合公式6和公式7，我们得到：

$$\boldsymbol{q}\boldsymbol{k}^{\intercal} \le \|W_{Q_{h}} W_{K_{h}}\|_{2} (2\|\boldsymbol{\gamma}\|^{2} + 2\|\boldsymbol{b}\|^{2})$$

第四步：注意力权重上界。为了给出注意力权重的上界，我们有：

$$\begin{aligned} \begin{aligned} \operatorname{Attn}_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k}) & =\frac{\exp \left(S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)}{\sum_{n=0}^{m} \exp \left(S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)} \\ & =\frac{\exp \left(\boldsymbol{q} \boldsymbol{k}^{\top}-l_{h}(m-n)\right)}{\sum_{n=0}^{m} \exp \left(S_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)} \\ & \leq \frac{\exp \left(\boldsymbol{q} \boldsymbol{k}^{\top}-l_{h}(m-n)\right)}{\exp \left(S_{n \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right)} \\ & \leq \frac{\exp \left(\boldsymbol{q} \boldsymbol{k}^{\top}-l_{h}(m-n)\right)}{\exp \left(\boldsymbol{q} \boldsymbol{q}^{\top}\right)} \\ & \leq \exp \left(\boldsymbol{q} \boldsymbol{k}^{\top}-l_{h}(m-n)\right) \\ & =\frac{\exp \left(\boldsymbol{q} \boldsymbol{k}^{\top}\right)}{\exp \left(l_{h}(m-n)\right)} . \end{aligned} \end{aligned}$$

第五步：推导距离条件。因此，为了确保$Attn_{m \to n}(q; k) \le \epsilon$，这等价于$\log(Attn_{m \to n}(q; k)) \le \log(\epsilon)$，我们需要：

$$\log \left(\operatorname{Attn}_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k})\right) \leq \boldsymbol{q} \boldsymbol{k}^{\top}-l_{h}(m-n) \leq \log (\epsilon)$$

第六步：最终结果。将公式8代入上述方程，我们得到：

$$\|W_{Q_{h}}W_{K_{h}}\|_{2}\left(2\|\boldsymbol{\gamma}\|^{2}+2\|\boldsymbol{b}\|^{2}\right)-l_{h}(m-n) \leq \log (\epsilon),$$

这给了我们：

$$m-n \geq \frac{2\left\|W_{Q_{h}} W_{K_{h}}\right\|_{2}\left(\|\boldsymbol{\gamma}\|^{2}+\|\boldsymbol{b}\|^{2}\right)-\log (\epsilon)}{l_{h}}.$$

在这种情况下，我们有：

$$\operatorname{Attn}_{m \rightarrow n}(\boldsymbol{q} ; \boldsymbol{k}) \leq \epsilon.$$

A.2 其他准确性评估

RULER和InfiniteBench评估。我们使用LLaMA3.1-8B-Instruct模型在RULER【Hsieh et al., 2024】和InfiniteBench【Zhang et al., 2024】数据集上对RazorAttention进行了额外评估，以证明其有效性。在我们的实验设置中，我们保留了30%的检索头（28%的归纳头和2%的回声头），并在序列超过4k个token时对非检索头实施了压缩机制：我们保留了大小为4的注意力窗口和覆盖序列长度20%的局部窗口，其余token被直接压缩。

表 5: RULER基准测试结果

表 6: InfiniteBench结果

A.3 效率评估

解码延迟和吞吐量评估。我们在8个Ascend 910B NPU上评估了GLM-9B-1M模型上RazorAttention的解码延迟和吞吐量，考虑了不同的预填充和解码输入长度，如表7所示。此外，由于KV缓存压缩后能够使用更大的批处理大小，我们在128k和256k的输入长度下，分别实现了64.2%和71.81%的最大吞吐量提升。

表 7: GLM-9B-1M模型在8个Ascend 910B NPU上的性能加速。