Elamur Research⌕
☰
home / papers / 2024 / Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

Paper · 2024 · William Brandon, Mayank Mishra, Aniruddha Nrusimha et al.

Summary

When large language models generate text autoregressively, they store keys and values (the KV cache) for every token and layer. This memory grows with context length and batch size and often limits serving. Earlier fixes such as Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) shrink the cache by letting several query heads share one key/value head within a layer. Brandon et al. add a separate dimension of sharing, across depth. In their Cross-Layer Attention (CLA), only some layers compute key/value projections, and neighbouring layers reuse those activations. The sharing factor sets how many layers use each projection: CLA2 groups pairs, CLA3 groups triples. The authors pretrained 1B- and 3B-parameter models from scratch on SlimPajama, compared many CLA configurations against tuned MQA/GQA/MHA baselines and mapped the accuracy-versus-memory frontier. Combining CLA2 with MQA halves the KV cache compared with plain MQA, and perplexity barely changes. This is a Pareto improvement. It matters because the change is simple and low-risk, and it lets models serve longer contexts or larger batches on the same hardware.

Key points

  • Method: CLA can be combined with MHA, GQA or MQA. It shrinks the KV cache by the sharing factor, or slightly less when that factor does not divide the layer count evenly. It also slightly reduces parameter count and FLOPs.
  • Setup: The 1B and 3B models were trained from scratch on SlimPajama with the GPT-NeoX BPE tokenizer. Configurations were labelled by head size and attention type, for example H128-MQA, H64-MQA-CLA2 and H128-GQA4.
  • Main result: MQA-CLA2 gives about a 2× smaller KV cache than an MQA baseline with nearly equal validation perplexity. The benefit held at both 1B and 3B scale.
  • Configuration guidance:
    • A sharing factor of 2 worked better than larger ones.
    • CLA was most reliably useful together with MQA.
    • Preliminary evidence suggests CLA models benefit from higher learning rates.
  • Limitations:
    • CLA does not lower per-step attention memory bandwidth or core attention latency.
    • Under pipeline parallelism, layers that share a cache must sit in the same stage, or KV activations must be sent between stages.
    • The paper does not evaluate end-to-end inference efficiency; the authors leave it for future work.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

🧭 Topics 1Research subjects and languages the work is about.Efficient models & compression · 328 works

Cite

Cite as William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda and Jonathan Ragan Kelly. 2024. Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. arXiv:2405.12981.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2405.12981 · arXiv 2405.12981

@misc{brandon2024reducing,
  title = {{Reducing Transformer Key-Value Cache Size with Cross-Layer Attention}},
  author = {William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda and Jonathan Ragan Kelly},
  year = {2024},
  doi = {10.48550/arxiv.2405.12981},
  eprint = {2405.12981},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2405.12981}
}

References

