DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs

TL;DR

This work tackles the memory bottleneck of KV cache in long-context LLMs by revealing that token retention across layers is task-dependent. It introduces DynamicKV, a layer-aware KV cache compression method that dynamically allocates per-layer budgets and uses progressive cache updates during the prefilling phase to adapt to the task at hand. Empirical results on LongBench show DynamicKV can preserve near-FullKV performance with a small token budget (as low as 6.9%), while Needle-in-a-Haystack evaluations demonstrate superiority over SOTA methods under extreme compression. The approach integrates with existing inference pipelines and maintains decoding efficiency, offering practical gains in latency and memory usage for long-context applications.

Abstract

Efficient KV cache management in LLMs is crucial for long-context tasks like RAG and summarization. Existing KV cache compression methods enforce a fixed pattern, neglecting task-specific characteristics and reducing the retention of essential information. However, we observe distinct activation patterns across layers in various tasks, highlighting the need for adaptive strategies tailored to each task's unique demands. Based on this insight, we propose DynamicKV, a method that dynamically optimizes token retention by adjusting the number of tokens retained at each layer to adapt to the specific task. DynamicKV establishes global and per-layer maximum KV cache budgets, temporarily retaining the maximum budget for the current layer, and periodically updating the KV cache sizes of all preceding layers during inference. Our method retains only 1.7% of the KV cache size while achieving ~85% of the Full KV cache performance on LongBench. Notably, even under extreme compression (0.9%), DynamicKV surpasses state-of-the-art (SOTA) methods by 11% in the Needle-in-a-Haystack test using Mistral-7B-Instruct-v0.2. The code will be released.

Figures (5)

• Figure 1: Comparison of DynamicKV with traditional methods in maintaining KV cache size across layers. Left: the structure difference: (a) Retain all KV cache. (b) Fixed KV cache for each layer (e.g., StreamingLLM, H2O, SnapKV). (c) Hierarchically decreasing pyramid KV cache retention. (d) Ours DynamicKV: layer-aware adaptive KV cache retention. Right: average accuracy on different KV cache retention.
• Figure 2: Analyzing the distribution of token retention across layers in LlaMA for different tasks, including Document QA, Summarization, Synthetic Task and Code Completion. (a) Each boxplot shows the distribution of token retention rates on different types of tasks across different layers. Results for different layers show that the token retention rates vary significantly across different tasks. (b) We visualize the token retention rates across different layers for four tasks, showing that the token retention rates exhibit different patterns across tasks.
• Figure 3: Overview of our DynamicKV structure and KV cache compression comparison. Left: Layer-wise KV cache retention mechanism in transformer architectures. Right: Our proposed DynamicKV framework employs stage-wise dynamic updating to maintain KV cache within predefined memory budgets, with task-specific visualization showing KV cache preservation patterns across layers.
• Figure 4: Performance of DynamicKV with different KV cache size on LongBench. The evaluation metrics are the average score of LongBench across datasets.
• Figure 5: Performance Comparison on the Needle in a Haystack Task using Mistral-7B-Instruct-v0.2 with 32k context size in 64 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the length.