NoMAD-Attention: Efficient LLM Inference on CPUs Through Multiply-add-free Attention

📄 arXiv: 2403.01273v1 📥 PDF

作者: Tianyi Zhang, Jonah Wonkyu Yi, Bowen Yao, Zhaozhuo Xu, Anshumali Shrivastava

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

发布日期: 2024-03-02

🔗 代码/项目: GITHUB

💡 一句话要点

提出NoMAD-Attention以解决CPU上大语言模型推理效率问题

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

关键词: 大语言模型 CPU推理 注意力机制 SIMD寄存器 高效算法 量化模型 自然语言处理

📋 核心要点

  1. 现有方法在CPU上进行大语言模型推理时，由于乘加操作的高计算成本，导致效率低下。
  2. 本文提出NoMAD-Attention，通过利用SIMD寄存器进行快速查找，替代传统的乘加运算，从而提高计算效率。
  3. 实验结果显示，NoMAD-Attention在不影响模型质量的前提下，能够将4位量化的LLaMA-7B模型推理速度提升至2倍。

📝 摘要（中文）

在中央处理单元（CPU）上进行大语言模型推理面临挑战，主要是由于注意力计算中大量昂贵的乘加（MAD）矩阵运算。本文提出NoMAD-Attention，一种高效的注意力算法，利用现代CPU中的单指令多数据（SIMD）寄存器进行批量超低延迟查找，替代MAD操作。通过硬件感知的算法设计，NoMAD-Attention在不需要模型微调的情况下，成功实现了注意力分数的计算。实验证明，该方法在保持原始模型质量的同时，在16k上下文长度下将4位量化的LLaMA-7B模型速度提升了最多2倍。

🔬 方法详解

问题定义：本文旨在解决在CPU上进行大语言模型推理时，由于乘加（MAD）操作的高计算成本而导致的效率低下问题。现有方法在处理注意力计算时，往往需要大量的计算资源，限制了其在CPU上的应用。

核心思路：论文的核心思路是利用现代CPU中的SIMD寄存器进行批量超低延迟查找，替代传统的乘加运算。通过这种设计，NoMAD-Attention能够在有限的寄存器空间内实现高效的注意力分数计算。

技术框架：NoMAD-Attention的整体架构包括数据预处理、SIMD寄存器的高效利用和注意力分数的计算模块。首先对输入数据进行预处理，然后通过SIMD寄存器进行快速查找，最后计算注意力分数并输出结果。

关键创新：NoMAD-Attention的最大创新在于将传统的乘加操作替换为寄存器内查找，显著降低了计算延迟。这一方法与现有基于MAD的注意力计算方法有本质区别，能够更好地适应CPU的硬件特性。

关键设计：在设计中，NoMAD-Attention对SIMD寄存器的使用进行了优化，确保在有限的寄存器空间内实现高效的注意力计算。此外，算法不需要对预训练的注意力模型进行微调，保持了模型的原始性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，NoMAD-Attention在16k上下文长度下，能够将4位量化的LLaMA-7B模型的推理速度提升至2倍，同时保持了模型的质量。这一显著提升为在CPU上高效运行大语言模型提供了新的解决方案。

🎯 应用场景

NoMAD-Attention的研究成果具有广泛的应用潜力，特别是在资源受限的环境中，如边缘计算和移动设备。其高效的推理能力能够使得大语言模型在更多实际场景中得以应用，推动自然语言处理技术的普及与发展。

📄 摘要（原文）

Large language model inference on Central Processing Units (CPU) is challenging due to the vast quantities of expensive Multiply-Add (MAD) matrix operations in the attention computations. In this paper, we argue that there is a rare gem in modern CPUs, Single-Instruction-Multiple-Data (SIMD) registers, which allow for ultra-low-latency lookups in batch. We leverage this unique capability of CPUs to propose NoMAD-Attention, an efficient attention algorithm that replaces MAD operations with in-register lookups. Through hardware-aware algorithmic designs, NoMAD-Attention achieves the computation of attention scores using repeated fast accesses to SIMD registers despite their highly limited sizes. Moreover, NoMAD-Attention works with pre-trained attention-based LLMs without model finetuning. Empirical evaluations demonstrate that NoMAD-Attention maintains the quality of the original LLMs well, and speeds up the 4-bit quantized LLaMA-7B-based model by up to 2$\times$ at 16k context length. Our results are reproducible at https://github.com/tonyzhang617/nomad-dist.