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.
翻译:在中央处理器(CPU)上进行大语言模型推理极具挑战性,因为注意力计算涉及大量昂贵的乘加(MAD)矩阵运算。本文指出现代CPU中蕴含着一项鲜为人知的技术——单指令多数据(SIMD)寄存器,该寄存器可实现批量超低延迟查找。我们利用CPU的独特能力提出NoMAD-Attention,这是一种用寄存器内查找替代MAD操作的高效注意力算法。通过硬件感知的算法设计,NoMAD-Attention即使在高限制的SIMD寄存器容量下,仍能通过重复快速访问实现注意力分数计算。此外,该算法无需微调即可直接应用于基于预训练注意力机制的大语言模型。实验结果表明,NoMAD-Attention能良好保持原始大语言模型的生成质量,并在16k上下文长度下,将基于4位量化LLaMA-7B的模型推理速度提升高达2倍。我们的结果可在https://github.com/tonyzhang617/nomad-dist复现。