NoMAD-Attention：通过无乘加注意力实现 CPU 上的高效 LLM 推理
机器学习
2024-03-05 v1 人工智能
计算与语言
摘要
大型语言模型在中央处理器（CPU）上的推理面临挑战，主要是注意力计算中大量昂贵的乘加（Multiply-Add, MAD）矩阵运算。在本文中，我们认为现代 CPU 中稀有的宝石——单指令多数据（SIMD）寄存器，允许对批量执行进行超低延迟查找。我们利用 CPU 的这一独特能力，提出了 NoMAD-Attention，这是一种高效注意力算法，用寄存器查找代替 MAD 操作。通过硬件感知的算法设计，NoMAD-Attention 能够在其高度受限大小下，通过重复快速访问 SIMD 寄存器来计算注意力得分。此外，NoMAD-Attention 可与无需模型微调的预训练注意力-based LLM 集成。实证评估表明，NoMAD-Attention 能很好地保持原始 LLM 的质量，将 4 位量化的 LLaMA-7B 模型在 16k 上下文长度下加速最高可达 2 倍。我们的结果可在 https://github.com/tonyzhang617/nomad-dist 复现。
引用
@article{arxiv.2403.01273,
title = {NoMAD-Attention: Efficient LLM Inference on CPUs Through Multiply-add-free Attention},
author = {Tianyi Zhang and Jonah Wonkyu Yi and Bowen Yao and Zhaozhuo Xu and Anshumali Shrivastava},
journal= {arXiv preprint arXiv:2403.01273},
year = {2024}
}