Adaptive method improves memory use and speed for large AI models
MetaKV: Adaptive KV Cache Compression for Constrained LLM Inference
Machine Learning
Summary
Large language models need a lot of memory to remember past words when answering questions, especially for long conversations. The researchers found that different ways to shrink this memory, called KV cache compression, each have their own problems with speed, accuracy, or memory use. They created MetaKV, which picks the best way to compress the memory for each question based on how much memory and speed is allowed. This makes the AI give more correct answers while using resources efficiently. Tests showed MetaKV worked better than using just one fixed method for all questions.
key-value cachecompressionlarge language modelsinference latencymemory utilizationadaptive methodspromptaccuracymachine learningresource constraints
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