初识华为 RazorAttention

论文原文：https://openreview.net/pdf?id=tkiZQlL04w最近学习华为的 RazorAttention，水平有限，根据论文做了初步了解。

1 背景：KV 缓存成为部署模型的主要瓶颈

长上下文大型语言模型（LLM）在不同任务的自然语言处理方面具有显著的先进性。在 LLM 模型的应用场景中，KV（Key-Value）缓存需要保存所有词元的 Key 与 Value，以便节省后续解码时的计算复杂度。然而，随着输入序列的增长，KV Cache 所占用的空间也会不断增加，总体达到 O(seq_len)的空间复杂度，导致显存消耗迅速增加（极端情况下 KV Cache 与权重的显存占比可达到 9:1），进而影响模型推理性能、限制其部署。

1.1 当前主流方式：

• 量化（Quantization）：量化方法通过降低存储每个数据点的精度来减少内存使用。使用较低的精度（如 4 位或 8 位）来表示数据。

• 窗口（Window）：这种方法通过限制缓存中存储的数据量来优化显存使用。例如，只缓存最近的或最重要的数据。

• 稀疏化（Sparse）：这种方法通过仅存储非零或重要的数据点来减少内存占用。它利用数据的稀疏性来优化存储。

目前业界流行的 KV Cache 压缩算法均为实时动态压缩，即实时地计算一些指标：注意力分数（与 FlashAttention 不兼容）或者 Topk 判断稀疏模式，对推理速度均有非常高的负面影响，在实际部署中不适用。

2 Razor Attention

基于 Attention 可解释性，提出了 RazorAttention 静态 KV Cache 压缩算法，是目前业界唯一的静态方法。RazorAttention 的动机来源于 Anthropic 2022 年的论文 In-context Learning and Induction Heads，该文章介绍了 Decoder-Only 架构 的语言大模型独有的上下文学习能力（In-Context Learning），并发现有一部分注意力头和模型的上下文学习能力强相关，文中称之为具有推断能力的 Induction Heads。在此基础上，本工作联想到大模型的长序列能力也是上下文学习能力的一个子集，并围绕此展开了研究。实验表明，RazorAttention 压缩了 70%的 KV Cache，并保持模型的长序列能力几近语义无损。

2.1 算法原理

业界最新的长序列 KV Cache 压缩算法大多数为在线动态压缩方法（实时计算⼀些指标），存在如下问题：

• 例如在线计算 Attention 得分，存在 Overhead，且与 FlashAttention 不兼容

• 最新技术在往轻量化重要性计算发展，但是依然存在 TopK 计算；随着序列长度增加，Top-K 计算的 overhead 将显著增加

• MInfer 等算法基于稀疏模式，要求在线执行稀疏计算

• 依赖当前输入的丢弃 token 无法还原，影响多轮对话等场景应用

文中先以基于 ALiBi 编码的模型为例，从理论上证明每个注意力头拥有不同的视野范围，因此可以通过动态调整每个头的 KV Cache 大小来实现压缩。具体来说，需要计算出每个头的有效视野 Lh，并只保留最近的 Lh 个 token 在 KV Cache 中，而将距离较远的 token 丢弃掉。这样可以保证在不损失准确率的情况下，大幅减少存储空间的需求。在此基础上，“不同注意力头拥有不同视野”的结论可以扩展到任意位置编码（如 RoPE）的语言大模型。本⼯作在 Induction Heads 基础上发现了 Echo Heads 同样对长文本任务起到关键作⽤，并统称为检索头（Retrieval Heads）。检索头通常能有效利用长距离信息，对输入拥有全局视野，其 KV Cache 也因此非常重要：

• Echo Head：关注前文中与当前 token 相同的 token

• Induction Head：关注前文中与当前 token 相同的下一个 token

在压缩算法设计上，RazorAttention 对检索头的 KV Cache 进行保护，确保上下文当中重要且主要信息不丢失，并直接减少了非检索头的 KV Cache 大小，以达到压缩的目的。在 Non-retrieval Head 上进行压缩，需要保留 Attention Sink+Local Attention。此外，本文提出了 Compensation token（补偿 token），对于丢弃掉的 KV Cache 以平均值形式保留在非检索头中，达到保护非检索头的局部视野的效果。基于 RoPE 位置编码的 KV Cache 压缩原理如下：

2.2 Alibi 模型的 attention

在传统的 Transformer 模型中，位置嵌入通常是预先计算并添加到输入序列的每个 token 上的。然而，Alibi 不采用这种方法。它是在计算注意力分数的过程中动态的考虑信息位置的。通俗来说，Alibi 类型的模型是使用相对位置来计算注意力分数的。论文中发现，在这种情况的“不同的注意力拥有不同的视野”这个重要结论：基于 Alibi 模型，从理论上证明每个注意力头拥有不同的视野范围，因此可以通过动态调整每个头的 KVcache 大小来实现压缩。具体来说，需要计算出每个头的有效视野 L.h，并只保留最近的 Lh 个 token 在 KVcache 中，而将距离较远的 token 丢掉。有点窗口化的意思，但是通过计算，在保证不损失准确率的情况下，大幅减少存储空间的需求。论文中计算注意力头视野的公式为：

可以看到，它设置了一个超参数，当前的词对较远的注意力分数低于这个设定的超参数的时候，就会将它舍弃，被认为是不需要的信息。只需要让每个注意力头加起来的“视野”能够覆盖所有信息。保证语义信息不丢失。

2.3 RoPE 模型 attention

在论文中表示，“不同注意力头拥有不同视野”的结论可以扩展到任意位置编码（如 ROPE）的语言大模型。本工作在 Induction Heads 基础上发现了 EchoHeads 同样对长文本任务起到关键作用，并统称为检索头（Retrieval Heads）。检索头通常能有效利用长距离信息，对输入拥有全局视野，其 KV Cache 也因此非常重要：

• Echo Head：关注前文中与当前 token 相同的 token

• Induction Head： 关注前文中与当前 token 相同的下一个 token

在压缩算法设计上，RazorAttention 对检索头的 KVCache 进行保护，确保上下文当中重要且主要信息不丢失，并直接减少了非检索头的 KV Cache 大小，以达到压缩的目的。在 Non-retrievalHead 上进行压缩的目的，需要保留 Attention Sink+Local Attention。此外，本文提出了 Compensation token（补偿 token），对于丢弃掉的 KV Cache 以平均值形式保留在非检索头中，达到保护非检索头的局部视野的效果。基于 RoPE 位置编码的 KV Cache 压缩原理如上图，将部分信息压缩。可以简单理解为，检索头作为小队长，负责统领大局；而非检索头作为专家，专注自己的这部分工作内容。从而做到性能几乎没有损失，而减少了 70%的存储。

3 总结

文中提出的 RazorAttention 是目前业界唯一的静态 KV Cache 压缩算法，一方面没有在线动态计算开销（既不依赖 Attention score，也不涉及在线 Topk 计算），类似 PTQ 仅需少量输入样本进行离线校准，便可轻量快捷的决定 Attention 稀疏模式；另一方面能与 FlashAttention 等主流融合算子兼容，无明显 overhead。

参考：https://blog.csdn.net/weixin_57667062/article/details/143898047