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

TL;DR

This paper presents SparAMX, a method leveraging AMX support and unstructured sparsity on Intel CPUs to significantly accelerate large language model token generation, reducing latency and enabling more efficient AI inference.

Contribution

It introduces novel sparse kernels and techniques that accelerate LLM inference on CPUs, combining AMX and unstructured sparsity for the first time in attention computation.

Findings

1.42x reduction in end-to-end latency with sparse linear layers.

1.14x speedup in attention computation without accuracy loss.

Open-source implementation for accelerating PyTorch models.

Abstract

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