CoChat
Sign Up
PreprintCited by 9Topic ModelingNatural Language Processing TechniquesAdvanced Neural Network Applications

Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

Open access
Adnan, Muhammad, Arunkumar, Akhil, Jain, Gaurav, Nair, Prashant J., Soloveychik, Ilya, Kamath, Purushotham·2024·arXiv
DOI 10.48550/arxiv.2403.09054arXiv 2403.09054
Open this paper in CoChat — free accountOpen PDF

Abstract

Transformers have emerged as the underpinning architecture for Large Language Models (LLMs). In generative language models, the inference process involves two primary phases: prompt processing and token generation. Token generation, which constitutes the majority of the computational workload, primarily entails vector-matrix multiplications and interactions with the Key-Value (KV) Cache. This phase is constrained by memory bandwidth due to the overhead of transferring weights and KV cache values from the memory system to the computing units. This memory bottleneck becomes particularly pronounced in applications that require long-context and extensive text generation, both of which are increasingly crucial for LLMs. This paper introduces "Keyformer", an innovative inference-time approach, to mitigate the challenges associated with KV cache size and memory bandwidth utilization. Keyformer leverages the observation that approximately 90% of the attention weight in generative inference focuses on a specific subset of tokens, referred to as "key" tokens. Keyformer retains only the key tokens in the KV cache by identifying these crucial tokens using a novel score function. This approach effectively reduces both the KV cache size and memory bandwidth usage without compromising model accuracy. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT, which employ various positional embedding algorithms. Our assessment encompasses a variety of tasks, with a particular emphasis on summarization and conversation tasks involving extended contexts. Keyformer's reduction of KV cache reduces inference latency by 2.1x and improves token generation throughput by 2.4x, while preserving the model's accuracy.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2403.09054

Cites

1 of 1 references verified0 citation errors · 0 not checkable

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

  • Rethinking Attention with Performers

    Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. (2020) · arXiv · doi:10.48550/arxiv.2009.14794 · Cited by 695 · cited 1× · see citing works

    “As the attention mechanism scales quadratically with the size of the input sequence, it tends to present the largest latency overhead during inference (Sukhbaatar et al., 2019; Dao et al., 2022; Choromanski et al., 2020).”
    Background

Flagged first, then by citations

