【论文精读】RazorAttention如何优化模型部署

最新推荐文章于 2025-11-10 19:01:56 发布

原创最新推荐文章于 2025-11-10 19:01:56 发布 · 919 阅读

18 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

论文原文：https://openreview.net/pdf?id=tkiZQlL04w

最近学习华为的RazorAttention，水平有限，初步学习阅读论文，不展开讲算法，以下自我理解。

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

长上下文大型语言模型（LLM）在不同任务的自然语言处理方面具有显著的先进性。在LLM 模型的应用场景中，KV（Key-Value）缓存需要保存所有词元的Key与Value，以便节省后续解码时的计算复杂度，但是在不断增加的输入长度下的增长已经成为部署的主要瓶颈。

二、当前主流方式：

量化（Quantization）：量化方法通过降低存储每个数据点的精度来减少内存使用。使用较低的精度（如4位或8位）来表示数据。
窗口（Window）：这种方法通过限制缓存中存储的数据量来优化显存使用。例如，只缓存最近的或最重要的数据。
稀疏化（Sparse）：这种方法通过仅存储非零或重要的数据点来减少内存占用。它利用数据的稀疏性来优化存储。

上述的方法，能有效的减少KV缓存的占比，但是由于降低了数据的精度或者丢弃了部分的tokens，无疑会导致性能的下降

E.G. 输入：“这部电影很棒，但结局让人失望”，请分析人物情感。如果使用稀疏化优化性能，可能导致：

模型基于这个稀疏化的特征向量进行推理，它可能错误地将情感分类为正面。

因为它也许会错误的将“结局” 这个向量舍弃，错误的认为这个词语并不重要，这会导致“让人失望”这个关键token它也会错误舍弃。最终模型读取为：这部电影很棒。而它丢弃的token无疑是重要的。

“我们能找到一种在不丢失语义的情况下减少KV缓存大小的方法？”

三、Razor Attention的出现

3.1、Alibi模型的attention

在传统的Transformer模型中，位置嵌入通常是预先计算并添加到输入序列的每个token上的。然而，Alibi不采用这种方法。它是在计算注意力分数的过程中动态的考虑信息位置的。

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

论文中计算注意力头视野的公式为：

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

3.2、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%的存储。