arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.02599cs.LGcs.AI
收藏

基于RoPE对齐的KV缓存压缩(RAP)

RAP: KV-Cache Compression via RoPE-Aligned Pruning

Jihao Xin, Tian Lyu, David Keyes, Hatem Ltaief, Marco Canini

AI总结:

RAP通过RoPE对齐剪枝技术,有效压缩KV缓存、注意力参数和FLOPs,同时保持模型精度,显著降低推理延迟。

摘要:

大语言模型的长上下文推理日益受到KV缓存的内存和计算成本的限制。低秩分解通过将W近似为A*B来压缩KV投影,其中A生成潜在的KV状态,B可以被吸收进下游权重中。在现代基于RoPE的LLM中,这种吸收失败:RoPE迫使潜在的KV状态被重建为全维度,重新引入了显著的内存和计算开销。我们提出了RoPE对齐的剪枝(RAP),它剪枝整个RoPE对齐的列对,以保留RoPE的2x2旋转结构,恢复B的吸收,并消除重建。我们在LLaMA-3-8B和Mistral-7B上的评估表明,RAP能够同时减少KV缓存、注意力参数和FLOPs约20-30%,同时保持强大的准确性。值得注意的是,RAP将注意力延迟降低到基线的83%(预填)和77%(解码)。

英文摘要:

Long-context inference in large language models (LLMs) is bottlenecked by the memory and compute of the key-value (KV) cache. Structured pruning is a direct way to shrink it: dropping the least useful channels of the W_k, W_v projection weights to reduce the output KV dimensions. However, modern LLMs apply Rotary Position Embedding (RoPE) after the QK projections, which rotates feature dimensions in pairs. Therefore, removing individual channels breaks these pairs, corrupting RoPE's positional semantics and rendering the pruned model unusable. We propose RoPE-Aligned Pruning (RAP), which constrains the pruning granularity to RoPE-aligned pairs rather than individual channels: removing whole pairs to keep the rotation intact. Our evaluation across Llama, Mistral, and Qwen models from 3B to 14B shows that RAP preserves accuracy at 30% KV compression (retain ratio \r{ho} = 0.7), far outperforms RoPE-blind channel pruning, stays near the strongest low-rank method at lower attention cost, and composes with orthogonal methods such as quantization.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