A Systematic Study of Cross-Layer KV Sharing for Efficient LLM Inference
A unified pizza–sandwich–lasagna framework shows that 2× KV-cache reduction usually preserves quality, while aggressive compression favors sandwich-middle.
Key Findings
Methodology
The paper defines kv(i), the layer whose K and V are paired with layer i's queries; kv(i)=i denotes a KV layer, while other layers reuse a target layer's cache. It evaluates three partitions—pizza, sandwich, and lasagna—and three target positions—bottom, top, and middle—yielding nine configurations. LCKV, YOCO, and CLA appear as special cases. Top and middle variants use iterative training and prefilling with m=7 and b=2.
Key Results
- On a 1.1B model with an RTX 3090, all nine configurations substantially outperform standard Llama for a short 5+2043 workload, and throughput rises as KV layers decrease. For 512+1024, iterative prefilling makes top and middle variants deteriorate sharply, with some falling below the baseline; bottom variants remain clearly faster.
- Using MiniPile's 1.7B tokens, the 110M model trained for one epoch and the 1.1B model for two. When half the layers remain KV layers, most configurations achieve perplexity comparable to the standard Transformer; with fewer KV layers, bottom variants degrade most, while top and middle variants are more robust.
- After training on 100B SlimPajama tokens, evaluation covered HellaSwag, OpenBookQA, WinoGrande, ARC-Easy, ARC-Challenge, BoolQ, PIQA, and SciQ. Sandwich-top generally beats bottom variants, whereas sandwich-middle is strongest when very few KV layers are retained.
Significance
This study turns fragmented cross-layer KV methods into a common design space and exposes the trade-off among cache capacity, target-layer placement, prefilling cost, and task quality. It gives practitioners actionable guidance: choose pizza-bottom or lasagna-bottom for approximately 2× compression without extra training or prefilling, and consider sandwich-middle when aggressive compression is needed and generated text is much longer than the prompt.
Technical Contribution
The main contribution is a formal kv(i) mapping combined with the pizza/sandwich/lasagna and bottom/top/middle axes, including previously unexplored lasagna-top and middle variants. To handle cyclic dependencies in top and middle designs, the authors use iterative approximation: the first n-b iterations are approximated with m steps, while only the final b iterations receive backpropagation. The same mechanism supports prompt prefilling, and KV-layer count directly controls WK/WV parameters and cache size.
Novelty
The novelty is systematic rather than a single new cache operator. The work provides a unified empirical comparison of major cross-layer KV strategies and new variants, mapping LCKV, YOCO, and CLA to sandwich-top, pizza-bottom, and lasagna-bottom respectively. It is especially valuable for revealing how target placement changes the speed–quality frontier under different prompt lengths.
Limitations
- The study evaluates only 110M and 1.1B models and at most 100B training tokens; rankings may change for larger models, longer contexts, or different data distributions.
- Top and middle variants incur extra training and prefilling cost and mask diagonal attention to break cyclic dependencies, which may create a quality gap relative to standard Transformers.
Future Work
Future work should test larger models, longer contexts, diverse hardware, and realistic serving workloads while reducing iterative prefilling cost. The authors specifically suggest separately computing each token's self-attention to recover the information lost by diagonal masking. Other directions include improved initialization, adaptive target placement, and more efficient iterative training.
AI Executive Summary
Large language models often bottleneck on the key–value cache used during autoregressive decoding. Every layer stores representations for every previous token, making memory capacity and bandwidth limiting factors in long-context and high-batch serving. Existing approaches such as GQA, H2O, StreamingLLM, LCKV, YOCO, and CLA reduce this burden, but they make different choices about which layers compute and share KVs. Their trade-offs had not been compared systematically.
Wu, Wu, and Tu introduce a unified mapping kv(i), specifying which layer supplies K and V to layer i's queries. Combining pizza, sandwich, and lasagna partitions with bottom, top, and middle target positions produces nine configurations, including LCKV, YOCO, CLA, and new variants. Bottom designs train conventionally; top and middle designs use iterative training and prefilling with m=7 and b=2 to approximate cyclic dependencies. On a 1.1B model and RTX 3090, all nine configurations outperform standard Llama on a short 5+2043 workload, while top and middle variants lose their advantage on 512+1024 because prompt encoding becomes iterative.
Quality experiments use 1.7B MiniPile tokens and 100B SlimPajama tokens, followed by perplexity and zero-shot tests on HellaSwag, OpenBookQA, WinoGrande, ARC, BoolQ, PIQA, and SciQ. Retaining half the KV layers usually preserves baseline quality; more aggressive compression hurts bottom variants most, while sandwich-middle becomes attractive when KV layers are scarce. The practical message is clear: use pizza-bottom or lasagna-bottom for low-risk 2× compression, and consider sandwich-middle when memory dominates latency and generated sequences greatly exceed prompts.
Deep Analysis
Background
During decoding, a Transformer stores K and V for every historical token at every layer. This creates severe memory and bandwidth pressure as context length and batch size grow. GQA/MQA reduce sharing across attention heads, while H2O and StreamingLLM reduce token retention. LCKV, YOCO, and CLA instead share KVs across layers, but differ in layer placement and dependency structure, motivating a unified comparison.
Core Problem
The central question is which layer partition and target position best balance cache reduction, throughput, language-modeling quality, and downstream accuracy for a fixed number of KV layers. The problem is difficult because top and middle sharing create sequential dependencies between lower queries and later-computed target KVs. Moreover, prompt length changes whether decoding savings outweigh iterative prefilling overhead.
Innovation
The authors define kv(i) and enumerate nine designs. Pizza keeps the first l−1 layers as KV layers; sandwich keeps KV layers at both ends; lasagna uniformly groups all L layers. Each group uses a bottom, top, or middle target, except the first lasagna group, whose bottom layer is always a KV target because making the first layer non-KV causes a substantial quality drop. Novel lasagna-top and middle configurations extend prior work.
Methodology
- �� Input: an L-layer Transformer and a budget of l KV layers; output: the kv(i) mapping.
- �� KV layers compute their own K and V. Non-KV layers remove WK and WV and reuse the target layer's cache, reducing parameters and memory.
- �� Bottom variants use standard parallel training. Top and middle variants use iterative training: m approximation iterations for the early n−b steps and backpropagation only through the final b steps.
- �� During prefilling, early exit is possible after the last KV layer. Top and middle layouts use m+b iterative encoding steps; diagonal self-attention is masked to break cyclic dependencies.
- �� Implementation uses HuggingFace Transformers, FlashAttention 2, fused RMSNorm, fused cross-entropy, and fused SwiGLU.
Experiments
The models contain 110M and 1.1B parameters, with 12 and 22 layers respectively. The 1.1B model has hidden size 2048, 32 attention heads, and 4 KV heads. Throughput is measured on an RTX 3090 using short 5+2043 and long 512+1024 prompt-plus-generation settings. MiniPile supplies 1.7B training tokens; SlimPajama supplies a 100B-token subset. Metrics are maximum generation throughput, perplexity, and LM Eval Harness zero-shot accuracy.
Results
For short prompts, every configuration is faster than standard Llama and fewer KV layers increase throughput. For long prompts, bottom layouts remain advantageous, whereas top and middle layouts suffer from iterative prompt encoding. Half-layer KV budgets generally preserve baseline perplexity; further reduction damages bottom variants most. On the large-data setting, sandwich-top is usually stronger than bottom layouts, while sandwich-middle performs best when the KV budget is very small.
Applications
The techniques target GPU-memory-limited chatbots, long-form generation, batch inference, and edge deployment. Bottom layouts fit short prompts and throughput-sensitive serving; sandwich-middle may suit workloads whose outputs greatly exceed inputs and can tolerate extra prompt processing. Deployment requires benchmarking prompt length, generation length, batch size, hardware, and retraining availability.
Limitations & Outlook
Only 110M and 1.1B models are studied, so larger-scale behavior is unknown. Top and middle designs add iterative training and prefilling overhead, and diagonal masking may reduce quality. Throughput is implementation- and hardware-dependent, particularly on the RTX 3090 setup. Future work should evaluate longer contexts, larger models, pretrained initialization, adaptive layouts, and explicit self-attention compensation.
Plain Language Accessible to non-experts
Imagine a factory producing custom products. In a standard factory, every workshop keeps its own full filing cabinet containing everything that happened earlier. This makes the factory reliable, but the cabinets consume enormous floor space and workers spend time opening them.
Cross-layer KV sharing lets several workshops share one cabinet. Pizza places most cabinets near the front, sandwich places them at both ends, and lasagna divides the factory into sections with one cabinet per section. Fewer cabinets save space and can speed production, but a workshop may have to wait for a cabinet located later in the production line. That waiting is why top and middle layouts struggle when the initial order is long.
The experiments show that short orders benefit from almost every layout. For long orders, bottom layouts are safer because the needed cabinets are ready without extra preparation. If a company only wants to halve storage, pizza-bottom and lasagna-bottom are practical choices. If it wants much more compression and can tolerate a slower setup phase, sandwich-middle may offer the better compromise.
ELI14 Explained like you're 14
Think of a giant video game. Every time your character moves, the game remembers earlier maps, enemies, items, and decisions. The normal system keeps a full memory pack for every game level, so long adventures use lots of computer memory.
This paper asks: what if several levels shared one memory pack? The researchers tried putting the pack near the beginning, the end, or the middle, and also tried splitting the game into sections. Less memory can make the game run faster—but if a level needs a pack that has not been prepared yet, everything waits. That is the main problem for some top and middle designs.
With a short opening scene and a long adventure, almost every design beat the normal system. With a long opening scene, bottom designs worked better because they avoided repeated preparation. Keeping about half the memory layers usually preserved answer quality; deleting too many made bottom designs weaker, while sandwich-middle stayed competitive.
So which one wins? There is no universal champion. For a safe 2× memory reduction, pizza-bottom or lasagna-bottom is a good pick. For extreme compression, and when you can accept extra setup time, sandwich-middle is promising. The next challenge is testing these tricks on much larger games—or, scientifically speaking, larger language models!
Glossary
KV cache
A memory of attention Keys and Values from previous tokens, allowing autoregressive decoding to avoid recomputing history. Its size grows with layers, context length, and KV heads.
The paper reduces it by computing KVs in only selected layers and sharing them across layers.
Cross-layer KV sharing
A mechanism in which multiple layers' queries use K and V representations produced by one target layer. Non-KV layers can therefore remove their WK and WV projections.
It is the common principle behind LCKV, YOCO, CLA, and all nine proposed configurations.
LCKV
Layer-Condensed KV Cache computes KVs only at the top layer and pairs them with queries from other layers. It strongly compresses the cache but introduces dependency and training considerations.
It corresponds to the paper's sandwich-top configuration.
YOCO
You Only Cache Once computes KVs at a middle layer and supplies them to upper-layer queries, while its original design uses efficient attention below. It targets very small cache growth.
The paper's pizza-bottom is close in sharing pattern but uses standard attention.
CLA
Cross-Layer Attention divides layers into adjacent groups and makes queries within each group use the bottom layer's KVs. It offers a regular and implementation-friendly sharing pattern.
It corresponds to lasagna-bottom in the unified framework.
Prefilling
The stage that processes the user prompt and constructs the KV cache before token-by-token decoding. Long prompts make its latency increasingly important.
Top and middle layouts lose throughput on long prompts because prefilling becomes iterative.
Open Questions Unanswered questions from this research
- 1 It remains unknown whether the observed ranking survives at much larger scales, longer contexts, and different hardware. Broad serving benchmarks are needed because memory bandwidth and kernel behavior may alter the trade-off.
- 2 No low-cost method fully removes the cyclic dependencies of top and middle sharing without quality loss. Better fixed-point solvers, self-attention correction, or adaptive placement may be required.
Applications
Immediate Applications
Memory-constrained batch generation
Serving teams can test pizza-bottom or lasagna-bottom on roughly 1.1B-scale models to target about 2× KV-cache reduction while preserving quality and improving short-prompt throughput. Actual deployment should remeasure with its own batch sizes, kernels, and GPU memory limits.
Long-output conversational systems
When generated text is much longer than the prompt, sandwich-middle is a candidate for aggressive cache reduction. It requires additional training and iterative prefilling, but its quality is comparatively robust when very few KV layers are retained.
Long-term Vision
Adaptive KV-layout serving
A future inference engine could choose bottom, top, or middle sharing dynamically from prompt length, expected output length, batch size, and memory pressure. Such routing could optimize latency, capacity, and quality per request instead of using one fixed layout.
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.