BUZZ: Beehive-structured Sparse KV Cache with Segmented Heavy Hitters for Efficient LLM Inference

Junqi Zhao, Zhijin Fang, Shu Li, Shaohui Yang, Shichao He
10/30/2024
cs.CLcs.AI

Official code: JunqiZhao888/buzz-llm (5 stars, last commit 2024-10-31).

Abstract

Large language models (LLMs) are essential in natural language processing but often struggle with inference speed and computational efficiency, limiting real-time deployment. The key-value (KV) cache mechanism reduces computational overhead in transformer models, but challenges in maintaining contextual understanding remain. In this paper, we propose BUZZ, a novel KV caching algorithm that leverages structured contextual information to minimize cache memory usage while enhancing inference speed. BUZZ employs a beehive-structured sparse cache, incorporating a sliding window to capture recent information and dynamically segmenting historical tokens into chunks to prioritize important tokens in local neighborhoods. We evaluate BUZZ on four real-world datasets: CNN/Daily Mail, XSUM, Wikitext, and 10-QA. Our results demonstrate that BUZZ (1) reduces cache memory usage by $\textbf{2.5}\times$ in LLM inference while maintaining over 99% accuracy in long-text summarization, and (2) surpasses state-of-the-art performance in multi-document question answering by $\textbf{7.69%}$ under the same memory limit, where full cache methods encounter out-of-memory issues. Additionally, BUZZ achieves significant inference speedup with a $\log{n}$ time complexity. The code is available at https://github.com/JunqiZhao888/buzz-llm.

View on arXivView PDF

Code Implementations(1)

51PythonSep 9, 20241 years ago

Similar Papers

Cite this paper

@article{zhao2024buzz,
  title  = {BUZZ: Beehive-structured Sparse KV Cache with Segmented Heavy Hitters for Efficient LLM Inference},
  author = {Junqi Zhao and Zhijin Fang and Shu Li and Shaohui Yang and Shichao He},
  year   = {2024},
  eprint = {2410.23079},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2410.23079v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2410.23079.svg)](https://codeofpaper.com/paper/2410.23079)

Discussion