Adaptive compression improves memory and speed for large language models
MetaKV: Adaptive KV Cache Compression for Constrained LLM Inference
Machine Learning
Summary
Large language models need to remember past words using key-value caches, which can use a lot of memory, especially for long texts. The authors found that different ways to shrink these caches work better for different texts and limits on memory and speed. They created MetaKV, a system that picks the best way to compress the memory for each input, balancing accuracy, speed, and memory use. This method works better than using just one fixed setting for all inputs, helping the model answer questions correctly while staying within memory and speed limits.
What this means in practice
- •For ai engineers: Optimize language model deployments by dynamically choosing cache compression to meet specific memory and speed constraints per input.
- •For mobile app developers: Support running large language model inference efficiently on devices with limited memory by adapting cache compression per request.
Authors
Michael Wang, Keith Li, Roozbeh Bostandoost
Abstract
Key--value (KV) cache compression is an effective way to reduce the memory overhead of large language model (LLM) inference, particularly for long-context workloads. However, existing compression methods make different trade-offs among accuracy, inference latency, and peak KV cache memory utilization, making a single fixed configuration unsuitable across different prompts and resource constraints. We introduce MetaKV, an adaptive framework that selects a KV cache compression configuration for each input prompt based on user-specified latency and peak memory budgets. MetaKV uses lightweight prediction models to estimate the end-to-end latency, peak memory, and probability of a correct response for each candidate configuration, and selects the configuration that best satisfies the latency-memory constraints while preserving accuracy. We evaluate MetaKV across ten configurations from three representative KV cache compression methods, KVQuant, H$_2$O, and RocketKV, together with an uncompressed FP16 configuration, on four datasets covering mathematics, science, commonsense reasoning, and reading comprehension. Across a wide range of latency and peak memory constraints, MetaKV consistently outperforms the best static configuration, improving constrained success rate (CSR), the fraction of prompts answered correctly while satisfying both constraints, by approximately 0.07 on average and up to 0.135. These results demonstrate the benefit of adapting KV cache compression to individual prompts and latency-memory constraints. Code is available at https://github.com/MichaelWang0505/MetaKV.git