🤖 AI

RAP: KV-Cache Compression via RoPE-Aligned Pruning

This paper introduces RoPE-Aligned Pruning (RAP), a structured pruning method that preserves Rotary Position Embedding semantics by removing Key-Value cache channels in aligned pairs, thereby achieving significant memory and compute savings for long-context LLM inference without sacrificing accuracy.

Original authors: Jihao Xin, Tian Lyu, David Keyes, Hatem Ltaief, Marco Canini

Published 2026-08-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jihao Xin, Tian Lyu, David Keyes, Hatem Ltaief, Marco Canini

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a very long story. Every time you hear a new sentence, you have to keep a mental list of all the previous sentences so you can understand how they connect. In the world of artificial intelligence, this mental list is called the "KV cache." It's like a giant notebook where a super-smart robot writes down the most important details of a conversation so it doesn't forget the beginning by the time it reaches the end. The problem is, as the story gets longer, this notebook gets huge. It takes up so much memory and requires so much brainpower to flip through the pages that the robot starts to slow down or even run out of space entirely. Scientists are always looking for ways to shrink this notebook without losing the important parts of the story, so the robot can chat about long topics without crashing.

One popular trick the robot uses to understand where things happen in the story is called "Rotary Position Embedding" (RoPE). Think of RoPE as a special dance the robot does with its memory. Instead of just writing down a number for "sentence 5," the robot pairs up two numbers and spins them around like a hula hoop. This spinning tells the robot exactly how far away that sentence is from the current one. The crucial rule of this dance is that the two numbers in the pair must stay together; if you separate them, the spin breaks, and the robot gets confused about the timeline.

Now, enter a new method called RAP (RoPE-Aligned Pruning). The researchers behind this idea noticed that previous attempts to shrink the robot's memory notebook were making a fatal mistake. Imagine you have a backpack full of these dancing pairs. Old methods tried to save space by randomly throwing out single items from the backpack. But because the items were dancing in pairs, throwing out just one left its partner spinning alone in the air with no one to hold hands with. The dance fell apart, and the robot's memory became useless.

The authors of this paper suggest a smarter way to pack the bag. Instead of tossing out random single items, RAP looks at the dancing pairs and decides to throw out whole pairs at once. If a pair isn't very important, the whole duo goes. If a pair is important, the whole duo stays. This keeps the dance intact. The researchers tested this on several popular AI models (ranging from 3 billion to 14 billion "brain cells") and found that this method works remarkably well. They were able to shrink the memory notebook by 30% (keeping only 70% of the original size) while maintaining high accuracy, though with a small, measurable drop in performance compared to the full, uncompressed model.

What makes this even cooler is that other methods that try to shrink the memory often have to do extra work every time the robot speaks to rebuild the missing pieces, which slows everything down. RAP, however, is like editing the notebook before the robot starts talking. It removes the extra pages permanently, so the robot doesn't have to do any extra math to fill in the gaps. It runs faster, uses less energy, and still remembers the story very well. The paper shows that while other tricks might save space, they often break the robot's sense of time or make it slower. RAP is one of the top methods that manages to shrink the memory, speed up the robot, and keep the story accurate all at once, staying very close to the strongest low-rank competitors while offering unique efficiency benefits. It's a bit like realizing you don't need to carry the whole encyclopedia to tell a story; you just need to carry the right chapters, and make sure you don't tear the pages in half.

Technical Summary: RAP (RoPE-Aligned Pruning) for KV-Cache Compression

1. Problem Statement

Long-context inference in Large Language Models (LLMs) is increasingly bottlenecked by the memory and computational requirements of the Key-Value (KV) cache. As context length scales, the KV cache dominates GPU memory usage and attention computation costs. While structured pruning offers a direct path to compress the cache by removing less useful channels from the $W_k$ and $W_v$ projection weights, existing methods face a critical incompatibility with Rotary Position Embedding (RoPE).

Modern LLMs apply RoPE after the Query/Key projections, rotating feature dimensions in pairs. Standard per-channel pruning methods (e.g., KVPruner) often remove individual channels without regard for this pairing structure. This "orphaning" of half-pairs breaks the rotational semantics of RoPE, leading to a silent but catastrophic collapse in model accuracy. Furthermore, existing low-rank decomposition methods (e.g., SVD-LLM, PaLU) that attempt to compress weights often cannot fuse their expansion matrices with the Query projection ( $W_q$ ) due to the non-commutativity of RoPE and matrix multiplication, forcing a costly per-step reconstruction of the full-dimensional Key vectors during decoding.

2. Methodology: RoPE-Aligned Pruning (RAP)

The authors propose RAP, a structured pruning framework designed specifically to respect the pair-wise structure imposed by RoPE.

Core Principle

Instead of pruning individual channels, RAP constrains the pruning granularity to RoPE-aligned pairs. By removing whole pairs of dimensions, the method ensures that the rotation operation remains mathematically valid and that the positional semantics are preserved.

