Mastering Long Contexts in LLMs with KVPress

One of the key features of Large Language Models (LLMs) is their context window —the maximum number of tokens they can process in a single request. As LLMs evolve, their context windows are becoming increasingly larger.

However, these extended windows come at a cost—memory taken by the long context in the KV Cache becomes hard to manage. For instance, handling 1M tokens with Llama 3-70B in float16 demands 330GB for KV Cache, rendering it infeasible for many applications.

In this blog post, we'll address one solution for this problem: compressing the KV Cache for more efficient generation. To achieve this, we'll explore:

The values stored in the KV Cache come from all the attention blocks used by the model. Therefore, its size depends on the model architecture, which dictates the number of attention heads. More concretely, the memory consumed by the KV Cache is determined by the following equation:

Since bfloat16 uses 2 bytes per parameter, the model weights alone require 140 GB (70B x 2 bytes). This means that running the model with a 1M token context size demands approximately 470 GB of memory, with the KV cache alone accounting for a staggering 70% of this total.

Many of these presses rely on a score that is used in each head to prune the KV pairs with the lowest importance. For instance the KnormPress prunes the KV pairs with lowest key value norm ( paper ), and SnapKVPress prunes the KV pairs associated with a low attention weights for the latest queries ( paper ).

The plot below demonstrates the GPU memory savings achieved with KVPress compression as prompt lengths increase. For shorter prompts, most of the memory is allocated to model weights—approximately 15GB for Llama 3.1 8B in bfloat16. However, as prompt lengths grow, the KV cache becomes a major contributor to memory consumption. For 128k context length, applying KVPress with a 50% compression ratio reduces peak memory usage from 45GB to 37GB. This smaller KV cache also improves decoding speed, from 11 tokens per second to 17 tokens per second on an A100 GPU ( source ). .

For researchers specializing in compression, KVPress offers a flexible and modular framework, making it easy to understand and extend with new methods. For developers , KVPress simplifies the process of deploying these cutting-edge techniques, enabling quick and efficient integration into real-world applications.

Before getting started, explore KVPress in this Space (you'll find examples at the end if needed):

In autoregressive models, text generation happens token by token , with each prediction relying on all preceding tokens for context. For example:

This repetitive computation becomes inefficient as the sequence grows, especially for large models. KV Cache optimizes this process by storing the intermediate results—keys (K) and values (V)—from the attention layers, so the model can reuse them for future tokens instead of recalculating them.