ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching

📄 arXiv: 2403.17312v1 📥 PDF

作者: Youpeng Zhao, Di Wu, Jun Wang

分类: cs.AI, cs.LG, cs.PF

发布日期: 2024-03-26

备注: ISCA 2024

💡 一句话要点

提出ALISA以解决长序列推理中的KV缓存效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 KV缓存 稀疏注意力 推理优化 动态调度

📋 核心要点

  1. 现有大型语言模型在推理时面临计算和内存密集的挑战，尤其是在处理长序列时，KV缓存的内存需求显著增加。
  2. ALISA通过引入稀疏窗口注意力算法，优先考虑重要标记，从而降低KV缓存的内存占用，同时保持准确率。
  3. 在单GPU-CPU系统中，ALISA在不同工作负载下显著提升了FlexGen和vLLM的吞吐量，分别达到3倍和1.9倍的提升。

📝 摘要（中文）

Transformer架构在自然语言处理领域取得了显著进展，尤其是在大型语言模型（LLMs）如LLaMA和OPT的开发中。然而，LLMs在实际推理中面临计算和内存密集的挑战。尽管KV缓存可以通过将二次复杂度计算替换为线性复杂度内存访问来加速推理，但在处理长序列时，内存需求的增加可能导致I/O瓶颈和内存溢出。为此，本文提出了ALISA，一种算法-系统协同设计的解决方案，通过稀疏窗口注意力算法（SWA）优先考虑生成新标记时最重要的标记，从而在保持准确率的前提下显著降低KV缓存的内存占用。此外，ALISA还通过三阶段的动态调度优化缓存与重计算之间的权衡，在单GPU-CPU系统中，ALISA在不同工作负载下将FlexGen和vLLM的吞吐量分别提升了3倍和1.9倍。

🔬 方法详解

问题定义：本文旨在解决大型语言模型推理中的KV缓存效率问题，现有方法在处理长序列时面临内存需求增加和I/O瓶颈，导致吞吐量下降和内存溢出。

核心思路：ALISA的核心思路是通过稀疏窗口注意力算法（SWA）优先考虑生成新标记时最重要的标记，从而在保持准确率的前提下显著降低KV缓存的内存占用。

技术框架：ALISA的整体架构包括算法层的稀疏窗口注意力模块和系统层的三阶段动态调度模块，后者优化了缓存与重计算之间的权衡。

关键创新：ALISA的主要创新在于引入了高稀疏度的注意力机制，显著减少了内存占用，同时通过动态调度提高了系统的整体性能，这与传统的KV缓存方法形成了鲜明对比。

关键设计：在设计中，ALISA采用了动态调度策略，考虑了不同工作负载下的性能需求，并通过优化参数设置来平衡缓存和重计算的效率。具体的损失函数和网络结构细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ALISA在单GPU-CPU系统中，通过引入稀疏窗口注意力算法和动态调度策略，在不同工作负载下将FlexGen的吞吐量提升了3倍，vLLM的吞吐量提升了1.9倍，显著提高了推理效率，展示了其在资源受限环境中的优越性。

🎯 应用场景

ALISA的研究成果在自然语言处理、对话系统和文本生成等领域具有广泛的应用潜力。通过提高大型语言模型的推理效率，ALISA能够支持更复杂的应用场景，尤其是在资源受限的环境中，提升用户体验和系统响应速度。未来，ALISA的技术可以扩展到其他深度学习模型的推理优化中，推动更多AI应用的发展。

📄 摘要（原文）

The Transformer architecture has significantly advanced natural language processing (NLP) and has been foundational in developing large language models (LLMs) such as LLaMA and OPT, which have come to dominate a broad range of NLP tasks. Despite their superior accuracy, LLMs present unique challenges in practical inference, concerning the compute and memory-intensive nature. Thanks to the autoregressive characteristic of LLM inference, KV caching for the attention layers in Transformers can effectively accelerate LLM inference by substituting quadratic-complexity computation with linear-complexity memory accesses. Yet, this approach requires increasing memory as demand grows for processing longer sequences. The overhead leads to reduced throughput due to I/O bottlenecks and even out-of-memory errors, particularly on resource-constrained systems like a single commodity GPU. In this paper, we propose ALISA, a novel algorithm-system co-design solution to address the challenges imposed by KV caching. On the algorithm level, ALISA prioritizes tokens that are most important in generating a new token via a Sparse Window Attention (SWA) algorithm. SWA introduces high sparsity in attention layers and reduces the memory footprint of KV caching at negligible accuracy loss. On the system level, ALISA employs three-phase token-level dynamical scheduling and optimizes the trade-off between caching and recomputation, thus maximizing the overall performance in resource-constrained systems. In a single GPU-CPU system, we demonstrate that under varying workloads, ALISA improves the throughput of baseline systems such as FlexGen and vLLM by up to 3X and 1.9X, respectively.