Reconstructing KV Caches with Cross-Layer Fusion for Enhanced Transformers
Hongzhan Lin, Zhiqi Bai, Xinmiao Zhang, Siran Yang, Jiamang Wang, Yunlong Xu, Jiaheng Liu, Yongchi Zhao, Xiang Li, Yuchi Xu, Wenbo Su, Bo Zheng
摘要
Transformer decoders have achieved strong results across tasks, but the memory required for the KV cache becomes prohibitive at long sequence lengths. Although Cross-layer KV Cache sharing (e.g., YOCO, CLA) offers a path to mitigate KV Cache bottleneck, it typically underperforms within-layer methods like GQA. To understand the root cause, we investigate the information flow of keys and values of the top-layers. Our preliminary reveals a clear distribution: values are predominantly derived from the bottom layer, while keys draw more information from both bottom and middle layers. Building upon this, we propose FusedKV, whose top-layer KV caches are a learnable fusion of the most informative ones from the bottom and middle layers. This fusion operates directly on post-RoPE keys, preserving relative positional information without the computational cost of re-applying rotary embeddings. To further improve efficiency, we propose FusedKV-Lite, an cross-layer sharing approach, where top-layer KV caches are directly derived from the bottom-layer values and the middle-layer keys. Compared to FusedKV, FusedKV-Lite reduces I/O overhead at the cost of a slight increase in perplexity. In experiments on LLMs ranging from 332M to 4B parameters, our proposed method reduce 50% cache memory while achieving lower validation perplexity than the standard Transformer decoder, establishing it as a memory-efficient, high-performance architectural alternative. We have made our Triton implementation available.
问问这篇 Paper
智能体会读完全文。
Lune 把这篇 Paper 索引到了每一个公式,引用它的顶会 Paper 也一样。你提问,回答直接引用原文。
它引用的顶会 Paper16
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou 等ICLR 2021 · 被引用 7,905 次
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra 等NeurIPS 2022 · 被引用 5,493 次
- A Simple and Effective Pruning Approach for Large Language ModelsMingjie Sun, Zhuang Liu, Anna Bair, J. Zico KolterICLR 2024 · 被引用 794 次
- KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV CacheZirui Liu, Jiayi Yuan, Hongye Jin, Shaochen (Henry) Zhong 等ICML 2024 · 被引用 436 次
- Anti-Oversmoothing in Deep Vision Transformers via the Fourier Domain Analysis: From Theory to PracticePeihao Wang, Wenqing Zheng, Tianlong Chen, Zhangyang WangICLR 2022 · 被引用 212 次
相关 Paper
- Reducing Transformer Key-Value Cache Size with Cross-Layer AttentionWilliam Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda 等NeurIPS 2024 · 被引用 150 次
- Improving Model Representation and Reducing KV Cache via Skip Connections with First Value HeadsZhoutong Wu, Yuan Zhang, Yiming Dong, Chenheng Zhang 等NeurIPS 2025 · 被引用 4 次
- You Only Cache Once: Decoder-Decoder Architectures for Language ModelsYutao Sun, Li Dong, Yi Zhu, Shaohan Huang 等NeurIPS 2024 · 被引用 162 次
- KV-Latent: Dimensional-level KV Cache Reduction with Frequency-aware Rotary Positional EmbeddingLuohe Shi, Zuchao Li, Lefei Zhang, Baoyuan Qi 等ACL 2025 · 被引用 5 次
- ContrastKV: Robust KV Cache Eviction via Contrastive Signal Fusion for Multi-Query GeneralizationXingchi Chen, Peiyuan Zong, Ziqiang Gao, Qing Li 等ACL 2026
