通过潜空间稀疏注意力压缩键值缓存,显著提升长文本推理速度。
SALS: Sparse Attention in Latent Space for KV cache Compression

- 将键值缓存投影到紧凑潜空间,用无位置编码的交互实现稀疏选择
- 4K序列下缓存压缩6.4倍,注意力算子提速5.7倍,端到端吞吐提升4.5倍
- 适合长上下文生成场景,尤其对部署资源受限的模型有实用价值
能够处理长上下文的大语言模型需求旺盛,但其推理受限于庞大的键值(KV)缓存和高内存带宽。已有研究发现KV缓存在隐藏维度具有低秩特性,具备压缩潜力。然而,现代LLM广泛采用旋转变换位置编码(RoPE),导致直接低秩压缩会严重损失精度或引入新瓶颈——需先重建完整缓存才能应用RoPE。本文提出两个关键洞察:一是RoPE作用于键向量会增加其方差,进而提高秩;二是键向量在潜空间中跨多数层保持稳定表征。基于此,我们提出潜空间稀疏注意力框架(SALS):通过低秩投影将KV缓存映射至紧凑潜空间,在该空间内进行无RoPE的查询-键交互并稀疏选择关键令牌。仅需重建少量重要令牌,避免全量重建开销。我们在两个大规模模型(LLaMA2-7b-chat 和 Mistral-7b)上全面评估,额外在RULER-128k基准测试中验证了LLaMA3.1-8B-Instruct的可扩展性。实验表明,SALS在保持竞争性准确率的同时达到当前最佳性能:4K序列下相比FlashAttention2实现6.4倍缓存压缩与5.7倍注意力算子加速;端到端吞吐方面,相比GPT-fast在4K和32K序列分别提升1.4倍和4.5倍。
原文摘要 · Abstract (English)
Large Language Models capable of handling extended contexts are in high demand, yet their inference remains challenging due to substantial Key-Value cache size and high memory bandwidth requirements. Previous research has demonstrated that KV cache exhibits low-rank characteristics within the hidden dimension, suggesting the potential for effective compression. However, due to the widely adopted Rotary Position Embedding mechanism in modern LLMs, naive low-rank compression suffers severe accuracy degradation or creates a new speed bottleneck, as the low-rank cache must first be reconstructed in order to apply RoPE. In this paper, we introduce two key insights: first, the application of RoPE to the key vectors increases their variance, which in turn results in a higher rank; second, after the key vectors are transformed into the latent space, they largely maintain their representation across most layers. Based on these insights, we propose the Sparse Attention in Latent Space framework. SALS projects the KV cache into a compact latent space via low-rank projection, and performs sparse token selection using RoPE-free query-key interactions in this space. By reconstructing only a small subset of important tokens, it avoids the overhead of full KV cache reconstruction. We comprehensively evaluate SALS on various tasks using two large-scale models: LLaMA2-7b-chat and Mistral-7b, and additionally verify its scalability on the RULER-128k benchmark with LLaMA3.1-8B-Instruct. Experimental results demonstrate that SALS achieves SOTA performance by maintaining competitive accuracy. Under different settings, SALS achieves 6.4-fold KV cache compression and 5.7-fold speed-up in the attention operator compared to FlashAttention2 on the 4K sequence. For the end-to-end throughput performance, we achieves 1.4-fold and 4.5-fold improvement compared to GPT-fast on 4k and 32K sequences, respectively.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。