论键值受限生成式语言模型推理中驱逐策略的有效性
计算与语言
2024-02-20 v2 人工智能
摘要
尽管 Large Language Models (LLMs) 近期取得了成功,但由于其过高的内存和计算需求,在资源受限的环境中部署它们成本极高。除了模型参数外,键值缓存也存储在 GPU 内存中,并随批大小和序列长度线性增长。作为补救措施,近期工作提出了各种驱逐策略,以将键值缓存的开销维持在给定预算之内。本文探讨了现有驱逐策略在重要性分数计算和驱逐范围构建方面的有效性。我们指出了先前策略在这两个方面的不足,并引入了 RoCo,一种基于时间注意力分数和鲁棒性度量的鲁棒缓存遗漏策略。涵盖预填充和自回归解码阶段的大量实验验证了 RoCo 的优越性。最后,我们发布了 EasyKV,一个专用于用户友好的键值受限生成式推理的多功能软件包。代码可在 https://github.com/DRSY/EasyKV 获取。
引用
@article{arxiv.2402.06262,
title = {On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference},
author = {Siyu Ren and Kenny Q. Zhu},
journal= {arXiv preprint arXiv:2402.06262},
year = {2024}
}