Research works printed but not linked: 61

  1. Adnan, M., Maboud, Y. E., Mahajan, D., and Nair, P. J. Ac- celerating recommendation system training by leveraging popular choices. Proc. VLDB Endow., 15(1):127-140, sep 2021. ISSN 2150-8097.
  2. Adnan, M., Maboud, Y. E., Mahajan, D., and Nair, P. J. Ad- rec: Advanced feature interactions to address covariate- shifts in recommendation networks. arXiv preprint arXiv:2308.14902, 2023. · arXiv
  3. Adnan, M., Maboud, Y. E., Mahajan, D., and Nair, P. J. Heterogeneous acceleration pipeline for recommenda- tion system training. In Proceedings of the 51st Inter- national Symposium on Computer Architecture (ISCA). ACM, 2024.
  4. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305.13245, 2023. · arXiv
  5. Anagnostidis, S., Pavllo, D., Biggio, L., Noci, L., Lucchi, A., and Hoffmann, T. Dynamic context pruning for effi- cient and interpretable autoregressive transformers. arXiv preprint arXiv:2305.15805, 2023. · arXiv
  6. Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. Longbench: A bilingual, multitask benchmark for long context under- standing. arXiv preprint arXiv:2308.14508, 2023. · arXiv
  7. Beltagy, I., Peters, M. E., and Cohan, A. Long- former: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. · arXiv
  8. Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
  9. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
  10. Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context large language models. arXiv:2309.12307, 2023. · arXiv
  11. Child, R., Gray, S., Radford, A., and Sutskever, I. Gen- erating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019. · arXiv
  12. Cooray, K. Generalized gumbel distribution. Journal of Applied Statistics, 37(1):171-179, 2010.
  13. Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Process- ing Systems, 35:16344-16359, 2022.
  14. de Souza Pereira Moreira, G., Rabhi, S., Lee, J. M., Ak, R., and Oldridge, E. Transformers4rec: Bridging the gap between nlp and sequential/session-based recommenda- tion. In Proceedings of the 15th ACM Conference on Recommender Systems, pp. 143-153, 2021.
  15. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020. · arXiv
  16. Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., and Gao, J. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801, 2024. · arXiv
  17. Geng, S., Yuan, J., Tian, Y., Chen, Y., and Zhang, Y. Hiclip: Contrastive language-image pretraining with hierarchy- aware attention. arXiv preprint arXiv:2303.02995, 2023. · arXiv
  18. Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V., Sabharwal, Y., and Verma, A. Power-bert: Accelerating bert inference via progressive word-vector elimination. In International Conference on Machine Learning, pp. 3690-3699. PMLR, 2020.
  19. Huang, L., Cao, S., Parulian, N., Ji, H., and Wang, L. Effi- cient attentions for long document summarization. arXiv preprint arXiv:2104.02112, 2021a. · arXiv
  20. Huang, L., Cao, S., Parulian, N., Ji, H., and Wang, L. Effi- cient attentions for long document summarization. In Pro- ceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1419-1436, Online, June 2021b. Association for Computational Linguistics. doi: 10.18653/v1/2021.naacl-main.112. · DOI · Paper page
  21. Jang, E., Gu, S., and Poole, B. Categorical repa- rameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016. · arXiv
  22. Kang, H., Zhang, Q., Kundu, S., Jeong, G., Liu, Z., Krishna, T., and Zhao, T. Gear: An efficient kv cache compression recipe for near-lossless generative inference of llm. arXiv preprint arXiv:2403.05527, 2024. · arXiv
  23. Kao, S.-C., Subramanian, S., Agrawal, G., Yazdanbakhsh, A., and Krishna, T. Flat: An optimized dataflow for mitigating attention bottlenecks. In Proceedings of the 28th ACM International Conference on Architectural Sup- port for Programming Languages and Operating Systems, Volume 2, pp. 295-310, 2023.
  24. Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
  25. Kim, H., Hessel, J., Jiang, L., West, P., Lu, X., Yu, Y., Zhou, P., Bras, R. L., Alikhani, M., Kim, G., Sap, M., and Choi, Y. Soda: Million-scale dialogue distillation with social commonsense contextualization. ArXiv, abs/2212.10465, 2022. · arXiv
  26. Kitaev, N. Kaiser, łl., levskaya, a.: Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020. · arXiv
  27. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. arXiv preprint arXiv:2309.06180, 2023. · arXiv
  28. Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mo- hamed, A., Levy, O., Stoyanov, V., and Zettlemoyer, L. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehen- sion. arXiv preprint arXiv:1910.13461, 2019. · arXiv
  29. Li, D., Shao, R., Xie, A., Sheng, Y., Zheng, L., Gonzalez, J., Stoica, I., Ma, X., and Zhang, H. How long can context length of open-source llms truly promise? In NeurIPS 2023 Workshop on Instruction Tuning and Instruction Following, 2023.
  30. Lin, C.-Y. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
  31. Liu, Z., Wang, J., Dao, T., Zhou, T., Yuan, B., Song, Z., Shrivastava, A., Zhang, C., Tian, Y., Re, C., et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, pp. 22137-22176.
  32. PMLR, 2023.
  33. Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712, 2016. · arXiv
  34. Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789, 2018. · arXiv
  35. Mohtashami, A. and Jaggi, M. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305.16300, 2023. · arXiv
  36. Mu, J., Li, X. L., and Goodman, N. Learning to compress prompts with gist tokens. arXiv preprint arXiv:2304.08467, 2023. · arXiv
  37. Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M. fairseq: A fast, exten- sible toolkit for sequence modeling. In Ammar, W., Louis, A., and Mostafazadeh, N. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demon- strations), pp. 48-53, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-4009. · DOI
  38. Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023.
  39. Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021. · arXiv
  40. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5551, 2020.
  41. Roemmele, M., Bejan, C. A., and Gordon, A. S. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011.
  42. Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
  43. See, A., Liu, P. J., and Manning, C. D. Get to the point: Summarization with pointer-generator networks. In Pro- ceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1073-1083, Vancouver, Canada, July 2017. Associ- ation for Computational Linguistics. doi: 10.18653/v1/ P17-1099. · DOI · Paper page
  44. Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019. · arXiv
  45. Sheng, Y., Zheng, L., Yuan, B., Li, Z., Ryabinin, M., Chen, B., Liang, P., Ré, C., Stoica, I., and Zhang, C. Flexgen: High-throughput generative inference of large language models with a single gpu. In International Conference on Machine Learning, pp. 31094-31116. PMLR, 2023.
  46. Strati, F., Mcallister, S., Phanishayee, A., Tarnawski, J., and Klimovic, A. Déjàvu: Kv-cache streaming for fast, fault-tolerant generative llm serving. arXiv preprint arXiv:2403.01876, 2024. · arXiv
  47. Sukhbaatar, S., Grave, E., Bojanowski, P., and Joulin, A. Adaptive attention span in transformers. arXiv preprint arXiv:1905.07799, 2019. · arXiv
  48. Sun, F., Liu, J., Wu, J., Pei, C., Lin, X., Ou, W., and Jiang, P. Bert4rec: Sequential recommendation with bidirectional encoder representations from transformer. In Proceedings of the 28th ACM international conference on information and knowledge management, pp. 1441-1450, 2019.
  49. Team, M. N. et al. Introducing mpt-7b: A new standard for open-source, commercially usable llms, 2023. URL www. mosaicml. com/blog/mpt-7b. Accessed, pp. 05-05, 2023.
  50. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
  51. Wang, H., Zhang, Z., and Han, S. Spatten: Efficient sparse attention architecture with cascade token and head prun- ing. In 2021 IEEE International Symposium on High- Performance Computer Architecture (HPCA), pp. 97-110. IEEE, 2021.
  52. Wang, I., Nair, P. J., and Mahajan, D. FLuID: Mitigating stragglers in federated learning using invariant dropout. In Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023.
  53. Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. · arXiv
  54. Wang, S., Wei, J., Sabne, A., Davis, A., Ilbeyi, B., Hecht- man, B., Chen, D., Murthy, K. S., Maggioni, M., Zhang, Q., et al. Overlap communication with dependent compu- tation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, pp. 93-106, 2022.
  55. Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019. · arXiv
  56. Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023. · arXiv
  57. Yan, Y., Chen, J., Qi, W., Bhendawade, N., Gong, Y., Duan, N., and Zhang, R. El-attention: Memory efficient lossless attention for generation. In International Conference on Machine Learning, pp. 11648-11658. PMLR, 2021.
  58. Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Al- berti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al. Big bird: Transformers for longer se- quences. Advances in neural information processing systems, 33:17283-17297, 2020.
  59. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H _2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048, 2023. · arXiv
  60. Zhao, W. X., Zhou, K., Li, J., Tang, T., Wang, X., Hou, Y., Min, Y., Zhang, B., Zhang, J., Dong, Z., et al. A survey of large language models. arXiv preprint arXiv:2303.18223, 2023. · arXiv
  61. Zhao, Y., Wu, D., and Wang, J. Alisa: Accelerating large language model inference via sparsity-aware kv caching. arXiv preprint arXiv:2403.17312, 2024. · arXiv

Cited by

Sorted by citations

Authors
Adnan, M., Arunkumar, A., Jain, G., Nair, P. J., Soloveychik, I., & Kamath, P.
Affiliations
  • University of British Columbia (CA)
  • d-Matrix Inc. (United States) (US)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2024
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2403.09054
arXiv
2403.09054
OpenAlex
W4392867129
Cited by
9
References
1
Citations by year
2025: 8 · 2024: 1
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by)
  • arXiv (Cornell University) (cc-by)
Funders
Natural Sciences and Engineering Research Council of Canada (rgpin-2019-05059, RGPIN-2019)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
  • Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
  • Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
Keywords
large language models, generative inference, memory bandwidth, long-context generation, inference latency

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review