A Systematic Study of Cross-Layer KV Sharing for Efficient LLM Inference

TL;DR

This work tackles the KV cache memory bottleneck in LLM inference by introducing a unified cross-layer KV sharing framework that subsumes prior methods such as LCKV, YOCO, and CLA into nine configurations defined by pizza, sandwich, and lasagna partitionings with different target-layer positions. The authors show that reducing the KV cache by $2\times$ can boost generation throughput beyond standard transformers while maintaining competitive perplexity, though further reductions shift the balance toward configurations that share KVs from upper layers at the cost of training and prefilling overhead. They provide iterative training and diagonal masking techniques to enable efficient training and inference for top/middle configurations, and they evaluate throughput, perplexity, and downstream task performance across 110M and 1.1B parameter models on diverse datasets. Key findings indicate that bottom configurations (e.g., pizza-bottom, lasagna-bottom) offer strong throughput with modest performance loss at $2\times$ KV reduction, while sandwich-middle shines when pursuing further throughput gains at higher training/prefilling costs. The work offers practical guidance for choosing cross-layer KV sharing patterns based on desired throughput and task requirements and provides code for reproducibility.

Abstract

Recently, sharing key-value (KV) cache across layers has been found effective in efficient inference of large language models (LLMs). To systematically investigate different techniques of cross-layer KV sharing, we propose a unified framework that covers several recent methods and their novel variants. We conduct comprehensive experiments on all the configurations of the framework, evaluating their generation throughput and performance in language modeling and downstream tasks. We find that when reducing the size of the KV cache by 2$\times$, most configurations can achieve higher throughput than standard transformers while maintaining competitive performance. When further reducing the size of the KV cache, however, pairing queries of all layers with KVs of upper layers performs better, at the expense of additional training cost and prefilling latency. We hope that this work will help users make more informed choices of cross-layer KV sharing approaches and facilitate future research on efficient LLM inference.

Figures (4)

• Figure 1: Experimental results.
• Figure 2: Throughput of 1.1B models at different batch sizes on an RTX 3090 (24GB) GPU with a prompt length of 5 and a generation length of 2043.
• Figure 3: Perplexity on the Minipile dataset of 1.1B models initialized with converted Tinyllama-2.5T weights.
• Figure 4: Perplexity on the Minipile dataset of 1.1B models with more options for target layer positioning.