SALS: Sparse Attention in Latent Space for KV cache Compression

TL;DR

The paper tackles the KV-cache memory bottleneck in long-context LLM inference by showing that Rotary Position Embedding (RoPE) increases key-variance and effective rank, complicating naive low-rank KV-cache compression. It introduces SALS, a Sparse Attention in Latent Space framework that first projects pre-RoPE KV-cache into a compact latent space via a joint low-rank projection, then estimates attention scores in latent space to select a small subset of tokens for full reconstruction and RoPE-enabled sparse attention. A key contribution is offline calibration to learn the latent projection, along with a latent-space scoring mechanism that preserves most full-attention mass while drastically reducing data movement and computation. Empirical results across LLaMA2-7B-Chat, Mistral-7B, and RULER-128k with LLaMA3.1-8B-Instruct show SALS achieving 6.4-fold KV-cache compression and 5.7-fold speed-up in the attention operator on 4K sequences, plus notable end-to-end throughput improvements (1.4×–4.5×) over strong baselines. These findings demonstrate SALS as a practical, scalable solution for enabling long-context inference with reduced KV-cache overhead without sacrificing accuracy.

Abstract

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.

Figures (4)

• Figure 1: (a) Increasing inference time due to low-rank matrix reconstruction. Low-rank KV cache with overhead leads to longer inference time than the standard attention due to the reconstruction overhead; (b) A simplifed key vector example with changed PCA direction after applying RoPE
• Figure 2: High overalp rate of pre-RoPE across layers
• Figure 3: Overall architecture of SALS. Three stages are introduced with stage 1 for multi-head KV Cache compression, stage 2 for token selection in latent space and stage 3 for sparse attention.
• Figure 4: (a)--(b): Eigenvalue distributions of key covariance matrices in LLaMA-2-7B-Chat and Mistral-7B-Instruct-v0.2 before and after applying Rotary Position Embedding (RoPE). (c)--(d): Number of principal components required to retain 90% of the total variance across transformer layers, indicating changes in effective rank after RoPE.

Theorems & Definitions (2)

• Lemma 1
• proof