🤖 AI Summary
This work addresses the severe memory and throughput bottlenecks in long-context large language model inference caused by the linear growth of key-value (KV) cache with context length. The authors propose MosaicKV, the first framework to enable dynamic, fine-grained two-dimensional KV cache compression that adaptively selects compression strategies based on the importance of elements within cache segments across both sequence and channel dimensions. MosaicKV introduces a co-designed cache management mechanism that efficiently orchestrates heterogeneous GPU/CPU resources to accelerate attention computation. Experiments on H800 demonstrate up to 16× speedup in attention computation, 7.3× higher throughput, 4.8× lower decoding latency, and 3× reduction in memory consumption, with only a 1.76% average accuracy drop on LongBench and RULER benchmarks.
📝 Abstract
Long-context LLM services now sustain prompts with hundreds of thousands to millions of tokens, making the key-value (KV) cache a first-order serving cost. Because the cache grows linearly with context length, it can exhaust GPU memory, force smaller batches, and reduce serving throughput. Prior KV cache compression techniques typically target only the sequence dimension or only the channel dimension, which leaves limited headroom as context windows scale. Compressing both dimensions promises higher memory reduction, but applying the two forms of compression directly leads to significant accuracy loss.
This paper introduces MosaicKV, a dynamic two-D (dimensional) KV cache compression system for extremely long-context serving. MosaicKV uses dynamic two-D compression to address the accuracy challenge, exploiting the non-uniform importance distribution of elements within the KV cache. Instead of applying one compression pattern globally, MosaicKV identifies important elements for each KV vector and selects compression strategies at the granularity of KV cache segments. To address the performance challenge, where fine-grained sparsity and compression management overhead can offset the gains from compression, MosaicKV introduces compressed KV cache management. This mechanism uses underutilized GPU and CPU resources to maintain compressed KV caches and accelerate attention computation.
Evaluation on an H800 GPU with multiple LLMs shows that MosaicKV delivers up to 16x attention speedup, 4.8x lower decode latency, and 7.3x higher throughput than the uncompressed baseline. At the same time, it reduces memory usage by 3x and incurs only 1.76% average accuracy loss on LongBench and RULER.