用CPU的AMX加速压缩大模型，推理速度提升1.42倍。
SparAMX: Accelerating Compressed LLMs Token Generation on AMX-powered CPUs
- 利用AMX和稀疏化优化线性层计算
- 端到端延迟降低1.42倍，注意力计算提速1.14倍
- 开源工具包可适配任意PyTorch模型
大语言模型计算开销高、延迟大、内存需求高。尽管专用加速器如GPU和TPU常用于运行这些任务，但CPU更普及且功耗更低。在最新Intel CPU支持的高级矩阵扩展（AMX）基础上，结合非结构化稀疏性，本工作通过在线性层中应用技术，相较当前PyTorch实现，实现了1.42倍的端到端延迟降低。我们提供一组开源自定义稀疏内核，可自动将任意PyTorch模型中的所有线性层替换为我们的稀疏实现以加速计算。此外，首次在注意力计算中使用非结构化稀疏性，实现了1.14倍的速度提升，且不损失精度。代码已开源。
原文摘要 · Abstract (English)
Large language models have high compute, latency, and memory requirements. While specialized accelerators such as GPUs and TPUs typically run these workloads, CPUs are more widely available and consume less energy. Accelerating LLMs with CPUs enables broader AI access at a lower cost and power consumption. This acceleration potential for CPUs is especially relevant during the memory-bound decoding stage of LLM inference, which processes one token at a time and is becoming increasingly utilized with reasoning models. We utilize Advanced Matrix Extensions (AMX) support on the latest Intel CPUs together with unstructured sparsity to achieve a $1.42 \times$ reduction in end-to-end latency compared to the current PyTorch implementation by applying our technique in linear layers. We provide a set of open-source customized sparse kernels that can speed up any PyTorch model by automatically replacing all linear layers with our custom sparse implementation. Furthermore, we demonstrate for the first time the use of unstructured sparsity in the attention computation achieving a $1.14 \times$ speedup over the current systems without compromising accuracy. Code: https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning/tree/main/SparAMX
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。