KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head

TL;DR

KV-Compress tackles the memory bottleneck of long-context inference by introducing a paged KV-cache eviction scheme that supports variable eviction rates across layers and heads. By aggregating eviction metrics over grouped queries and evicting contiguous KV blocks within a PagedAttention framework, the method reduces memory footprint while sustaining performance. It achieves state-of-the-art results on LongBench for Mistral-7B-Instruct-v0.2 and Llama-3.1-8B-Instruct, with up to $8\times$ compression and significant throughput gains when integrated with vLLM (up to $5.18\times$). The approach hinges on query-group compression, on-device block management, and a squared-attention metric, enabling high compression (up to $64\times$ in some cases) with minimal degradation and practical deployment in large-scale inference systems.

Abstract

Context lengths of Large Language Models (LLMs) have exploded in recent years, with 128k-token context becoming a standard and million-token context becoming a reality. Efficiently supporting long-context inference remains challenging as the memory that must be allocated in key-value (KV) cache for a generation scales with its context length, limiting the number of long-context requests that can be served concurrently under a given memory budget. KV cache compression can mitigate this issue by removing under-utilized KVs from each attention head's cache and reducing its memory footprint. Higher theoretical compression rates can be achieved when the number of removed KVs varies across attention heads, but application of such a strategy within existing inference frameworks adds fragmentation and cannot realize the theoretical compression rates in physical memory. We introduce KV-Compress, a novel compression method that evicts contiguous KV blocks within a PagedAttention framework, reducing the memory footprint of the KV cache proportionally to this theoretical compression rate. Our method achieves state-of-the-art performance on LongBench for both Mistral-7B-Instruct-v0.2 and Llama-3.1-8B-Instruct while lowering the total number of compressed KVs by 4x compared with prior methods. Evaluations on Llama-3.1-8B-Instruct and Llama-3.1-70B-Instruct-FP8 achieve compression rates up to 8x with negligible impact on performance, and up to 64x while retaining over 90% of full-cache performance for all but three of the suite's subsets. We benchmark an integration of our method with vLLM that increases total throughput by up to 5.18x by enabling larger decoding batches.

Figures (8)

• Figure 1: Llama-3.1-8B-Instruct performance for different rates of compression. \ref{['fig:llama-8b-throughput-acc']} plots average LongBench subtask performance for each subtask category, measured as a percentage of the accuracy achieved without compression. \ref{['fig:llama-8b-throughput-thrpt']} displays the throughput that can be achieved on an NVIDIA L4 for the same range of compression rates. Plots are shown for varying input lengths.
• Figure 2: Llama-3.1-70B-Instruct-FP8 performance for different rates of compression. \ref{['fig:llama-70b-throughput-acc']} plots average LongBench subtask performance for each subtask category, measured as a percentage of the accuracy achieved without compression. \ref{['fig:llama-70b-throughput-thrpt']} displays the throughput that can be achieved on an NVIDIA H100 for the same range of compression rates. Plots are shown for varying input lengths.
• Figure 3: Comparison of block management in vanilla vLLM versus with KV-Compress. In vLLM (left) each cache block stores KVs for every attention head of every layer, while with KV-Compress (right) each block only holds KVs for a single head. The KV-Compress block tables are therefore expanded by $l \times H$ so that the unique block for each specific KV head and layer can be retrieved. Due to the larger number of blocks, we move block management to the GPU so scheduling operations can be done in parallel. Note that every cell of KV cache in the above diagram represents a $d$-dimensional vector, while all others represent scalars.
• Figure 4: Llama-3.1-8B-Instruct percent of full-cache performance by compression rate on all LongBench subtasks. Results are grouped by subtask category.
• Figure 5: Visualization of our compression algorithm for a simplified example with two KV heads and a block size of two. KV metrics are visualized for a given cache state, highlighting blocks of a particular sequence in the decoding batch that is scheduled to evict two blocks. Logical indices are displayed under the corresponding metrics slot.