RAP: KV-Cache Compression via RoPE-Aligned Pruning

TL;DR

RAP introduces a novel pruning method for RoPE-based large language models that significantly reduces memory, compute, and latency costs of KV-Cache without sacrificing accuracy.

Contribution

It proposes RoPE-Aligned Pruning (RAP), a technique that preserves RoPE structure to enable effective low-rank compression and reduce resource usage.

Findings

Achieves 20-30% reduction in KV-Cache, attention parameters, and FLOPs.

Reduces attention latency to 83% (prefill) and 77% (decode) of baseline.

Maintains strong accuracy despite compression.

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.