49 references, 20 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebron, and S. Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing,…
  2. J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023b.
  3. J. Ansel, E. Yang, H. He, N. Gimelshein, A. Jain, M. Voznesensky, B. Bao, P. Bell, D. Berard, E. Burovski, G. Chauhan, A. Chourdia, W. Constable, A. Desmaison, Z. DeVito, E. Ellison, W. Feng, J. Gong, M. Gschwind, B. Hirsh, S. Huang, K. Kalambarkar, L. Kirsch, M. Lazos, M. Lezcano, Y. Liang, J. Lian… — PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation
  4. S. Black, S. Biderman, E. Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, M. Pieler, U. S. Prashanth, S. Purohit, L. Reynolds, J. Tow, B. Wang, and S. Weinbach. Gpt-neox-20b: An open-source autoregressive language model, 2022.
  5. S. Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, G. B. Van Den Driessche, J.-B. Lespiau, B. Damoc, A. Clark, D. De Las Casas, A. Guy, J. Menick, R. Ring, T. Hennigan, S. Huang, L. Maggiore, C. Jones, A. Cassirer, A. Brock, M. Paganini, G. Irving, O. Vinyals, S. Osindero, K. S… find
  6. R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509. ↗
  7. A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-… — PaLM: Scaling Language Modeling with Pathways
  8. Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. Le, and R. Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In A. Korhonen, D. Traum, and L. Màrquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978–2988, F… find
  9. T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023.
  10. T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  11. DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024. find
  12. B. Gao, Z. He, P. Sharma, Q. Kang, D. Jevdjic, J. Deng, X. Yang, Z. Yu, and P. Zuo. Attentionstore: Cost-effective attention reuse across multi-turn conversations in large language model serving, 2024. find
  13. L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. A framework for few-shot language m… find
  14. S. Ge, Y. Zhang, L. Liu, M. Zhang, J. Han, and J. Gao. Model tells you what to discard: Adaptive kv cache compression for llms, 2024.
  15. Google. Context caching guide. https://ai.google.dev/gemini-api/docs/caching, 2024. Accessed: 2024-05-20. ↗
  16. A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023. find
  17. K. Guu, K. Lee, Z. Tung, P. Pasupat, and M.-W. Chang. Realm: retrieval-augmented language model pre-training. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020.
  18. C. Hooper, S. Kim, H. Mohammadzadeh, M. W. Mahoney, Y. S. Shao, K. Keutzer, and A. Gholami. Kvquant: Towards 10 million context length llm inference with kv cache quantization, 2024.
  19. Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. X. Chen, D. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and Z. Chen. GPipe: efficient training of giant neural networks using pipeline parallelism. Curran Associates Inc., Red Hook, NY, USA, 2019. find
  20. G. Izacard, P. Lewis, M. Lomeli, L. Hosseini, F. Petroni, T. Schick, J. Dwivedi-Yu, A. Joulin, S. Riedel, and E. Grave. Atlas: few-shot learning with retrieval augmented language models. J. Mach. Learn. Res., 24(1), mar 2024. ISSN 1532-4435.
  21. D. Kalamkar, D. Mudigere, N. Mellempudi, D. Das, K. Banerjee, S. Avancha, D. T. Vooturi, N. Jammalamadaka, J. Huang, H. Yuen, J. Yang, J. Park, A. Heinecke, E. Georganas, S. Srinivasan, A. Kundu, M. Smelyanskiy, B. Kaul, and P. Dubey. A study of bfloat16 for deep learning training, 2019. find
  22. A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020. find
  23. Y. Liu, H. Li, Y. Cheng, S. Ray, Y. Huang, Q. Zhang, K. Du, J. Yao, S. Lu, G. Ananthanarayanan, M. Maire, H. Hoffmann, A. Holtzman, and J. Jiang. Cachegen: Kv cache compression and streaming for fast language model serving, 2024. find
  24. Z. Liu, A. Desai, F. Liao, W. Wang, V. Xie, Z. Xu, A. Kyrillidis, and A. Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time, 2023. find
  25. I. Loshchilov and F. Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id= Skq89Scxx. ↗
  26. I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=Bkg6RiCqY7. ↗
  27. S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models, 2016.
  28. P. Micikevicius, S. Narang, J. Alben, G. Diamos, E. Elsen, D. Garcia, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu. Mixed precision training, 2018. find
  29. A. Mohtashami and M. Jaggi. Random-access infinite context length for transformers. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 54567–54585. Curran Associates, Inc., 2023. URL https://proceedings… find
  30. T. Munkhdalai, M. Faruqui, and S. Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention, 2024.
  31. A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
  32. B. Peng, D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, T. Ferdinan, H. Hou, P. Kazienko, K. K. GV, J. Kocoń, B. Koptyra, S. Krishna, R. M. J. au2, N. Muennighoff, F. Obeid, A. Saito, G. Song, H. Tu, S. Woźniak, R. Zhang, B. Zhao, Q. Zhao, P. Zhou, J. Zhu, and R.-J… find
  33. P. Ramachandran, B. Zoph, and Q. V. Le. Searching for activation functions, 2017. find
  34. N. Shazeer. Fast transformer decoding: One write-head is all you need, 2019. find
  35. N. Shazeer. Glu variants improve transformer, 2020. find
  36. Y. Sheng, L. Zheng, B. Yuan, Z. Li, M. Ryabinin, B. Chen, P. Liang, C. Ré, I. Stoica, and C. Zhang. Flexgen: high-throughput generative inference of large language models with a single gpu. In Proceedings of the 40th International Conference on Machine Learning, ICML’23. JMLR.org, 2023.
  37. M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020. find
  38. D. Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://www.cerebras.net/blog/ slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama, 2023. URL https://huggingface.co/datasets/cerebr… ↗
  39. J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
  40. Y. Sun, L. Dong, Y. Zhu, S. Huang, W. Wang, S. Ma, Q. Zhang, J. Wang, and F. Wei. You only cache once: Decoder-decoder architectures for language models, 2024. find
  41. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan,… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  42. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_ files/paper/2017/file/3f5ee243… ↗
  43. C. Wang, K. Cho, and J. Gu. Neural machine translation with byte-level subwords, 2019.
  44. S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768, 2020. URL https://arxiv.org/abs/2006.04768. ↗
  45. Y. Wu, M. N. Rabe, D. Hutchins, and C. Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id= TrjbxzRcnf-. ↗
  46. S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efficient training, 2024. find
  47. T. Zhang, J. Yi, Z. Xu, and A. Shrivastava. Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization, 2024. find
  48. Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. A. Wang, and B. Chen. H2o: Heavy-hitter oracle for efficient generative inference of large language models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu… find
  49. X. Zhu, J. Li, Y. Liu, C. Ma, and W. Wang. A survey on model compression for large language models, 2023.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/brandon-2024-reducing-transformer-key-value-cache/