SparAMX：面向 AMX 加速 CPU 的压缩大语言模型加速生成技术
机器学习
2025-02-19 v1 人工智能
硬件体系结构
性能
摘要
大型语言模型具有高计算、延迟和内存需求。虽然 GPU 和 TPU 等专用加速器通常用于运行这些工作负载，但 CPU 更广泛可用且耗能更低。利用 CPU 加速大语言模型可在更低成本和能耗下实现更广泛的 AI 访问。这种 CPU 加速潜力在大语言模型推理的内存受限解码阶段尤为重要，该阶段逐 token 处理，随着推理模型的普及而日益受到重视。我们利用最新 Intel CPU 上的高级矩阵扩展（AMX）支持结合非结构化稀疏性，通过在线性层中应用技术实现端到端延迟降低 1.42 倍。我们提供了一套开源定制稀疏内核，可通过自动替换所有线性层为自定义稀疏实现来加速任意 PyTorch 模型。此外，我们首次在注意力计算中应用非结构化稀疏性，在不牺牲准确性的前提下实现 1.14 倍加速。代码：https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning/tree/main/SparAMX
引用
@article{arxiv.2502.12444,
title = {SparAMX: Accelerating Compressed LLMs Token Generation on AMX-powered CPUs},
author = {Ahmed F. AbouElhamayed and Jordan Dotzel and Yash Akhauri and Chi-Chih Chang and Sameh Gobriel and J. Pablo Muñoz and Vui Seng Chua and Nilesh Jain and Mohamed S. Abdelfattah},
journal= {arXiv preprint arXiv:2502.12444},
year = {2025}
}