# SALS: Sparse Attention in Latent Space for KV cache Compression
**Authors**: Junlin Mu, Hantao Huang, Jihang Zhang, Minghui Yu, Tao Wang, Yidong Li
1]Beijing Jiaotong University 2]ByteDance Seed [*]Work done at ByteDance Seed [†]Corresponding authors
(October 28, 2025)
## Abstract
Large Language Models (LLMs) capable of handling extended contexts are in high demand, yet their inference remains challenging due to substantial Key-Value (KV) 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 (RoPE) 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 (SALS) 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. The source code will be publicly available in the future.
Tao Wang at , Hantao Huang at , Minghui Yu at , Yidong Li at
## 1 Introduction
The groundbreaking success of Large Language Models (LLMs), such as ChatGPT [15] and Claude [1], is transforming information retrieval in many areas. The exponential growth in LLM service requests is creating an unprecedented demand for inference optimization algorithms, especially for long-context applications. In particular, many research efforts [14, 20, 33, 28] show that the Key-Value cache (KV cache) acts as a major performance bottleneck in LLM serving. As the sequence length increases, the KV cache consumes a large portion of GPU device memory. This high resource requirement underscores the urgent need to mitigate the KV cache memory overhead and further improve attention efficiency.
To tackle the KV cache overhead challenge, previous works [4, 20] show that the KV cache exhibits low-rank characteristics within the hidden dimension, and thus can be effectively compressed in the latent space. Palu [4] further shows that there exists a clear trade-off between accuracy and reconstruction overhead for multi-head attention. Compressing each KV head separately reduces the overhead but at the cost of losing accuracy. However, compressing all heads together by low-rank projection improves accuracy by preserving the global information but results in significantly increased reconstruction cost. In this work, we show that sparsely selecting a small subset of tokens from the KV cache significantly reduces reconstruction error while preserving model accuracy.
Recently, Rotary Position Embedding (RoPE) [21] has been widely adopted in LLMs; it introduces sinusoidal positional information by multiplying query and key states with rotation matrices. This not only prevents the fusion of low rank weights into the query state, but also requires the reconstruction of key states from the latent space. In this work, we observe that the application of RoPE to the key vectors increases their variance, which in turn results in a higher rank. Therefore, the KV cache must be compressed before applying RoPE, which leads to high reconstruction complexity. We further observe that the key vectors in the latent space largely preserve their representation across most layers. Inspired by Double Sparse [29], we use these compressed key vectors as token selection guidance, so that only the selected tokens need to be reconstructed, which significantly reduces the reconstruction complexity.
In this paper, we propose the Sparse Attention in Latent Space (SALS) framework. SALS projects all attention heads into a shared single-head latent space via low-rank projection, for both pre-RoPE queries and pre-RoPE keys. The approximation attention scores are computed in the latent space and the top-k tokens are selected accordingly. These selected tokens are then reconstructed and applied with RoPE for the final attention computation. By reconstructing only a small subset of all tokens, it avoids the overhead of full KV cache reconstruction.
In summary, this paper makes the contributions as follow:
- We observe that applying RoPE to key vectors leads to increased variance, resulting in higher rank and lower compression rates. In , key vectors before RoPE can still maintain the representation for the critical token selection.
- We propose the Sparse Attention in Latent Space (SALS) framework to utilize the low-rank pre-ROPE KV cache compression and further to adopt the KV cache in latent space to select critical tokens. The sparse attention is then proposed with reduced KV data movement and computational complexity.
- 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 show that we can achieves much higher accuracy comparing to the low-rank based KV cache compression. Comparing to sparse attention works, under the similar benchmark accuracy, SALS achieves 6.4-fold KV cache compression and 5.7-fold speed-up in attention operator compared to FlashAttention2 on 4K sequence. For the end-to-end throughput performance, SALS achieves 1.4-fold and 4.5-fold improvement compared to GPT-fast on 4k and 32K sequences respectively.
Table 1: KV data movement, memory cost and complexity comparison of quantization, low rank and token sparse method that is with fixed, dynamic and local token selection strategy.
| Name | Methods | KV data movement | Memory size | Computation Complexity | Accuracy |
| --- | --- | --- | --- | --- | --- |
| Palu [4] | Low Rank | Median | Low | High | Low |
| Loki [20] | Dynamic + Low Rank | Low | Median | Median | Median |
| StreamingLLM [33] | Fixed pattern | Low | Median | Low | Low |
| Quest [22] | Dynamic | Low | High | Median | Median |
| DS [29] | Dynamic | Low | Median | Median | High |
| Hshare [25] | Dynamic | Low | Median | Median | High |
| SALS (Ours) | Dynamic+Low Rank | Low | Low | Midian | High |
## 2 Related Work
### 2.1 Attention in Latent Space
The low rank matrix decomposition such as singular value decomposition (SVD) to compress LLMs has been actively studied. Existing works such as ASVD [31], SVD-LLM [24] and CALDERA [18] are proposed to compress LLM weights. However, as the context sequence length and batch size increase, the KV cache size is getting larger than LLM weights, urging the need for KV cache compression. Eigen Attention [19] has been proposed to compress the KV cache after applying RoPE but suffers a relatively large accuracy loss. To improve the accuracy, the most naive way is to compress the pre-RoPE KV cache. However, as Palu [4] points out, this will greatly introduce additional computation for recovering the key vectors. Palu further proposes the grouped-head low-rank decomposition optimization to reduce the computation overhead. We agree with Palu’s observation but tackle the computation from the sparse perspective. By only selecting a subset of all the tokens (the critical tokens), the low-rank key vectors reconstruction complexity is greatly reduced. By integrating sparse attention, we observe that the KV cache size, data movement and computation complexity are all reduced as summarized in Table 1.
### 2.2 Sparse Attention
Early efforts to reduce attention complexity in transformers, such as Sparse Transformer [5], Reformer [12], and Longformer [3], primarily focused on training-based approaches to achieve sparse attention. In contrast, recent studies [17, 33, 22] have explored post-training sparse attention, leveraging the inherent sparsity of attention scores to identify the most salient tokens without retraining. For instance, StreamingLLM [26] identifies initial and recent tokens as critical, while others [33, 32] accumulate attention scores to locate important tokens in the KV cache. Quest [22] assesses KV cache page importance via min/max metrics, and Double-sparse [29] selects key KV tokens across important channels. However, most of these post-training approaches mainly reduce bandwidth and computation cost, leaving the KV cache uncompressed and thus a potential bottleneck [28, 20].
Natively Sparse Attention (NSA) [30] introduces a dynamic hierarchical sparse strategy combining coarse-grained token compression with fine-grained token selection. Unlike SALS, NSA requires training from scratch to learn a compressed attention module. In contrast, SALS can be calibrated post-training, performing top- $k$ token selection based on low-rank latent approximations rather than a trained sparse module. This allows SALS to achieve comparable sparsity behavior without retraining and with precise control over the KV-cache compression process.
Overall, our work addresses the limitation of prior sparse attention methods by storing and operating on a low-rank KV cache, then selectively reconstructing high-dimensional representations for exact sparse attention. This design preserves accuracy while substantially reducing memory and computation overhead.
<details>
<summary>x1.png Details</summary>
### Visual Description
## Bar Chart: Inference Time per Module Across Methods
### Overview
The bar chart compares inference times (in milliseconds) for different methods across varying input sizes (k = 8k, 16k, 32k, 64k, 128k). Each bar is segmented into four methods: TopK (only SALS), Recover, Softmax, and p=qk. Annotations highlight key components like "Low-rank reconstruction" and "SALS attention."
### Components/Axes
- **X-axis**: Input size (k) labeled as 8k, 16k, 32k, 64k, 128k.
- **Y-axis**: Inference Time (ms), ranging from 0 to 25.
- **Legend**:
- Purple: TopK (only SALS)
- Green: Softmax
- Red: y=pv
- Orange: p=qk
- **Annotations**:
- "Low-rank reconstruction" (arrow pointing to Recover method at 64k and 128k).
- "SALS attention" (arrow pointing to TopK at 16k).
- "Low-rank reconstructed attention" (arrow pointing to Recover at 32k).
### Detailed Analysis
- **TopK (only SALS)**:
- Inference time remains consistently low (~1-2 ms) across all k values.
- At 128k, the bar is ~1.5 ms.
- **Recover**:
- Dominates inference time at larger k values.
- At 64k: ~12 ms (with ~8 ms attributed to "Low-rank reconstruction").
- At 128k: ~25 ms (with ~15 ms attributed to "Low-rank reconstruction").
- **Softmax**:
- Minimal contribution (~0.5-1 ms) across all k values.
- **p=qk**:
- Moderate contribution (~2-3 ms) at 16k and 32k, increasing to ~5 ms at 64k and 128k.
### Key Observations
- **Recover method** has the highest inference time, especially at 64k and 128k, due to "Low-rank reconstruction" overheads.
- **TopK (only SALS)** is the most efficient, with minimal time increases as k grows.
- "SALS attention" contributes ~1-2 ms to TopK at 16k.
### Interpretation
The chart demonstrates that the **Recover method** introduces significant computational overhead, particularly for large k values, due to low-rank reconstruction. In contrast, **TopK (only SALS)** remains efficient, suggesting it is optimized for scalability. The annotations imply that "Low-rank reconstruction" is a critical bottleneck in the Recover method, while "SALS attention" adds minor overhead to TopK. This highlights trade-offs between model complexity and inference efficiency.
---
## Scatter Plot: Effect of RoPE on PCA Directions and Variance
### Overview
The scatter plot compares PCA directions and variance before and after applying RoPE (Rotary Positional Encoding). Blue points represent data before RoPE, and red points represent data after RoPE. Trend lines are fitted to each dataset.
### Components/Axes
- **X-axis**: X (PCA direction 1), ranging from -10 to 25.
- **Y-axis**: Y (PCA direction 2), ranging from -15 to 20.
- **Legend**:
- Blue: Before Applying RoPE
- Red: After Applying RoPE
- **Trend Lines**:
- Blue line (before RoPE) slopes upward (~0.5 slope).
- Red line (after RoPE) slopes downward (~-0.3 slope).
### Detailed Analysis
- **Before RoPE (Blue)**:
- Points are clustered around the origin, with a positive correlation (slope ~0.5).
- Variance is lower, with points concentrated in a narrow band.
- **After RoPE (Red)**:
- Points are more dispersed, with a negative correlation (slope ~-0.3).
- Variance increases, as red points spread further from the origin.
### Key Observations
- **RoPE application** reduces the positive correlation (blue trend line) and introduces a negative correlation (red trend line).
- **Variance increases** after RoPE, as red points are more spread out.
- The blue trend line is steeper, indicating stronger linear relationships before RoPE.
### Interpretation
The scatter plot suggests that **RoPE reduces variance** in PCA directions and alters the directionality of components. The negative slope after RoPE implies that positional encoding introduces rotational invariance, which may improve model generalization. However, the increased variance could indicate a trade-off between robustness and computational efficiency. The blue trend line’s steeper slope highlights that pre-RoPE data had stronger linear dependencies, which RoPE disrupts to enhance model flexibility.
</details>
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
## 3 Challenges on Latent Space Transformation
In this section, we first review the RoPE mechanism preventing the fusion of query and key states in latent space. We also show that after RoPE, the key vector variance increases as well as the rank. We then analyze the pre-RoPE representational capability of tokens in the latent space, focusing on their ability to preserve key semantic features.
### 3.1 Latent Space Transformation with RoPE
The introduction of the RoPE mechanism complicates the low-rank transformation into the latent space. We begin by considering the case without RoPE. Let $K∈ℝ^s× d$ denote the key matrix, where $s$ is the number of tokens and $d$ is the dimensionality of each head. By applying an orthonormal projection matrix $U∈ℝ^d× r$ , with $r\ll d$ , we can transforme origin Key matrix into latent space:
$$
\widetilde{K}=K U \tag{1}
$$
so that each original Key vector $k_t∈ℝ^d$ is mapped to latent space $\widetilde{k}_t∈ℝ^r$ leading to kv cache compression rate $d/r$ . This method works well on the basis of that
$$
UU^T≈I, QK^T≈QUU^TK^T \tag{2}
$$
However, once RoPE is introduced, the relative order between the RoPE rotation matrix $R$ and the low‑rank projector $UU^⊤$ becomes crucial. In Equation 3, we show both design choices in a single expression: (i) pre-‑RoPE transformation, where keys are first transformed and then rotated, and (ii) post-‑RoPE transformation, where keys are rotated first and transformed afterwards. For a query at position $i$ and a key at position $j$ , the attention score be approximated in two low‑rank forms:
$$
Q_i K_j^⊤ ≈ \underbrace{Q R_i\bigl(R^⊤_j UU^⊤\bigr) K^⊤}_pre‑-RoPE or \underbrace{Q R_i\bigl(UU^⊤ R^⊤_j\bigr) K^⊤}_post‑-RoPE, \tag{3}
$$
where $R_i$ and $R_j$ are the RoPE rotation matrices for positions $i$ and $j$ , respectively.
Ideally, we can adopt the post-RoPE transformation to cache $\widetilde{K}=KR_jU$ and transform the rotated query into a latent space $\widetilde{Q}=QR_iU$ to avoid the reconstruction complexity. However, as we will discuss later, applying RoPE will rotate the key vectors with larger variance, making them difficult to approximate using a low-rank matrix. On the other hand, in the pre‑-RoPE transformation, caching $\widetilde{K}$ and reconstructing the full-‑rank keys as $\widetilde{K}U^⊤$ will result in substantial overhead.
Variance amplification for post–RoPE rotation. Previous works [4, 20] observe that the key cache exhibits low-rank characteristics with a relatively small rank to retain approximately 90% of the energy. As such, we can use a single projection weight $U$ as shown in Equation 1 to compress the key cache. However, we find that this phenomenon only holds for the key cache prior to applying RoPE. From Figure 1 (b), we observe that applying RoPE to a set of key vectors pushes the data points outward and rotates their principal axes. As shown in Figure 1 (b), the principal component is rotates away from its original direction and the points become more scattered with two main principal components. This indicates the increased rank due to applying RoPE. Moreover, the principal components rotate based on the token position, suggesting that a single fixed projection matrix may no longer approximate all of these rotated subspaces. Therefore, a pre–RoPE key for latent space transformation is preferable for maintaining accuracy.
Compute overhead during pre–RoPE reconstruction From Figure 1 (a), we compare full attention with pre–RoPE low–rank compression across sequence lengths ranging from 1 K to 32 K tokens. While low–rank compression can reduce memory size, the time spent reconstructing low-rank keys/values from their rank- $r$ form and applying the RoPE rotation grows linearly with sequence length and soon dominates the total attention runtime at 32 K tokens. This aligns with our analysis of Equation 3, which intruoduces a trade-off of memory savings and computational cost.
### 3.2 Token Representation in Latent Space
<details>
<summary>x2.png Details</summary>
### Visual Description
## Line Chart: Overlap Score Across Layers
### Overview
The chart visualizes the overlap scores (in percentage) of three language models (LLaMA2-7B-Chat, LLaMA3-8B, Mistral-7B-v0.2) across 30 layers. Scores range from 0% to 100%, with distinct trends for each model.
### Components/Axes
- **X-axis (Layer)**: Discrete values from 0 to 30, labeled "Layer."
- **Y-axis (Overlap Score %)**: Continuous scale from 0% to 100%, labeled "Overlap Score (%)".
- **Legend**: Located at the bottom-right corner, mapping:
- Blue circles: LLaMA2-7B-Chat
- Orange squares: LLaMA3-8B
- Green triangles: Mistral-7B-v0.2
### Detailed Analysis
1. **LLaMA2-7B-Chat (Blue Circles)**:
- **Layer 0**: ~17% overlap (lowest initial score).
- **Layer 1**: Sharp rise to ~95%.
- **Layers 2–30**: Stabilizes between ~90–95%, with minor fluctuations (~92–94%).
- **Layer 30**: ~84% (notable drop from peak).
2. **LLaMA3-8B (Orange Squares)**:
- **Layer 0**: ~52% overlap.
- **Layer 1**: Peaks at ~99%.
- **Layers 2–29**: Maintains ~98–99% overlap.
- **Layer 30**: ~92% (slight decline from peak).
3. **Mistral-7B-v0.2 (Green Triangles)**:
- **Layer 0**: ~36% overlap.
- **Layer 1**: Rises to ~99%.
- **Layers 2–29**: Fluctuates between ~95–99%.
- **Layer 30**: ~80% (significant drop from peak).
### Key Observations
- **Initial Layer Disparity**: All models show low overlap at Layer 0, with LLaMA2-7B-Chat having the lowest (~17%) and LLaMA3-8B the highest (~52%).
- **Convergence**: By Layer 1, all models achieve near-100% overlap, suggesting architectural or training similarities in early layers.
- **Divergence at Final Layer**: Mistral-7B-v0.2 experiences the largest drop (~19% from Layer 29 to 30), while LLaMA3-8B retains the highest final score (~92%).
- **Stability**: LLaMA3-8B demonstrates the most consistent performance across layers, maintaining ~98–99% overlap.
### Interpretation
The data suggests that while all three models exhibit high layer overlap after the first layer, their architectural or training differences manifest in the final layers. LLaMA3-8B’s sustained high overlap indicates robust consistency, whereas Mistral-7B-v0.2’s sharp decline at Layer 30 may reflect divergent design choices in later layers. LLaMA2-7B-Chat’s rapid convergence to ~95% overlap by Layer 1 highlights its efficiency in aligning with other models early in processing. These trends could inform model selection based on layer-specific performance requirements.
</details>
Figure 2: High overalp rate of pre-RoPE across layers
To quantify the token representations of the query and key in the latent space, we first compute the latent‑-space attention vector $\widetilde{p}=\widetilde{Q}\widetilde{K},$ where in the pre‑-RoPE setting, $\widetilde{Q}=QU.$ Let $C=\{i_1,i_2,…,i_N_{c}\}⊆\{1,\dots,s\}, N_c\ll s,$ denote the index set corresponding to the top‑- $N_c$ entries of $\widetilde{p}$ . The overlap score (OS) measures the token representation in the latent space, which is fraction of the full attention mass that these indices capture. $OS=∑_i∈Cp_i/∑_i=1^sp_i$ , where $p_i$ denotes the $i$ -th entry of the full attention distribution $p∈R^s$ .
Figure 2 presents the overlap score for LLaMA–7B‑-chat, LLaMA–8B and Mistral-‑7B‑-v0.2. We observe that the average overlap score mostly remains above $90\$ for layers 2–-29 but drops below $50\$ in layers 0, 1. A similar observation has been reported for hybrid RoPE attention design [27]. This finding indicates that, even when RoPE is completely omitted, latent space tokens can still preserve almost all of the full attention scores across the majority of layers.
## 4 Sparse Attention in Latent Space Framework
<details>
<summary>x3.png Details</summary>
### Visual Description
## Diagram: Technical Architecture of a Sparse Attention Mechanism
### Overview
The diagram illustrates a multi-stage technical architecture for optimizing attention mechanisms, likely in a transformer-based model. It emphasizes key-value (KV) compression, token selection, and sparse attention operations. The process flows from left to right, with distinct stages and color-coded components.
### Components/Axes
- **Color Coding**:
- **Orange**: Pre-rope Query, Post-rope Query, Approx Attention Score, Fused Recover Kernel
- **Blue**: Pre-rope Key, Latent Key Cache, Critical Post-rope Key
- **Green**: Key covariance matrix, EVD, Reconstruct & Rope
- **Red**: Approx Attention Score (highlighted)
- **Key Elements**:
- **Inputs**: Pre-rope Query (1×(head_num×head_dim)), Pre-rope Key
- **Stage 1 (KV Compression)**:
- Offline Calibration (processes Pre-rope Keys into Key covariance matrix and EVD)
- Latent Key Cache (output of Offline Calibration)
- **Stage 2 (Token Selection)**:
- Approx Attention Score (computes attention scores between Query and Key)
- Top-k Index [2,7] (selects top-k indices for sparse attention)
- Transpose (reorients data for Stage 3)
- **Stage 3 (Sparse Attention)**:
- Fused Recover Kernel (combines Reconstruct & Rope operations)
- Outputs: Post-rope Query, Critical Post-rope Key
- **Final Output**: FA (likely Feature Aggregation or Final Attention)
### Detailed Analysis
1. **Stage 1: KV Compression**
- Pre-rope Keys undergo **Offline Calibration** to compute a **Key covariance matrix** and **EVD** (Eigenvalue Decomposition), stored in the **Latent Key Cache**.
- This stage reduces dimensionality or precomputes statistics for efficiency.
2. **Stage 2: Token Selection**
- The **Approx Attention Score** (orange) computes attention scores between the Query and Key.
- **Top-k Index [2,7]** selects sparse indices (e.g., 2nd and 7th tokens) for further processing.
- Data is **Transposed** to align with Stage 3 requirements.
3. **Stage 3: Sparse Attention**
- The **Fused Recover Kernel** (orange) performs **Reconstruct & Rope** operations on the selected tokens.
- Outputs include the **Post-rope Query** (orange) and **Critical Post-rope Key** (blue), which feed into the final **FA** block.
### Key Observations
- **Color Consistency**: Orange dominates Stage 1 and 3, blue tracks Key-related components, and green represents intermediate matrices.
- **Flow Direction**: Left-to-right progression mirrors typical transformer attention workflows but with added compression/optimization steps.
- **Sparse Attention**: The Top-k Index [2,7] suggests selective token retention, reducing computational load.
### Interpretation
This architecture optimizes attention mechanisms by:
1. **Precomputing Key Statistics** (Stage 1) to accelerate attention score calculations.
2. **Sparsifying Attention** (Stage 2) via Top-k selection, reducing quadratic complexity.
3. **Fusing Operations** (Stage 3) to combine reconstruction and rope operations, minimizing memory overhead.
The **Critical Post-rope Key** (blue) likely retains essential information after sparse processing, ensuring model accuracy despite compression. The **FA** block may aggregate features or finalize attention outputs. This design balances efficiency (via sparsity) and accuracy (via offline calibration), critical for large-scale models.
</details>
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.
### 4.1 SALS Framework
In this section, we introduce our proposed SALS, which leverages a low-rank projection matrix to transform multi-head pre-RoPE keys into a compact single-head latent space, thereby markedly shrinking the memory footprint. To avoid the prohibitive cost of reconstructing the entire key cache, SALS estimates attention scores directly in the latent space, selects the top-‑ $N_c$ most critical tokens, and reconstructs only their corresponding keys. RoPE is only applied to the reconstructed keys, which are then used to compute the exact attention weights. Figure 3 provides a high-‑level summary of the SALS pipeline. Here, we use the top-k index $[2,7]$ to indicate the selected tokens. The selected tokens are then reconstructed and reshaped to multi-head keys for sparse attention. The SALS framework is pipelined with three stages: KV cache compression to latent space, critical token selection in latent space and selective reconstruction from latent space for the later sparse attention.
### 4.2 Multi-head KV cache Compression to Latent Space
In the offline calibration process, we select a small calibration dataset from the pre-training corpus and collect its pre-RoPE key tensors. Let $K∈ℝ^s× nd$ denote the stacked multi‑head keys of the $s$ calibration sequences. We begin by computing the empirical covariance of the calibration keys: $C=K^⊤K$ . Applying an eigenvalue decomposition, $C=U\boldsymbol{Σ} U^⊤$ , and selecting the leading $r$ eigenvectors $U_r$ yields the optimal rank-‑ $r$ projection matrix for the latent space. Note that the key vectors $K∈ℝ^s× nd$ are reshaped by merging the head number with head dimension. We provide the rational using the following lemma.
**Lemma 1**
*For any column-orthonormal matrix $U∈ℝ^h× r$ satisfying $U^⊤U=I_r$ , we define the captured variance by $U$ as $E(U)$ . For the multi–head joint projection, all possible projection matrix set is defined as $U_r:=\{U∈ℝ^nd× r\midU^⊤U=I_r\}$ , and $B_r:=\{\bigl(U_1,…,U_n\bigr) | U_i∈ℝ^d× r^{\prime}, U_i^⊤U_i=I_r^\prime\},$ where $r^\prime=r/n$ , is all possible projection matrix set for per–head projection. For all situations, the optimal $U$ from multi–head joint projection matrix can capture more enery then per–head projection matrix:
$$
\max_U∈U_rE(U) ≥ \max_U∈B_rE(U)
$$*
* Proof*
Pick any $U=\bigl(U_1,\dots,U_n\bigr)∈B_r$ with $U_i^⊤U_i=I_r^\prime$ . Because the blocks occupy disjoint coordinate ranges, their product is block‑diagonal:
$$
U^⊤U=\bigl(U_1^⊤U_1,\dots,U_n^⊤U_n\bigr)=(I_r^\prime,\dots,I_r^\prime)=I_r.
$$
Hence $U$ is column‑-orthonormal and thus $U∈U_r$ ; consequently $B_r⊆U_r$ . Since every candidate in $B_r$ is feasible for the joint search $U_r$ , maximizing over the larger set cannot yield a smaller value. ∎
### 4.3 Critical Token Selection in Latent Space
As discussed in previous works [25, 29], token‑wise attention is inherently sparse. However, identifying the relevant (critical) tokens fast enough remains the main bottleneck of most sparse-‑attention schemes. Guided by the observation in Sec 3.2 that pre‑-RoPE latent keys already reveal the sparsity pattern of the middle layers, we forgo the expensive full‑-rank query–-key multiplication and instead work in the latent space.
More specifically, let $U_r∈ℝ^nd× r$ be the joint projector from Equation 1 and let $r^⋆\ll r$ for scoring. We project the query once, $\tilde{q}=U_r^⊤q∈ℝ^r,$ and keep only its leading coordinates $\tilde{q}_: r^⋆∈ℝ^r^{⋆}.$ Each cached key already stores its full latent vector $\tilde{k}_j=U_r^⊤k_j,$ from which we extract the first $r^⋆$ dimensions, $\tilde{k}_j,: r^⋆∈ℝ^r^{⋆}.$ The approximate attention score of token in position $j$ is a cheap inner product
$$
s_j = \tilde{q}_: r^⋆^⊤ \tilde{k}_j,: r^⋆
$$
Thus, critical tokens are identified directly from the existing key cache, without extra storage and at a fraction of the original compute.
Algorithm 1 Sparse Attention in Latent Space
0: Pre-RoPE query $q∈ℝ^nd$ , key $k∈ℝ^nd$ , value $v∈ℝ^nd$ , Projection matrix $U_r∈ℝ^nd× r$ , Low-rank key cache $\widetilde{K}∈ℝ^(S-1)× r$ , quantised value cache $\widehat{V}∈ℝ^(S-1)× nd$ , Sparsity budget $k$ , Approx latent rank $r^⋆$
0: Attended output $y∈ℝ^nd$
1: Function SparseAttention $(q,k,v,U_r,\widetilde{K},\widehat{V},k)$
2: $\tilde{q}←qU_r$ ; $\tilde{k}←kU_r$ # Project new token to $r$ -‑dim. latent space
3: $\widetilde{K}←(\widetilde{K},\tilde{k})$ ; $\widehat{V}←(\widehat{V},v)$
4: $p^\prime←\tilde{q}_r^⋆\widetilde{K}^⊤_r^⋆$ # Cheap similarity on top‑- $r^⋆$ dims
5: $C←(p^\prime,k)$ # Selectr top– $k$ critical token index
6: $K_C←\widetilde{K}_CU^⊤_r$ # Reconstruct critical latent key cache
7: $q^R←(q)$ ; $K^R_C←(K_C)$ # Standard flash attention (FA) computation
8: $p←\bigl(q^R{K^R_C}^⊤/√{d}\bigr)$
9: $y←p \widehat{V}_C$
10: return $y$
11: End Function
### 4.4 Selective Reconstruction from Latent Space for Sparse Attention
Sparse attention is defined by restricting the standard attention computation to a selected subset of keys and values. Given a query matrix $Q∈ℝ^n× d$ and a full set of key and value matrices of length $s$ , we identify a subset of salient indices
$$
C=\{i_1,i_2,\dots,i_N_{c}\}⊆\{1,\dots,s\}, N_c\ll s, \tag{4}
$$
where $N_c$ denotes the number of selected positions on the token sequence. The corresponding sub-matrices $K_C∈ℝ^N_c× n× d$ and $V_C∈ℝ^N_c× n× d$ are extracted based on the selected indices $C$ . Here, the salient indices $C$ are obtained via critical token selection in latent space. The sparse attention output is then computed by restricting the standard attention to the selected tokens:
$$
p_C=≤ft(\frac{QK_C^⊤}{√{d}}\right), y=p_CV_C. \tag{5}
$$
This method works well since the softmax operation normalizes the input distribution so that the output sums to 1 and only a few positions (or tokens) in $C$ dominate with large values. The rest scores after softmax closes to 0. The whole SALS algorithm is summarized in Algorithm 1.
### 4.5 Performance Discussion
As the roofline model [8] suggests, the performance of a GPU is bounded by either computation or memory bandwidth. In attention computation, performance is commonly bounded by memory bandwidth [7], which is primarily consumed by KV cache data movement. For a context of $s$ tokens and feature dimension $d$ , full attention needs to transfer $2sd$ elements of keys and values. On the other hand, for the SALS framework, critical token selection and sparse attention both involve data movement. Critical token selection requires $r^⋆\lld$ latent dimensions, so the query–key dot products read $sr^⋆$ elements. After selecting the top– $k$ tokens, SALS only needs to access low‑-rank key and value caches, each of size $kr$ . The second pass therefore moves $2kr$ elements.
Combining this two phases, the sparse variant moves $sr^⋆+2kr$ scalars, so its memory‑-bound speed‑-up is: $\frac{2sd}{sr^⋆+2kr}=\frac{1}{d_r^⋆/2 + d_r k_s}$ . Here $d_r^⋆=r^⋆/d, d_r=r/d, k_s=k/s$ denote the latent‑-approximated-score ratio, the low-‑rank ratio, and the token sparsity, respectively. We develop a Triton kernel that fuses token selection, reconstruction and RoPE rotation into a single pass. This fused reconstruct–RoPE kernel reduces memory traffic by $7.69×$ to $14.28×$ , depending on the chosen sparsity level and low-rank compression ratio, compared to the standard FlashAttention implementation.
## 5 Experiment
### 5.1 Setup
Table 2: Evaluation on GSM8K and CoQA datasets for LLaMA2–7B–chat
| Method | GSM8K (strict/flexible) $↑$ | CoQA $↑$ | Memory Access $↓$ | Comp. ratio $↓$ |
| --- | --- | --- | --- | --- |
| baseline | 0.2335 / 0.2335 | 0.5997 | 1.00 | 1.00 |
| KIVI-4 | 0.2297 / 0.2305 | 0.5992 | 0.31 | 0.31 |
| KIVI-2 | 0.2047 / 0.2047 | 0.6010 | 0.19 | 0.19 |
| Palu-30%(3bit) | 0.1713/0.1759 | 0.5938 | 0.14 | 0.14 |
| Palu-50%(3bit) | 0.0614/0.0879 | 0.5560 | 0.09 | 0.09 |
| SALS-25% | 0.2312 / 0.2343 | 0.5975 | 0.13 | 0.28 |
| SALS-12.5% | 0.2176 / 0.2229 | 0.6070 | 0.07 | 0.15 |
Models and tasks. We evaluate our SALS framework on two mainstream 7B chat models: LLaMA2–7B–Chat [23], that employs multi–head attention (MHA), and Mistral–7B–v0.2 [11], that uses grouped–query attention (GQA). To assess accuracy, we report results on the reasoning benchmark GSM8K [6], the conversational benchmark CoQA [16], and the 16 English subsets of LongBench [2] that probe long–context understanding. All experiments are conducted on a machine with Xeon(R) Platinum 8336C CPU, one GPU (ampere architecture), and 128G RAM.
Baselines. We compare SALS with several state of the art baseline in two directions. For KV–cache compression we select Palu (low–rank projection) [4] and KIVI (quantization) [14]. For token–sparse decoding we include Double Sparse [29], Hshare [25], and Loki [20]. For the KV–cache compression baselines, we evaluate Palu at every reported compression ratio, including its variants that apply the optional quantisation step. For KIVI we evaluate both official settings: 4 bit and 2 bit. For the token–sparse baselines, sparsity is introduced only during the decoding phase, and we enforce the same overall sparsity ratio as Hshare to ensure a fair comparison.
Compression setup and calibration. To obtain the latent projection matrix, we randomly sample 512 sequences of length 4096 from the C4 corpus [9] and compute it offline. We apply multi-head joint compression to the key cache at two ratios: $d_r=25\$ and $d_r=12.5\$ . For latent scoring, we simply set $r^⋆=0.5r$ for all models. Since value tensors are almost full rank and play a pivotal role in attention, we forgo low-rank projection for them and instead perform channel-wise group quantisation that mirrors the key-cache setting (4 bit at 25% and 2 bit at 12.5%).
Following KIVI, we adopt a mixed high–precision / low–precision scheme: tokens in the most recent window are compressed by only 50%, whereas all preceding tokens are compressed according to the target ratio of the experiment. The high–precision window is aligned with the sparsity window: when the sparse mechanism always selects the most recent $w$ tokens, the compression stage likewise keeps those same $w$ tokens in high precision. Based on the observations in Figure 2, we skip the sparsification for layers 0, 1, and 31 across all models to ensure more accurate compression and sparsity results.
Table 3: Evaluation on LongBench datasets
| Model | Method | Single- QA | Multi- QA | Summari- zation | Few-shot | Synthetic | Code | Avg | Memory Access $↓$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLaMA2-7B-chat | baseline | 24.50 | 22.18 | 23.64 | 62.80 | 6.50 | 56.29 | 32.65 | 1.00 |
| KIVI-4bit | 24.76 | 22.22 | 23.40 | 62.67 | 6.75 | 56.42 | 32.70 | 0.31 | |
| Palu-30% (3bit) | 22.36 | 22.01 | 22.66 | 60.39 | 7.50 | 50.53 | 30.91 | 0.13 | |
| SALS-25% | 24.37 | 22.50 | 23.40 | 63.01 | 6.50 | 53.80 | 32.26 | 0.11 | |
| KIVI-2bit | 23.11 | 21.40 | 22.67 | 62.81 | 5.75 | 55.84 | 31.93 | 0.19 | |
| Palu-50% (3bit) | 20.48 | 22.18 | 21.44 | 54.62 | 5.75 | 37.07 | 26.92 | 0.09 | |
| SALS-12.5% | 24.81 | 22.52 | 23.08 | 62.84 | 5.00 | 53.58 | 31.97 | 0.06 | |
| Mistral-7B-v0.2 | baseline | 36.43 | 29.65 | 28.06 | 66.72 | 44.87 | 52.96 | 43.12 | 1.00 |
| KIVI-4bit | 36.34 | 29.85 | 28.09 | 66.81 | 44.67 | 52.90 | 43.11 | 0.31 | |
| Palu-30% (3bit) | 29.48 | 36.40 | 27.20 | 65.73 | 53.19 | 40.77 | 34.74 | 0.13 | |
| SALS-25% | 36.61 | 29.92 | 27.57 | 66.63 | 43.54 | 52.49 | 42.79 | 0.11 | |
| KIVI-2bit | 35.34 | 28.63 | 27.72 | 66.77 | 39.68 | 52.63 | 41.80 | 0.19 | |
| Palu-50% (3bit) | 26.73 | 32.72 | 25.73 | 63.25 | 18.57 | 44.43 | 35.71 | 0.09 | |
| SALS-12.5% | 35.18 | 29.88 | 27.18 | 66.98 | 35.32 | 51.39 | 40.99 | 0.06 | |
### 5.2 KV‑cache Compression Comparison
Implementation details. For the dataset GSM8K and CoQA, we always keep the most recent $w=128$ tokens and decode the remaining context at a fixed sparsity of $1/4$ . For the dataset LongBench, LLaMA2–7B–Chat supports a 4 k context window, whereas Mistral–7B–v0.2 supports 32 k. To equalise the average sparsity at $1/8$ on this benchmark, we retain $N_c=512$ tokens for LLaMA2–7B–Chat and $N_c=1024$ tokens for Mistral–7B–v0.2. For LLaMA2–7B–Chat we follow the Hshare [25] configuration ( $x=16$ sink tokens, $y=432$ critical tokens, $z=64$ recent tokens) and simply double each count for Mistral–7B–v0.2.
Results. Tables 2 shows the evaluation on the benchmark (GSM8K, CoQA) on LLaMA2–7B–chat model. We can see that although Palu achieves the highest compression rate, the acurracy drop is non-negligible, especillay for the GSM8K dataset. SALS provides two setting results with 25% KV cache compression rate (SALS–25%) and 12.5% compression rate (SALS–12.5%). SALS–25% achieves the best accuracy compared to KIVI and Palu, with neligible loss compared to the baseline. Tables 3 compares task-level performance across six LongBench categories, alongside normalized memory access costs. The first three columns reflect core reasoning tasks, where SALS achieves consistently strong performance. Notably, SALS-12.5% retains competitive accuracy while reducing memory access to just 6% of the baseline. This indicates that SALS preserves informative tokens more effectively. The memory-access savings translate into practical latency improvements, as fewer KV memory lookups reduce attention computation overhead.
Table 4: Comparison of Token Sparse Methods on LLaMA2-7B-chat Using LongBench Tasks
| Method | Single- QA | Multi- QA | Summari- zation | Few-shot | Synthetic | Code | Avg | Memory Access $↓$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| baseline | 24.50 | 22.18 | 23.64 | 62.80 | 6.50 | 56.29 | 32.65 | 1.00 |
| Double Sparse | 24.78 | 22.72 | 24.70 | 61.84 | 4.17 | 51.66 | 31.64 | 0.16 |
| HShare | 24.59 | 22.18 | 24.54 | 61.69 | 4.74 | 53.26 | 31.83 | 0.14 |
| Loki | 24.57 | 18.09 | 24.37 | 63.43 | 4.75 | 56.49 | 31.95 | 0.19 |
| SALS-25% | 24.37 | 22.50 | 23.40 | 63.01 | 6.50 | 53.80 | 32.26 | 0.11 |
| SALS-12.5% | 24.81 | 22.52 | 23.08 | 62.84 | 5.00 | 53.58 | 31.97 | 0.06 |
### 5.3 Token‑-Sparse Comparison
Since SALS incorporates a sparsification stage, we also compare it with state of the art sparse decoders, such as Double Sparse, HShare and Loki in LongBench datasets.
Implementation details. We use the same sparse setting as Sec 5.2 for all methods, which have $x=16$ sink tokens, $y=432$ critical tokens and $z=64$ recent tokens, achieving sparsity on LLaMA2–7b–chat at $1/8$ .
Result. Table 4 shows that, on LongBench, SALS with 25% key compression matches the average accuracy of the baseline while issuing only about $11\$ of the baseline’s memory traffic. When the compression is tightened to 12.5%, the Memory–Access Ratio drops to just $5.8\$ , yet SALS still outperforms all competing sparse methods in accuracy. These results suggest that the latent‐-space scoring mechanism selects critical tokens more accurately than existing heuristics, enabling KV-‐cache compression and token sparsity to coexist without loss of accuracy.
Table 5: Performance comparison of baseline and SALS methods on the RULER dataset with Llama3.1-8B-Instruct (4k sequence length).
| Method | avg | S1 | S2 | MK1 | MK2 | MV | MQ | FEW | QA1 | QA2 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Baseline | 81.60 | 99.6 | 96.0 | 94.6 | 73.6 | 93.85 | 96.9 | 68.47 | 69.6 | 41.8 |
| SALS-25% | 80.81 | 99.6 | 95.2 | 94.2 | 65.8 | 93.2 | 96.4 | 71.53 | 70.2 | 41.2 |
| SALS-12.5% | 75.86 | 97.4 | 93.8 | 92.8 | 42.2 | 84.05 | 93.05 | 72.53 | 67.8 | 39.14 |
### 5.4 RULER Benchmark
Dataset. RULER [10] is a recently proposed long-context benchmark designed to evaluate the reasoning, retrieval, and compositional understanding capabilities of large language models across a wide range of input lengths and retrieval patterns. It decomposes the evaluation into fine-grained retrieval types, including single-key, multi-key, multi-value, and multi-query tasks, along with few-shot and question-answering (QA) subtasks. Compared with LongBench, RULER provides a more detailed breakdown of retrieval behaviors, making it well-suited for assessing token selection accuracy and KV-cache compression performance.
Implementation details. The experimental setup follows the same configuration as in Sec. 5.2. We evaluate all methods on a 128k-token context, applying an 8 $×$ sparsity ratio, i.e., selecting 16k active tokens during inference. The retained-to-pruned token ratio is identical to that used in the LongBench experiments, ensuring consistent compression levels and comparable evaluation metrics. All experiments are conducted on the LLaMA 3.1–8B–Instruct model with a 4k sequence length.
Result. Table 5 summarizes the performance of SALS on the RULER dataset. At a 25% compression ratio, SALS achieves nearly identical average accuracy to the baseline (80.81 vs. 81.60), demonstrating its high fidelity in token selection and KV-cache reconstruction. Performance remains consistent across most retrieval subtasks, including NIAH-Single-1/2, Multi-Value, and Multi-Query, indicating that latent-space sparsification preserves key contextual semantics even under substantial cache reduction.
When the compression is further tightened to 12.5%, accuracy degradation mainly occurs in NIAH-Single-1 and Multi-Key-2, where the retrieval dependency is strongest. Nevertheless, the model still achieves competitive overall performance (75.86 average) and stable scores on Few-shot and QA tasks. These results suggest that the proposed latent-space attention mechanism can maintain accurate token importance estimation even under extreme sparsity, enabling effective KV-cache compression with minimal accuracy loss across retrieval-oriented benchmarks.
Table 6: Attention Operator Latency (ms) across Methods and Input Configurations
| Config (ms) | Flash-attn [7] | Loki [20] | Double-sparse [29] | Hshare [25] | SALS-25% | SALS-12.5% |
| --- | --- | --- | --- | --- | --- | --- |
| bs=8, 1k | 0.230 | 0.248 | 0.138 | 0.134 | $0.409± 0.093$ | $0.357± 0.051$ |
| bs=8, 2k | 0.830 | 0.464 | 0.237 | 0.430 | $0.430± 0.018$ | $0.359± 0.012$ |
| bs=8, 4k | 1.630 | 1.102 | 0.724 | 0.576 | $0.530± 0.008$ | $0.439± 0.016$ |
| bs=16, 1k | 0.440 | 0.452 | 0.223 | 0.134 | $0.415± 0.008$ | $0.360± 0.011$ |
| bs=16, 2k | 1.630 | 0.864 | 0.434 | 0.246 | $0.552± 0.040$ | $0.444± 0.019$ |
| bs=16, 4k | 3.230 | 2.101 | 1.319 | 1.067 | $0.757± 0.026$ | $0.565± 0.013$ |
### 5.5 Efficiency Evaluation
Inference details. Following Hshare [25], all speed tests are performed on the PyTorch backend, with the Triton–based fused kernel as discussed in Sec 4.5. We evaluate two aspects: (i) self–attention latency and (ii) end–to–end generation speed–up. Self–attention latency is compared against FlashAttention–v2 [7], whereas end–to–end throughput is benchmarked against GPT–Fast [13]. Experiments use batch sizes of $8$ and $16$ and sequence lengths of $1$ k, $2$ k, and $4$ k tokens; the sparsity ratio is fixed at $1/8$ for all methods.
Result. Tables 6 reports our attention latency comparisons. Mean and varaince are reported using 1000 repetions. Thanks to the added token sparsity, SALS accelerates the stand‑alone attention operator by $7,46×$ speed‑up for batch size (bs) =8 and 4k sequence length (4k). Similar performance speed is also observed when batch size is 16. SALS will introduce some overhead for short sequences (e.g. 1k) but for longer sequences the speed-up is significantly better than the state-of-the-art sparse algorithm. As discussed in Section 4.5, reducing memory access will improve the attention performance.
Table 7: Long Context End-to-End Performance Throughput Comparison (token/second)
| Bsz | Seq (k) | GPT-Fast | SALS-25% | SALS-12.5% |
| --- | --- | --- | --- | --- |
| 8 | 4 | 118 | $154.1± 7.1$ | $163.5± 5.1$ |
| 8 | 8 | 70.6 | $128.6± 5.1$ | $150.1± 3.3$ |
| 8 | 16 | 38.3 | $102.5± 0.3$ | $122.7± 7.7$ |
| 8 | 32 | 19.8 | $67.97± 0.8$ | $89.47± 1.4$ |
| 4 | 64 | 9.2 | $32.92± 0.3$ | $42.96± 0.9$ |
Table 7 shows the end-to-end throughput speed-up of SALS comparing to the GPT-Fast results. Note that batch size is chosen 4 for sequence 64k to avoid GPU out of memory. Mean and varaince are reported using 50 repetions. We observe that when sequence is 8k, the speed-up of SALS-12.5% is $2.13×$ although the sparsity ratio is 1/8. This is due to the overhead of reconstructing the selected token. When the sequence is set 32k, the speed-up of SALS-12.5% is increasing and reaches $4.57×$ . Coupled with the strong accuracy of SALS, these results confirm that combining latent‑space KV‑cache compression with token sparsification offers an effective trade‑off between speed and accuracy.
## 6 Conclusion
This paper investigates the KV cache compression based on the low-rank characteristics in the hidden dimension. We further discuss the increasing variance of keys due to applying rotary position embedding and propose to use the pre-RoPE keys to select the critical tokens. We observe that after the key vectors are transformed into the latent space, they largely maintain their representation across most layers and thereby the salient tokens are selected with high accuracy. Based on these insights, we propose the Sparse Attention in Latent Space (SALS) framework. The SALS framwork achieves compressed KV cache to the latent space, selects the critical token in the latent space with significantly less computation load and performs sparse attention on the selected tokens. By reconstructing only a small subset of important tokens, it avoids the overhead of full KV cache reconstruction. Experimental results demonstrate that SALS achieves SOTA performance by maintaining competitive accuracy under different settings and achieving 6.4-fold KV cache compression and 5.7-fold speed-up in attention compared to FlashAttention2 on 4K sequences. For the end-to-end throughput performance, we achieve 1.4-fold and 4.5-fold improvement compared to GPT-fast on 4k and 32K sequences respectively.
## References
- Anthropic [2023] Anthropic. Claude. Large language model, May 2023. URL https://www.anthropic.com/claude. Accessed May 5, 2025.
- Bai et al. [2023] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508, 2023.
- Beltagy et al. [2020] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
- Chang et al. [2024] Chi-Chih Chang, Wei-Cheng Lin, Chien-Yu Lin, Chong-Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed S Abdelfattah, and Kai-Chiang Wu. Palu: Compressing kv-cache with low-rank projection. arXiv preprint arXiv:2407.21118, 2024.
- Child et al. [2019] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
- Cobbe et al. [2021] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
- Dao [2023] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
- Ding and Williams [2019] Nan Ding and Samuel Williams. An instruction roofline model for gpus. IEEE, 2019.
- Habernal et al. [2016] Ivan Habernal, Omnia Zayed, and Iryna Gurevych. C4corpus: Multilingual web-size corpus with free license. In Proceedings of the Tenth International Conference on Language Resources and Evaluation (LREC’16), pages 914–922, 2016.
- Hsieh et al. [2024] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models?, 2024. URL https://arxiv.org/abs/2404.06654.
- Jiang et al. [2024] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
- Kitaev et al. [2020] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
- Labs [2023] PyTorch Labs. gpt-fast: Simple and efficient pytorch-native transformer text generation in <1000 loc of python, 2023. URL https://github.com/pytorch-labs/gpt-fast. Accessed: 2025-05-14.
- Liu et al. [2024] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750, 2024.
- OpenAI [2023] OpenAI. Chatgpt. Large language model, May 2023. URL https://chat.openai.com/. Accessed May 5, 2025.
- Reddy et al. [2019] Siva Reddy, Danqi Chen, and Christopher D Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249–266, 2019.
- Ribar et al. [2023] Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficient llm inference. arXiv preprint arXiv:2312.04985, 2023.
- Saha et al. [2024] Rajarshi Saha, Naomi Sagan, Varun Srivastava, Andrea Goldsmith, and Mert Pilanci. Compressing large language models using low rank and low precision decomposition. Advances in Neural Information Processing Systems, 37:88981–89018, 2024.
- Saxena et al. [2024] Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, and Kaushik Roy. Eigen attention: Attention in low-rank space for kv cache compression. arXiv preprint arXiv:2408.05646, 2024.
- Singhania et al. [2024] Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arXiv preprint arXiv:2406.02542, 2024.
- Su et al. [2024] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
- Tang et al. [2024] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context llm inference, 2024.
- Touvron et al. [2023] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
- Wang et al. [2024] Xin Wang, Yu Zheng, Zhongwei Wan, and Mi Zhang. Svd-llm: Truncation-aware singular value decomposition for large language model compression. arXiv preprint arXiv:2403.07378, 2024.
- Wu et al. [2025] Huaijin Wu, Lianqiang Li, Hantao Huang, Tu Yi, Jihang Zhang, Minghui Yu, and Junchi Yan. Hshare: Fast llm decoding by hierarchical key-value sharing. In The Thirteenth International Conference on Learning Representations, 2025.
- Xiao et al. [2023] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023.
- Yang et al. [2025] Bowen Yang, Bharat Venkitesh, Dwarak Talupuru, Hangyu Lin, David Cairuz, Phil Blunsom, and Acyr Locatelli. Rope to nope and back again: A new hybrid attention strategy. arXiv preprint arXiv:2501.18795, 2025.
- Yang et al. [2024a] Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference. arXiv preprint arXiv:2405.12532, 2024a.
- Yang et al. [2024b] Shuo Yang, Ying Sheng, Joseph E Gonzalez, Ion Stoica, and Lianmin Zheng. Post-training sparse attention with double sparsity. arXiv preprint arXiv:2408.07092, 2024b.
- Yuan et al. [2025] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025. URL https://arxiv.org/abs/2502.11089.
- Yuan et al. [2023] Zhihang Yuan, Yuzhang Shang, Yue Song, Qiang Wu, Yan Yan, and Guangyu Sun. Asvd: Activation-aware singular value decomposition for compressing large language models. arXiv preprint arXiv:2312.05821, 2023.
- Zhang et al. [2024a] Tianyi Zhang, Jonah Yi, Zhaozhuo Xu, and Anshumali Shrivastava. Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization. arXiv preprint arXiv:2405.03917, 2024a.
- Zhang et al. [2024b] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024b.
## 7 Key Dimensianlity Analysis with RoPE
<details>
<summary>appendix_image/llama_eigen.png Details</summary>
### Visual Description
## Histogram: LLaMa-2-7B-Chat Eigenvalues Distribution of Keys
### Overview
The chart compares eigenvalue distributions of keys in LLaMa-2-7B-Chat before and after applying Rotary Positional Encoding (RoPE). Two histograms are overlaid: "post-RoPE" (green) and "pre-RoPE" (orange). A secondary inset provides a zoomed view of the dominant eigenvalue range (0–2).
---
### Components/Axes
- **X-axis (Eigenvalue)**: Logarithmic scale from 0 to 250, with gridlines at 0, 0.5, 1.0, 1.5, 2.0, 50, 100, 150, 200, 250.
- **Y-axis (Count)**: Logarithmic scale from 10⁰ to 10³.
- **Legend**:
- Green: post-RoPE
- Orange: pre-RoPE
- **Inset**: Focuses on eigenvalue range 0–2 with the same axes and legend.
---
### Detailed Analysis
#### Main Chart
- **pre-RoPE (orange)**:
- Dominates the lower eigenvalue range (0–0.5), with a peak count of ~10³ at eigenvalue 0.
- Counts drop sharply after 0.5, with sparse values beyond 1.0.
- A single outlier at ~250 with a count of ~10⁰.
- **post-RoPE (green)**:
- Lower counts than pre-RoPE in 0–0.5 (~10² at 0).
- More evenly distributed across 0–2, with gradual declines.
- Sparse values beyond 1.0, with a notable peak at ~1.5 (~10¹).
#### Inset (0–2 Eigenvalue Range)
- **pre-RoPE**:
- Tallest bars in 0–0.5 (~10³ at 0, ~10² at 0.5).
- Minimal presence beyond 1.0.
- **post-RoPE**:
- Lower counts than pre-RoPE in 0–0.5 (~10² at 0, ~10¹ at 0.5).
- More pronounced distribution in 0.5–2, peaking at ~1.5 (~10¹).
---
### Key Observations
1. **Concentration vs. Spread**:
- pre-RoPE eigenvalues are highly concentrated near 0, suggesting uniform key initialization.
- post-RoPE eigenvalues are more spread out, indicating positional encoding introduces variability.
2. **Dominant Eigenvalue Behavior**:
- pre-RoPE dominates eigenvalues <1.0, while post-RoPE shows non-negligible values up to 2.0.
3. **Outliers**:
- A single pre-RoPE eigenvalue at ~250 with negligible count (~10⁰), likely an artifact.
---
### Interpretation
The data suggests RoPE modifies eigenvalue distributions to enhance positional awareness:
- **pre-RoPE**: Uniform initialization (low eigenvalues) may limit positional differentiation.
- **post-RoPE**: Broader eigenvalue spread (up to 2.0) implies improved handling of positional relationships, though most dominant eigenvalues (<1.0) remain pre-RoPE-dominated. The presence of post-RoPE eigenvalues in the 0–2 range confirms partial effectiveness in redistributing key magnitudes. The outlier at 250 is insignificant but highlights potential edge cases in eigenvalue computation.
</details>
(a)
<details>
<summary>appendix_image/mistral_eigen.png Details</summary>
### Visual Description
## Chart: Eigenvalues Distribution of Keys for Mistral-7B-instruct-v0.2
### Overview
The chart compares the distribution of eigenvalues for keys in the Mistral-7B-instruct-v0.2 model before and after applying RoPE (Rotary Positional Encoding). Two histograms are overlaid: "post-RoPE" (green) and "pre-RoPE" (orange). A secondary inset chart provides a zoomed-in view of the dominant eigenvalue range (0–2).
---
### Components/Axes
- **X-axis (Eigenvalue)**: Linear scale from 0 to 60, with gridlines at integer intervals.
- **Y-axis (Count)**: Logarithmic scale from 10⁰ to 10³, with gridlines at 10¹, 10², and 10³.
- **Legend**: Located in the top-right corner, with green representing "post-RoPE" and orange representing "pre-RoPE."
- **Inset Chart**: Positioned in the top-right corner of the main chart, focusing on eigenvalues 0–2 with the same logarithmic y-axis.
---
### Detailed Analysis
#### Main Chart
- **pre-RoPE (orange)**:
- Dominates at eigenvalue 0, with a count exceeding 10³.
- Sharp decline after eigenvalue 0, with minimal counts beyond eigenvalue 10.
- A secondary peak near eigenvalue 10 (~10¹ count).
- **post-RoPE (green)**:
- Lower counts at eigenvalue 0 (~10¹) compared to pre-RoPE.
- More evenly distributed across eigenvalues 0–20, with counts decreasing gradually.
- Sparse counts beyond eigenvalue 20, with isolated peaks at ~30, 40, and 60.
#### Inset Chart (0–2 Eigenvalue Range)
- **pre-RoPE (orange)**:
- Peaks at eigenvalue 0 (~10² count).
- Gradual decline to ~10¹ at eigenvalue 1.0.
- **post-RoPE (green)**:
- Higher counts than pre-RoPE across the range.
- Peaks at eigenvalue 0.5 (~10¹ count) and 1.0 (~10² count).
- Slightly higher counts at eigenvalue 1.5 (~10¹ count) compared to pre-RoPE.
---
### Key Observations
1. **pre-RoPE** exhibits a highly concentrated distribution, with over 90% of counts clustered near eigenvalue 0.
2. **post-RoPE** shows a more dispersed distribution, with counts spread across eigenvalues 0–20 and isolated peaks at higher eigenvalues.
3. In the dominant range (0–2), post-RoPE consistently outperforms pre-RoPE in count density, particularly at eigenvalues 0.5 and 1.0.
4. The logarithmic y-axis emphasizes the disparity in counts at lower eigenvalues, where pre-RoPE dominates.
---
### Interpretation
The data suggests that applying RoPE significantly alters the eigenvalue distribution of keys in the Mistral-7B-instruct-v0.2 model:
- **pre-RoPE** eigenvalues are tightly concentrated near 0, indicating a lack of positional diversity in key representations.
- **post-RoPE** eigenvalues are more evenly distributed, implying improved positional encoding and potentially better handling of sequence order in the model's attention mechanism.
- The inset confirms that RoPE enhances the model's ability to distinguish between different positional relationships in the dominant eigenvalue range, which is critical for tasks requiring fine-grained context awareness.
This distribution shift highlights RoPE's role in regularizing key representations, likely contributing to the model's improved performance on tasks involving long-range dependencies.
</details>
(b)
<details>
<summary>appendix_image/llama.png Details</summary>
### Visual Description
## Line Chart: LLaMa-2-7B-Chat Keys
### Overview
The chart visualizes the "Rank keep" metric across 31 layers (0–30) for three configurations: post-RoPE (green), pre-RoPE (orange), and Full Dim (4096, gray dashed). The y-axis ranges from 0 to 4000, with the Full Dim line acting as a constant reference at 4096.
### Components/Axes
- **X-axis (layer)**: Integer values from 0 to 30, representing model layers.
- **Y-axis (Rank keep)**: Scale from 0 to 4000, with gridlines at 1000, 2000, 3000.
- **Legend**:
- Green: post-RoPE
- Orange: pre-RoPE
- Gray dashed: Full Dim (4096)
- **Line styles**: Solid for post-RoPE/pre-RoPE, dashed for Full Dim.
### Detailed Analysis
1. **post-RoPE (green)**:
- Starts at ~100 at layer 0.
- Peaks at ~2200 at layer 25.
- Drops to ~1400 at layer 30.
- Shows gradual increase until layer 25, then a sharp decline.
2. **pre-RoPE (orange)**:
- Starts at 0 at layer 0.
- Rises to ~1300 at layer 25.
- Declines to ~800 at layer 30.
- Steady growth until layer 25, followed by a gradual drop.
3. **Full Dim (4096)**:
- Constant dashed line at 4096 across all layers.
### Key Observations
- **Divergence**: post-RoPE consistently outperforms pre-RoPE in Rank keep until layer 25, where both approach ~1300–2200.
- **Peak at layer 25**: post-RoPE reaches its maximum Rank keep (~2200), while pre-RoPE peaks at ~1300.
- **Decline at layer 30**: post-RoPE drops sharply (~2200 → 1400), while pre-RoPE declines more gradually (~1300 → 800).
- **Full Dim gap**: Both configurations remain below the Full Dim threshold (4096) throughout.
### Interpretation
The data suggests that post-RoPE significantly enhances Rank keep compared to pre-RoPE, particularly in mid-layers (layers 10–25). The sharp decline in post-RoPE at layer 30 may indicate a limitation in later layers, while pre-RoPE’s gradual decline suggests diminishing returns. The Full Dim line (4096) represents an upper bound, implying that neither configuration fully utilizes the model’s dimensional capacity. The peak at layer 25 could highlight an optimal layer for post-RoPE performance, warranting further investigation into layer-specific optimizations.
</details>
(c)
<details>
<summary>appendix_image/mistal.png Details</summary>
### Visual Description
## Line Chart: Mistral-7B-Instruct-v0.2 Keys
### Overview
The chart compares the "Rank keep" metric across layers (0-30) for two configurations of the Mistral-7B-Instruct-v0.2 model: "post-RoPE" (green line) and "pre-RoPE" (orange line), against a constant "Full Dim (1024)" baseline (gray dashed line). The y-axis represents rank keep values (0-1000), while the x-axis represents model layers.
### Components/Axes
- **X-axis (Layer)**: Discrete values from 0 to 30, labeled "layer."
- **Y-axis (Rank keep)**: Continuous scale from 0 to 1000, labeled "Rank keep."
- **Legend**: Located in the top-left corner, with:
- Green: "post-RoPE"
- Orange: "pre-RoPE"
- Gray dashed: "Full Dim (1024)"
- **Lines**:
- Post-RoPE: Green solid line with circular markers.
- Pre-RoPE: Orange solid line with circular markers.
- Full Dim: Gray dashed line at y=1024.
### Detailed Analysis
1. **Post-RoPE (Green)**:
- Starts at ~50 at layer 0.
- Peaks at ~300 at layer 2.
- Fluctuates between 200-400, with a secondary peak at ~450 at layer 25.
- Ends at ~350 at layer 30.
2. **Pre-RoPE (Orange)**:
- Starts near 0 at layer 0.
- Rises to ~150 at layer 2.
- Fluctuates between 100-300, peaking at ~300 at layer 25.
- Ends at ~200 at layer 30.
3. **Full Dim (Gray Dashed)**:
- Constant at y=1024 across all layers.
### Key Observations
- **Post-RoPE vs. Pre-RoPE**: Post-RoPE consistently maintains higher rank keep values than pre-RoPE, especially in later layers (e.g., layer 25: ~450 vs. ~300).
- **Trend**: Both lines show gradual improvement in rank keep up to layer 25, followed by stabilization or slight decline.
- **Full Dim Baseline**: The dashed line at 1024 remains far above both data series, suggesting rank keep values are significantly lower than the full dimensionality.
### Interpretation
The data suggests that the "post-RoPE" configuration enhances rank keep performance compared to "pre-RoPE," particularly in deeper layers (layers 20-25). The Full Dim baseline (1024) indicates that rank keep values are constrained relative to the model's full dimensionality. The divergence between post-RoPE and pre-RoPE implies that RoPE modifications (likely positional encoding adjustments) improve the model's ability to retain rank information, which could correlate with better performance in tasks requiring positional awareness. The stabilization after layer 25 may reflect diminishing returns or architectural saturation.
</details>
(d)
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.
In this section, we conduct a numerical analysis to quantify how rotary position embedding (RoPE) alters the principal-component structure of the key states. We perform principal component analysis (PCA) on the key tensors before and after applying RoPE, denoted as pre-RoPE and post-RoPE, respectively. Concretely, we first compute the covariance matrix of the key states and then carry out an eigenvalue decomposition. The magnitude of each eigenvalue reflects the contribution of its associated eigenvector: the presence of many small eigenvalues indicates that the corresponding representation is effectively low rank. To relate rank to the spectrum more systematically, we adopt the metrix introduced in Loki [20]:
$$
_l(v) = \min\Bigl\{ d∈ℤ_+:∑_i=1^dλ_l^(i) ≥ \frac{v}{100}\Bigr\},
$$
where $λ_l^(i)$ denotes the $i$ -th largest eigenvalue of the covariance matrix for the keys in layer $l$ . This definition specifies the smallest number of principal components needed to explain at least $v\$ of the total variance, enabling a direct comparison of the intrinsic dimensionality before and after RoPE.
Figure 4 (a)–(b) illustrate the eigen–value spectra of the first layer ( $l=0$ ) for Llama–2–7B–Chat and Mistral–7B–Instruct–v0.2 under the pre–RoPE and post–RoPE conditions. The pre–RoPE spectra exhibit a markedly larger number of small eigenvalues, whereas the post–RoPE spectra are shifted upward, corroborating our earlier finding that RoPE increases the overall variance. Figure 4 (c)–(d) present the layer–wise $_l(90)$ values, i.e. the minimal dimensionality needed to retain $90 \$ of the variance. For both models, the post–RoPE condition consistently requires a higher rank, reflecting the broader spectra observed in panels (a)–(b). In addition, the required rank varies substantially across layers, indicating that a layer–adaptive rank selection scheme could further enhance compression efficiency.