RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

TL;DR

RazorAttention tackles the KV-cache explosion in long-context LLMs by introducing a head-wise caching scheme that preserves all information inside a small subset of retrieval heads while compressing non-retrieval heads via remote-token dropping and compensation tokens. The method applies to both ALiBi and RoPE positional embeddings, identifying retrieval heads through theoretical bounds (ALiBi) or data-free head-scores (RoPE) and achieving about 70% KV-cache reduction with negligible accuracy loss, all without retraining and with compatibility to FlashAttention. Comprehensive experiments on LongBench and Needle In A Haystack across multiple models demonstrate near-baseline performance at significant compression, highlighting the approach as a practical, plug-and-play solution for efficient long-context inference. Ablation studies confirm the importance of echo/induction heads and compensation tokens in maintaining information fidelity during compression.

Abstract

The memory and computational demands of Key-Value (KV) cache present significant challenges for deploying long-context language models. Previous approaches attempt to mitigate this issue by selectively dropping tokens, which irreversibly erases critical information that might be needed for future queries. In this paper, we propose a novel compression technique for KV cache that preserves all token information. Our investigation reveals that: i) Most attention heads primarily focus on the local context; ii) Only a few heads, denoted as retrieval heads, can essentially pay attention to all input tokens. These key observations motivate us to use separate caching strategy for attention heads. Therefore, we propose RazorAttention, a training-free KV cache compression algorithm, which maintains a full cache for these crucial retrieval heads and discards the remote tokens in non-retrieval heads. Furthermore, we introduce a novel mechanism involving a "compensation token" to further recover the information in the dropped tokens. Extensive evaluations across a diverse set of large language models (LLMs) demonstrate that RazorAttention achieves a reduction in KV cache size by over 70% without noticeable impacts on performance. Additionally, RazorAttention is compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM inference efficiency without overhead or retraining of the original model.

Figures (6)

• Figure 1: RazorAttention achieves comparable performance to the original model, even with $70\%$ KV cache compressed. To demonstrate this, we tested Llama2-13B-64K fu2024data on the Needle in A Haystack benchmark gkamradt2023needle.
• Figure 2: Importance-based token-dropping methods cannot work when querying the less relevant information to the main theme. Here, we use an 8K document from LongBench bai2023longbench and add two sentences that are not relevant to the main theme. In this case, H2O discards tokens that are less relevant to the main theme, leading to failures in both Q1 and Q2. SnapKV discards tokens based on the first query, making it effective for Q1 but failing in subsequent queries like Q2. Only RazorAttention successfully outputs the exact information from the lengthy input even when we compress $70\%$ of the KV cache.
• Figure 3: In Figure \ref{['fig:show_rope_ra']} we present the illustration of how RazorAttention compress the KV cache. For retrieval heads, we maintain a full cache for retaining all the tokens' information. For non-retrieval heads, we directly discard remote tokens and compress the discarded tokens into a compensation token whose KV cache is denoted as $\{\hat{\bm{k}}, \hat{\bm{v}}\}$. In Figure \ref{['fig:induction+echo']} we provide an illustration example of the echo head and induction head. The current token is "B" and the generated token is "C". In this case, the echo head would mainly attend to token "B" while the induction head mainly attend to token "C" in previous context.
• Figure 4: Performance comparison of RazorAttention and other compression algorithms on Llama2-7b-80K, Needle In A Haystack. Notice that H2O is incompatible with FlashAttention so we get OOM errors when tested on longer sequences, and its performance has already become unusable in this case.
• Figure 5: Adding $1\%$ of the echo heads can significantly enhances the retrieving performance of RazorAttention on Llama2-7B-80k.

Theorems & Definitions (2)

• Theorem 1
• proof