SparAMX: Accelerating Compressed LLMs Token Generation on AMX-powered CPUs

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: AbouElhamayed, Ahmed F., Dotzel, Jordan, Akhauri, Yash, Chang, Chi-Chih, Gobriel, Sameh, Muñoz, J. Pablo, Chua, Vui Seng, Jain, Nilesh, Abdelfattah, Mohamed S.
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866915157536407552
author AbouElhamayed, Ahmed F.
Dotzel, Jordan
Akhauri, Yash
Chang, Chi-Chih
Gobriel, Sameh
Muñoz, J. Pablo
Chua, Vui Seng
Jain, Nilesh
Abdelfattah, Mohamed S.
author_facet AbouElhamayed, Ahmed F.
Dotzel, Jordan
Akhauri, Yash
Chang, Chi-Chih
Gobriel, Sameh
Muñoz, J. Pablo
Chua, Vui Seng
Jain, Nilesh
Abdelfattah, Mohamed S.
contents 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
format Preprint
id arxiv_https___arxiv_org_abs_2502_12444
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle SparAMX: Accelerating Compressed LLMs Token Generation on AMX-powered CPUs
AbouElhamayed, Ahmed F.
Dotzel, Jordan
Akhauri, Yash
Chang, Chi-Chih
Gobriel, Sameh
Muñoz, J. Pablo
Chua, Vui Seng
Jain, Nilesh
Abdelfattah, Mohamed S.
Machine Learning
Artificial Intelligence
Hardware Architecture
Performance
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
title SparAMX: Accelerating Compressed LLMs Token Generation on AMX-powered CPUs
topic Machine Learning
Artificial Intelligence
Hardware Architecture
Performance
url https://arxiv.org/abs/2502.12444