Table of Contents
Fetching ...

SmallKV: Small Model Assisted Compensation of KV Cache Compression for Efficient LLM Inference

Yi Zhao, Yajuan Peng, Cam-Tu Nguyen, Zuchao Li, Xiaoliang Wang, Hai Zhao, Xiaoming Fu

TL;DR

SmallKV tackles the memory- and compute-bottleneck of long-context LLM inference by addressing two main issues in KV-cache eviction: saliency shift and marginal information overcompression. It introduces a small-model assisted compensation framework that uses an SLM to preserve globally important attention patterns and to approximate the attention of marginal tokens, forming a hierarchical KV-cache compression strategy. Through extensive benchmarks (GSM8K, BBH, MT-Bench, LongBench) and model pairs (Qwen and LLaMA series), SmallKV achieves 1.75–2.56× higher throughput while maintaining accuracy under aggressive KV-cache budgets. The approach leverages attention similarity across model scales and remains compatible with Flash Attention and speculative decoding, enabling efficient and scalable LLM inference in resource-constrained environments.

Abstract

KV cache eviction has emerged as an effective solution to alleviate resource constraints faced by LLMs in long-context scenarios. However, existing token-level eviction methods often overlook two critical aspects: (1) their irreversible eviction strategy fails to adapt to dynamic attention patterns during decoding (the saliency shift problem), and (2) they treat both marginally important tokens and truly unimportant tokens equally, despite the collective significance of marginal tokens to model performance (the marginal information over-compression problem). To address these issues, we design two compensation mechanisms based on the high similarity of attention matrices between LLMs of different scales. We propose SmallKV, a small model assisted compensation method for KV cache compression. SmallKV can maintain attention matching between different-scale LLMs to: 1) assist the larger model in perceiving globally important information of attention; and 2) use the smaller model's attention scores to approximate those of marginal tokens in the larger model. Extensive experiments on benchmarks including GSM8K, BBH, MT-Bench, and LongBench demonstrate the effectiveness of SmallKV. Moreover, efficiency evaluations show that SmallKV achieves 1.75 - 2.56 times higher throughput than baseline methods, highlighting its potential for efficient and performant LLM inference in resource constrained environments.

SmallKV: Small Model Assisted Compensation of KV Cache Compression for Efficient LLM Inference

TL;DR

SmallKV tackles the memory- and compute-bottleneck of long-context LLM inference by addressing two main issues in KV-cache eviction: saliency shift and marginal information overcompression. It introduces a small-model assisted compensation framework that uses an SLM to preserve globally important attention patterns and to approximate the attention of marginal tokens, forming a hierarchical KV-cache compression strategy. Through extensive benchmarks (GSM8K, BBH, MT-Bench, LongBench) and model pairs (Qwen and LLaMA series), SmallKV achieves 1.75–2.56× higher throughput while maintaining accuracy under aggressive KV-cache budgets. The approach leverages attention similarity across model scales and remains compatible with Flash Attention and speculative decoding, enabling efficient and scalable LLM inference in resource-constrained environments.

Abstract

KV cache eviction has emerged as an effective solution to alleviate resource constraints faced by LLMs in long-context scenarios. However, existing token-level eviction methods often overlook two critical aspects: (1) their irreversible eviction strategy fails to adapt to dynamic attention patterns during decoding (the saliency shift problem), and (2) they treat both marginally important tokens and truly unimportant tokens equally, despite the collective significance of marginal tokens to model performance (the marginal information over-compression problem). To address these issues, we design two compensation mechanisms based on the high similarity of attention matrices between LLMs of different scales. We propose SmallKV, a small model assisted compensation method for KV cache compression. SmallKV can maintain attention matching between different-scale LLMs to: 1) assist the larger model in perceiving globally important information of attention; and 2) use the smaller model's attention scores to approximate those of marginal tokens in the larger model. Extensive experiments on benchmarks including GSM8K, BBH, MT-Bench, and LongBench demonstrate the effectiveness of SmallKV. Moreover, efficiency evaluations show that SmallKV achieves 1.75 - 2.56 times higher throughput than baseline methods, highlighting its potential for efficient and performant LLM inference in resource constrained environments.

Paper Structure

This paper contains 22 sections, 21 equations, 10 figures, 9 tables, 1 algorithm.

Figures (10)

  • Figure 1: (a) The illustration of saliency shift. (b) The quantification of saliency shift issue by measuring Jaccard similarity of important tokens between real drop and global view of critical tokens. Lower Jaccard similarity indicates that more important tokens have been wrongly evicted. (c) The visualization of highly consistent attention patterns between LLMs with different scaling, indicating that SLM can assist in capturing and preserving global attention information. (d) Comparison of 10% KV cache budget for saliency shift compensation using Qwen2-0.5B to assist Qwen2-7B.
  • Figure 2: Left: The distribution of attention sparsity (blue bar) along with the drop of model accuracy (red line), indicating the necessity of maintaining marginal tokens. Horizontal line (Full) represents the baseline of full KV cache. Right: Marginal information compensation by the hierarchical compression that differentiates tokens based on their varying levels of importance.
  • Figure 3: The attention illustration of the SmallKV method in decoding.
  • Figure 4: The system architecture of SmallKV. (The detail of SmallKV attention compensation is in Figure \ref{['fig:framework']})
  • Figure 5: Benchmark results of SmallKV with KV cache budget varying from 100% to 5%. Full represents using the full KV cache without compression.
  • ...and 5 more figures