ZSMerge: Zero-Shot KV Cache Compression for Memory-Efficient Long-Context LLMs

TL;DR

ZSMerge tackles the memory and compute bottlenecks of long-context LLMs by introducing a zero-shot KV cache compression framework that combines fine-grained head-level budget allocation, residual merging, and compensated attention scoring. The method achieves sublinear KV cache growth while preserving generation quality, delivering substantial memory savings (e.g., ~82% VRAM reduction at 54K tokens) and throughput gains (over threefold at extreme contexts) without retraining. Empirical results across multiple models (LLaMA2-7B, Falcon-7B, Mistral-7B-Instruct) and benchmarks (LongBench, InfiniteBench, XSum) show ZSMerge outperforms eviction-based baselines and generalizes across architectures and tasks, maintaining robust performance under tight cache budgets. The approach enables scalable long-context inference on resource-constrained devices, reduces energy consumption, and supports broad deployment without architectural changes or fine-tuning.

Abstract

The linear growth of key-value (KV) cache memory and quadratic computational in attention mechanisms complexity pose significant bottlenecks for large language models (LLMs) in long-context processing. While existing KV cache optimization methods address these challenges through token pruning or feature merging, they often incur irreversible information loss or require costly parameter retraining. To this end, we propose ZSMerge, a dynamic KV cache compression framework designed for efficient cache management, featuring three key operations: (1) fine-grained memory allocation guided by multi-dimensional token importance metrics at head-level granularity, (2) a residual merging mechanism that preserves critical context through compensated attention scoring, and (3) a zero-shot adaptation mechanism compatible with diverse LLM architectures without requiring retraining. ZSMerge significantly enhances memory efficiency and inference speed with negligible performance degradation across LLMs. When applied to LLaMA2-7B, it demonstrates a 20:1 compression ratio for key-value cache retention (reducing memory footprint to 5\% of baseline) while sustaining comparable generation quality, coupled with triple throughput gains at extreme 54k-token contexts that eliminate out-of-memory failures. The code is available at https://github.com/SusCom-Lab/ZSMerge.

Figures (5)

• Figure 1: Comparative analysis of contextually adaptive KV cache optimization strategies. (A) Sparse Approximation (Eviction): Low-contribution tokens are permanently discarded, reducing memory but risking error propagation. (B) Token Merging: Similar tokens are fused to mitigate information loss, but may require fine-tuning or auxiliary networks. (C) ZSMerge (Proposed): Achieves zero-shot "sparse+residual" compression without added parameters or tuning, balancing efficiency and performance.
• Figure 2: ZSMerge Online Compression
• Figure 3: VRAM Utilization and Decoding Throughput Across Sequence Lengths. ZSMerge enforces constant memory footprint (43GB) and sustains 9 tokens/sec decoding rate beyond 54K tokens, eliminating out-of-memory (OOM) via dynamic KV cache compression.
• Figure 4: Numerical error analysis comparing eviction vs. residual merging across compression ratios. The log-scale visualization demonstrates that ZSMerge's residual merging consistently reduces approximation error relative to pure eviction, with benefits amplified under aggressive compression ($\le 20\%$ cache size).
• Figure 5: Hyperparameter Sensitivity Analysis: (top) Proximity Maintenance Ratio ($B_p/B$), (middle) Residual Budget Ratio ($B_r/(B-B_p)$), (bottom) Scale Factor ($\alpha$)

Theorems & Definitions (2)

• Theorem 1
• proof : Proof of Theorem \ref{['thm:main']}