初识华为RazorAttention

最近学习华为的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。

参考：

【论文精读】RazorAttention如何优化模型部署-CSDN博客