ALISA：通过稀疏感知 KV 缓存加速大型语言模型推理
人工智能
2024-03-27 v1 机器学习
性能
摘要
Transformer 架构显著推动了自然语言处理（NLP）的发展，并已成为开发大型语言模型（LLMs）的基础，例如 LLaMA 和 OPT，这些模型已在广泛的 NLP 任务中占据主导地位。尽管 LLMs 具有卓越的准确性，但其计算和内存密集型特性给实际推理带来了独特挑战。得益于 LLM 推理的自回归特性，Transformer 中注意力层的 KV 缓存能够通过以线性复杂度的内存访问替代二次复杂度的计算，有效加速 LLM 推理。然而，随着处理更长序列的需求增长，这种方法需要增加内存。由此产生的开销会导致因 I/O 瓶颈而降低吞吐量，甚至引发内存不足错误，尤其是在资源受限的系统（如单个商用 GPU）上。在本文中，我们提出了 ALISA，一种新颖的算法-系统协同设计方案，以应对 KV 缓存带来的挑战。在算法层面，ALISA 通过稀疏窗口注意力（SWA）算法，优先处理对生成新 token 最重要的 token。SWA 在注意力层中引入了高稀疏性，并在可忽略不计的精度损失下减少了 KV 缓存的内存占用。在系统层面，ALISA 采用三阶段 token 级动态调度，并优化了缓存与重计算之间的权衡，从而在资源受限的系统中最大化整体性能。在单个 GPU-CPU 系统中，我们证明，在不同的工作负载下，ALISA 将基线系统（如 FlexGen 和 vLLM）的吞吐量分别提高了最多 3 倍和 1.9 倍。
引用
@article{arxiv.2403.17312,
title = {ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching},
author = {Youpeng Zhao and Di Wu and Jun Wang},
journal= {arXiv preprint arXiv:2403.17312},
year = {2024}
}
备注
ISCA 2024