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

2025-02-18

AI Summary

Memory-bound token-by-token decoding severely limits large language model (LLM) inference efficiency on CPU platforms. Method: This paper proposes a co-optimization framework leveraging Intel Advanced Matrix Extensions (AMX) and unstructured sparsity—introducing unstructured sparsity into attention computation for the first time, designing hardware-aware sparse GEMM and sparse attention kernels, and building an open-source sparse kernel library that enables end-to-end, automatic replacement of PyTorch linear layers. Contribution/Results: Evaluated against native PyTorch, our approach achieves 1.42× end-to-end decoding latency reduction and 1.14× acceleration in attention computation, with zero accuracy loss. The solution is fully compatible with standard PyTorch workflows and requires no model retraining. By eliminating memory bottlenecks without sacrificing precision, this work delivers a practical, system-level optimization for energy-efficient and cost-effective LLM deployment on commodity CPUs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 imes$ 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 imes$ speedup over the current systems without compromising accuracy. Code: https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning/tree/main/SparAMX

Problem

• Accelerate LLMs on CPUs
• Reduce latency and power consumption
• Leverage AMX and unstructured sparsity

Innovation

• Utilizes AMX on Intel CPUs
• Implements unstructured sparsity in layers
• Custom sparse kernels for PyTorch