Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RAP: KV-Cache Compression via RoPE-Aligned Pruning

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

TL;DR

RAP tackles the KV-Cache memory bottleneck in long-context LLM inference by introducing RoPE-Aligned Pruning, a structured pruning approach that preserves RoPE's 2×2 rotation blocks to maintain RoPE-commutativity. By pruning entire RoPE-aligned column pairs and absorbing the remaining factor into downstream weights, RAP eliminates reconstruction overhead while jointly reducing KV-Cache, attention parameters, and FLOPs. The method employs Fisher-information-based pair scoring, adaptive budget allocation, and a KD/LoRA-based accuracy recovery, with practical considerations for RoPE variants and hybrid compression. Empirical results on LLaMA-3-8B-Instruct and Mistral-7B-v0.3 show 20–30% joint savings and substantial latency reductions (prefill 83%, decode 77% of baseline at moderate compression) while preserving accuracy, making RAP a drop-in upgrade for RoPE-based LLMs.

Abstract

Long-context inference in large language models is increasingly bottlenecked by the memory and compute cost of the KV-Cache. Low-rank factorization compresses KV projections by writing $W \approx A * B$, where A produces latent KV states and B can be absorbed into downstream weights. In modern RoPE-based LLMs, this absorption fails: RoPE forces latent KV states to be reconstructed to full dimension, reintroducing substantial memory and compute overhead. We propose RoPE-Aligned Pruning (RAP), which prunes entire RoPE-aligned column pairs to preserve RoPE's 2x2 rotation structure, restore B absorption, and eliminate reconstruction. Our evaluation on LLaMA-3-8B and Mistral-7B shows that RAP enables joint reduction of KV-Cache, attention parameters, and FLOPs by 20-30%, all at once, while maintaining strong accuracy. Notably, RAP reduces attention latency to 83% (prefill) and 77% (decode) of baseline.

RAP: KV-Cache Compression via RoPE-Aligned Pruning

TL;DR

RAP tackles the KV-Cache memory bottleneck in long-context LLM inference by introducing RoPE-Aligned Pruning, a structured pruning approach that preserves RoPE's 2×2 rotation blocks to maintain RoPE-commutativity. By pruning entire RoPE-aligned column pairs and absorbing the remaining factor into downstream weights, RAP eliminates reconstruction overhead while jointly reducing KV-Cache, attention parameters, and FLOPs. The method employs Fisher-information-based pair scoring, adaptive budget allocation, and a KD/LoRA-based accuracy recovery, with practical considerations for RoPE variants and hybrid compression. Empirical results on LLaMA-3-8B-Instruct and Mistral-7B-v0.3 show 20–30% joint savings and substantial latency reductions (prefill 83%, decode 77% of baseline at moderate compression) while preserving accuracy, making RAP a drop-in upgrade for RoPE-based LLMs.

Abstract

Long-context inference in large language models is increasingly bottlenecked by the memory and compute cost of the KV-Cache. Low-rank factorization compresses KV projections by writing , where A produces latent KV states and B can be absorbed into downstream weights. In modern RoPE-based LLMs, this absorption fails: RoPE forces latent KV states to be reconstructed to full dimension, reintroducing substantial memory and compute overhead. We propose RoPE-Aligned Pruning (RAP), which prunes entire RoPE-aligned column pairs to preserve RoPE's 2x2 rotation structure, restore B absorption, and eliminate reconstruction. Our evaluation on LLaMA-3-8B and Mistral-7B shows that RAP enables joint reduction of KV-Cache, attention parameters, and FLOPs by 20-30%, all at once, while maintaining strong accuracy. Notably, RAP reduces attention latency to 83% (prefill) and 77% (decode) of baseline.
Paper Structure (40 sections, 2 theorems, 21 equations, 24 figures, 17 tables, 2 algorithms)

This paper contains 40 sections, 2 theorems, 21 equations, 24 figures, 17 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Transformers and Attention Mechanism
  4. KV-Cache Compression and Model Pruning
  5. Motivation: RoPE Incompatibility
  6. RAP: RoPE-Aligned Pruning
  7. Step 1: RoPE Pair Scoring
  8. Step 2: Budget Allocation
  9. Step 3: RAP Construction
  10. Step 4: Accuracy Recovery
  11. Practical Considerations
  12. RoPE-Invariant Error Bounds
  13. Experiments
  14. Setup
  15. Results
  16. ...and 25 more sections

Key Result

Theorem 5.1

Let $\mathcal{L}(\cdot)$ denote the model's training loss as a function of the key projection matrix, with all other parameters held fixed. Let $A_k$ denote the pruned matrix obtained by RAP by removing RoPE pairs $\mathcal{P} \setminus \mathcal{S}$. Under a second-order Taylor approximation of $\ma where $\sigma_p$ is the Fisher score defined in Eq. eq:score.

Figures (24)

  • Figure 1: SVD-based KV-Cache compression requires full-dimensional reconstruction at runtime.
  • Figure 2: Optimizations of KV-Cache/FLOPs/Parameters.
  • Figure 3: RAP Example: Weight matrix $\mathbf{W}$ with 12 columns is grouped into 6 RoPE pairs, then pruned 50% based on the RoPE scores.
  • Figure 4: LLaMA PPL of pruning one layer at a time.
  • Figure 5: Parameters Relative to Baseline
  • ...and 19 more figures

Theorems & Definitions (4)

  • Definition 1.1
  • Theorem 5.1: Loss induced by RoPE-aware pair pruning
  • Corollary 5.2: Optimality of Fisher-based RoPE pair selection
  • proof