中文

BUZZ:面向高效 LLM 推理的蜂巢结构稀疏 KV 缓存与分段重击项

计算与语言 2024-10-31 v1 人工智能

摘要

大语言模型(LLM)在自然语言处理中至关重要,但常面临推理速度与计算效率的挑战,限制了实时部署。键值(KV)缓存机制降低了 Transformer 模型中的计算开销,但在保持上下文理解方面仍存在挑战。本文提出 BUZZ,一种利用结构化上下文信息最小化缓存内存使用并提升推理速度的新型 KV 缓存算法。BUZZ 采用蜂巢结构稀疏缓存,结合滑动窗口捕获近期信息,并将历史 token 动态分段至块中,以优先处理局部邻域中的重要 token。我们在四个真实数据集上评估 BUZZ:CNN/Daily Mail、XSUM、Wikitext 和 10-QA。结果表明 BUZZ(1)在 LLM 推理中将缓存内存使用减少 2.5×\textbf{2.5}\times,同时在长文本摘要中保持超过 99% 的准确率,(2)在相同内存限制下超越最先进性能 7.69%\textbf{7.69}\%,而全缓存方法会遭遇内存溢出问题。此外,BUZZ 实现了显著的推理加速,具有 logn\log{n} 时间复杂度。代码地址:https://github.com/JunqiZhao888/buzz-llm。

关键词

引用

@article{arxiv.2410.23079,
  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},
  journal= {arXiv preprint arXiv:2410.23079},
  year   = {2024}
}