Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ahmed F. AbouElhamayed, Jordan Dotzel, Yash Akhauri, Chi-Chih Chang, Sameh Gobriel, J. Pablo Muñoz, Vui Seng Chua, Nilesh Jain, Mohamed S. Abdelfattah

TL;DR

SparAMX targets CPU-based acceleration of LLM decoding by combining unstructured sparsity with Intel AMX on Sapphire Rapids. The authors implement dense and sparse GEMM kernels, plus INT8 variants, to reduce memory transfers during memory-bound decode while maintaining accuracy, and extend sparsity to the KV cache within attention. Key results include up to $1.42\times$ end-to-end speedup over stock PyTorch, up to $1.46\times$ gains for INT8 versus DeepSparse, and $1.14\times$ speedups for KV-cache sparsity with minimal accuracy loss at long contexts. The approach is general-purpose, open-source, and demonstrates practical CPU-based deployment potential for LLMs, especially in memory-bound decoding scenarios where GPUs are not available or cost/power constraints are tight.

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

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

TL;DR

SparAMX targets CPU-based acceleration of LLM decoding by combining unstructured sparsity with Intel AMX on Sapphire Rapids. The authors implement dense and sparse GEMM kernels, plus INT8 variants, to reduce memory transfers during memory-bound decode while maintaining accuracy, and extend sparsity to the KV cache within attention. Key results include up to end-to-end speedup over stock PyTorch, up to gains for INT8 versus DeepSparse, and speedups for KV-cache sparsity with minimal accuracy loss at long contexts. The approach is general-purpose, open-source, and demonstrates practical CPU-based deployment potential for LLMs, especially in memory-bound decoding scenarios where GPUs are not available or cost/power constraints are tight.

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

Paper Structure

This paper contains 23 sections, 18 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. LLM Inference
  4. DNN Sparsity
  5. GEMM Mapping
  6. Advanced ISA Extensions
  7. Related Work
  8. Kernel Design
  9. Dense Kernel
  10. Sparse Format
  11. AMX Sparse Kernel
  12. AVX Sparse Kernel
  13. INT8 Kernels
  14. Results
  15. Attention Kernel
  16. ...and 8 more sections

Figures (18)

  • Figure 1: SparAMX Performance: SparAMX consistently leads to a speedup in decode latency over PyTorch. The improvement tends to be greater as the model size increases. Context length is set to 512.
  • Figure 2: Linear GEMM Mapping: A linear layer can be computed with matrix multiplication. In this example hidden_dim = 4, num_neurons = 3, and batch_size = 2.
  • Figure 3: Inference Breakdown: Linear layers dominate the latency during LLM inference. Results profiled from Llama 3 8B running on Intel(R) Xeon(R) Gold 6430L CPU with 512 context length.
  • Figure 4: AMX Matrix Multiplication: AMX tiles support BF16 (blue) and INT8 (red) values. First two tiles contain the input matrix tiles, which are then multiplied and accumulated to the third tile.
  • Figure 5: AMX Dense Kernel: (1) Tiles 0, 1, 2, and 3 act as accumulators for the results computed by multiplying (4x6), (4x7), (5x6), and (5x7) respectively. (2) Tiles 4 and 5 are utilized to load all columns of the input rows (denoted as out_rows) while tiles 6 and 7 are utilized to load all rows of the weight columns. (3) After the loop over the inner dimension ends, results are stored in memory and a new set of result tiles is initialized and computed in the same way. (4) Some boundary conditions might occur near the end of the matrix.
  • ...and 13 more figures