The RAP Pipeline

RoPE-Pair Scoring:
- For the Key projection weights ( $W_k$ ), columns are grouped into RoPE pairs according to the model's specific pairing strategy (either adjacent or half-apart).
- Importance is measured using Fisher information scores calculated on a small calibration dataset. For a pair $p = (j, j')$ , the score $\sigma_p$ is the aggregated root-Fisher magnitude of the two channels.
- For the Value projection ( $W_v$ ), which lacks RoPE pairing, individual channel Fisher scores are computed.
Budget Allocation:
- A global retain ratio $\rho$ is distributed across layers and heads.
- Layer-Adaptive: Layers are assigned budgets based on their sensitivity (aggregated Fisher scores), giving more sensitive layers a larger budget.
- Head-Uniform: Within a layer, all heads receive the same number of retained pairs to facilitate batched GEMM operations.
- Balanced K/V: To prevent the Key dimension from being crushed (which harms long-context performance), RAP allocates the same per-head channel budget to both $K$ and $V$ . Specifically, if $K$ retains $m$ pairs, $V$ retains $2m$ channels.
Construction and Fusion:
- A binary, pair-preserving selection matrix $E_k$ is constructed to map the retained pairs back to their original indices.
- Crucially, because $E_k$ preserves the pair structure, it commutes with the RoPE operation: $RoPE(X \tilde{W}_k E_k) = RoPE(X \tilde{W}_k) E_k$ .
- This commutativity allows the pruning mask to be fused offline into the Query projection: $\tilde{W}_q = W_q E_k^\top$ .
- Result: The model runs on pruned dimensions, and the KV cache stores only the low-dimensional latent keys. No online reconstruction is required.
Accuracy Recovery:
- To recover accuracy lost during pruning, RAP employs a lightweight LoRA (Low-Rank Adaptation) update combined with Knowledge Distillation (KD).
- A small rank ( $r=8$ ) LoRA adapter is trained on the pruned model to align its outputs with the dense teacher model, then merged before deployment.

3. Key Contributions

Identification of RoPE Incompatibility: The paper highlights that standard channel pruning breaks RoPE pairs, causing accuracy collapse, and that this failure is often silent (no runtime errors, only degraded performance).
RoPE-Aligned Pruning Strategy: RAP introduces a pruning granularity that strictly adheres to RoPE pairs, ensuring the rotation operation remains valid without requiring complex retraining or reconstruction.
Reconstruction-Free Decoding: Unlike low-rank decomposition methods that must reconstruct full-dimensional keys at every decode step, RAP fuses the pruning mask into $W_q$ offline, eliminating the "reconstruction tax" and reducing both memory and compute overhead.
Comprehensive Evaluation: The method is evaluated across multiple model families (Llama, Mistral, Qwen) and scales (3B to 14B), demonstrating robustness.

4. Experimental Results

Evaluations were conducted under an iso-KV setting (matching the retained KV channel ratio $\rho$ ) to ensure fair comparison of quality and efficiency.

Accuracy vs. Compression:
- At a 30% compression ratio ( $\rho=0.7$ ), RAP maintains usable accuracy (e.g., WikiText-2 PPL ~9.12 for Llama-3-8B vs. baseline 8.28), significantly outperforming RoPE-blind channel pruning (KVPruner), which suffers massive accuracy collapse (PPL ~397).
- RAP achieves performance comparable to or better than strong low-rank baselines (PaLU, SVD-LLM) while operating at lower attention costs.
Efficiency Gains:
- Parameters & FLOPs: RAP reduces attention parameters and FLOPs linearly with the retain ratio $\rho$ (e.g., ~70% at $\rho=0.7$ ). In contrast, low-rank methods often retain higher parameter counts due to the need for online expansion factors.
- Latency: RAP reduces Prefill and Decode latency (e.g., 83%/77% of dense baseline at $\rho=0.7$ ) because it avoids per-step reconstruction. Low-rank methods often show increased latency, particularly during decoding.
Robustness:
- RAP maintains performance across different model sizes and architectures.
- It is orthogonal to and composable with 4-bit weight quantization.
- The method handles non-contiguous RoPE indexing efficiently via a custom Triton kernel, avoiding the overhead of standard PyTorch indexing.

5. Significance and Claims

The paper claims that RAP is the only method that simultaneously:

Reduces the KV cache size.
Reduces model parameters and projection FLOPs.
Maintains high accuracy at high compression ratios.
Operates without per-step reconstruction overhead.

The authors emphasize that RAP is a "drop-in" solution for RoPE-based LLMs. By respecting the structural constraints of Rotary Position Embedding, it solves the "silent failure" problem of existing pruning methods and offers a more efficient alternative to low-rank decomposition for long-context inference. The work suggests that structural alignment with model-specific mechanisms (like RoPE) is critical for effective model compression.