ZigZagkv: Dynamic KV Cache Compression for Long-context Modeling based on Layer Uncertainty

TL;DR

ZigZagKV tackles the memory bottleneck of KV caches in long-context LLM inference by introducing a training-free, layer-uncertainty–driven dynamic budgeting scheme. It analyzes how attention and hidden-state information loss varies across layers and models, formalizes LMBA and LMBO as per-layer minimum budgets, and allocates budgets with a bounded minimum using LMBA-derived uncertainty. The method dynamically assigns per-layer budgets and selects KV entries via cumulative attention signals, achieving memory usage around ~20% of FullKV while preserving near-lossless performance on Needle-in-a-Haystack and LongBench across Mistral and LLaMa backbones. This approach offers practical benefits for scalable long-context reasoning without additional training or heavy computation, enabling more memory-efficient inference in real-world deployments.

Abstract

Large Language models (LLMs) have become a research hotspot. To accelerate the inference of LLMs, storing computed caches in memory has become the standard technique. However, as the inference length increases, growing KV caches might lead to out-of-memory issues. Many existing methods address this issue through KV cache compression, primarily by preserving key tokens throughout all layers to reduce information loss. Most of them allocate a uniform budget size for each layer to retain. However, we observe that the minimum budget sizes needed to retain essential information vary across layers and models based on the perspectives of attention and hidden state output. Building on this observation, this paper proposes a simple yet effective KV cache compression method that leverages layer uncertainty to allocate budget size for each layer. Experimental results show that the proposed method can reduce memory usage of the KV caches to only $\sim$20\% when compared to Full KV inference while achieving nearly lossless performance.

Figures (6)

• Figure 1: Comparison of the proposed method (ZigZagKV, Right) with previous PartialKV inference methods (Left). Most existing PartialKV methods allocate a uniform cache size per layer, whereas ZigZagKV dynamically adjusts the cache size based on layer uncertainty. "L0" and "L31" refer to Layer 0 and Layer 31, respectively. The numbers in brackets represent importance scores from the current token to prefix tokens, sorted from high to low.
• Figure 2: LMBA across various Layers of Mistral and LLaMa on 2WikiMQA dataset.
• Figure 3: LMBO across various Layers of Mistral and LLaMa on 2WikiMQA dataset.
• Figure 4: The evaluation results from Needle-in-a-HayStack testing across 96, 128, 256, 512 and 1024 budget sizes on Mistral-7B-Instruct-v0.3 and LLaMa-3.1-8B-Instruct. Proposed method ZigZagKV outperforms H2O, SnapKV, PyramidKV and StreamLM, especially in limited budget sizes.
• Figure 5: Performance comparison for Needle-in-a-Haystack testing with a budget size of 256 on Mistral-7B-Instruct-v0.3 and LLaMa-3.1-8B-Instruct. The x-axis represents the length of the document, while the y-axis indicates the position where the needle is located. A red cell indicates that the model fails to recall the information in the needle, whereas a green cell indicates success.