SME (Scalable Matrix Extension)

Computation Graphs & ML Infrastructure advanced How AI Works 11 min read

Overview

ARM’s Scalable Matrix Extension (SME) is a CPU architecture extension that adds native hardware support for matrix operations — specifically outer products — directly into the processor’s instruction set. Introduced in the Armv9-A architecture, SME represents a fundamental shift in how CPUs handle AI workloads: instead of decomposing matrix operations into sequences of vector instructions (as SIMD does), SME provides dedicated 2D tile storage and single instructions that accumulate entire outer products into these tiles.

The significance of SME for AI engineering is that it enables efficient on-device inference without dedicated GPU or NPU hardware. Where traditional SIMD (NEON, SVE) processes data as 1D vectors, SME adds a 2D “ZA array” — a square matrix register — and instructions that stream vector inputs from two sides and accumulate the results directly into the matrix. This maps perfectly to the multiply-accumulate loops at the heart of every neural network layer.

SME2, the second generation introduced in Armv9.3-A, extends this further with multi-vector instructions, compressed neural network support, and higher-throughput operations. In real-world benchmarks, SME2 has achieved a 6x speedup in chatbot response times when running the Gemma 3 LLM with int4 quantization — bringing local LLM inference on mobile devices from impractical to viable.

How It Works

The ZA Array: 2D Tile Storage

At the core of SME is the ZA (Z Array) storage — a 2D square register file that holds matrix tiles. Unlike traditional SIMD registers which are 1D vectors, ZA is a square byte array with dimensions of SVL × SVL (where SVL is the Streaming Vector Length in bytes). For a processor with SVL = 256 bits (32 bytes), the ZA storage is 32×32 = 1,024 bytes.

ZA Tile Organization

The ZA array can be accessed as differently-typed tiles. With 8-bit elements, ZA is one tile (ZA0.B). With 16-bit elements, ZA splits into 2 tiles (ZA0.H, ZA1.H). With 32-bit: 4 tiles (ZA0.S-ZA3.S). With 64-bit: 8 tiles (ZA0.D-ZA7.D). This allows the same physical storage to be partitioned based on the data type precision required.

Outer Product Accumulation

The fundamental operation of SME is the outer product. Given two vectors — one loaded as a “column” and one as a “row” — the outer product produces a matrix where each element [i][j] is column[i] × row[j]. SME computes this entire matrix in a single instruction and accumulates the result into a ZA tile.

// Conceptual view of SME outer product accumulation
// Computing C += A * B^T for matrix multiplication

// Traditional approach (SIMD): decompose into dot products
for each row i of A:
    for each column j of B:
        C[i][j] += dot_product(A[i,:], B[:,j])  // Many vector instructions

// SME approach: decompose into outer products
for each k:
    ZA_tile += outer_product(A[:,k], B[:,k])  // One instruction per k
    // Computes: ZA[i][j] += A[i][k] * B[j][k] for ALL i,j simultaneously

This is a fundamental algorithmic reframing. Traditional SIMD computes matrix multiplication as a series of dot products (reducing two vectors to a scalar). SME computes it as a series of outer products (expanding two vectors into a matrix). The outer product approach maps naturally to the 2D tile architecture and allows the entire accumulation to happen within the ZA register without writing partial results back to memory.

Streaming SVE Mode

SME introduces a “Streaming SVE mode” that reuses the existing SVE vector registers (Z0-Z31) as inputs to outer product instructions. When in streaming mode, the processor feeds SVE vectors as the column and row inputs to the outer product unit, which accumulates results into the ZA tiles. This means SME builds on the existing SVE ecosystem rather than introducing an entirely new register file for input data.

Vector-Length Agnostic Design

Like SVE, SME is vector-length agnostic. Code written for SME runs correctly on any implementation, regardless of the actual SVL width. A processor with 128-bit SVL has a 16×16-byte ZA array; one with 512-bit SVL has a 64×64-byte ZA array. The same code runs optimally on both — the hardware automatically scales throughput with the available tile size.

SME2 Enhancements

SME2 (Armv9.3-A) adds several capabilities that specifically target neural network inference:

Architectural Comparison: SME vs GPU Tensor Cores

SME serves a similar role to NVIDIA’s Tensor Cores but in the CPU domain. Tensor Cores operate on 4×4 or 8×8 matrix tiles within a GPU’s streaming multiprocessor. SME’s ZA tiles serve the same purpose within the CPU. The key difference is power efficiency and accessibility: SME operates within the CPU’s power envelope (typically 5-15W on mobile), while Tensor Cores require a discrete GPU (75-400W). For edge AI, embedded systems, and mobile devices, SME provides matrix acceleration where a GPU is unavailable or impractical.

Practical Applications

On-Device LLM Inference

SME2’s primary impact is making local LLM inference practical on ARM-based devices. Arm and Google collaborated to optimize Int4 matrix multiplications for the Gemma 3 model using Int8 outer product instructions in SME2, achieving a measured 6x speedup in chatbot response times. This brings inference from “technically possible but too slow” to “responsive conversational experience” on mobile hardware.

The key enabler is that SME2’s outer product instructions handle the quantized multiply-accumulate operations that dominate LLM inference. A single outer product instruction on 8-bit inputs generates 32-bit accumulated outputs in the ZA tile, maintaining numerical precision through the accumulation chain while operating on compact quantized weights.

Framework Integration via KleidiAI

Arm developed KleidiAI, a library of optimized compute kernels that integrates SME2 into major AI frameworks. Through KleidiAI, developers get SME2 acceleration in PyTorch, TensorFlow Lite, MediaPipe, XNNPACK, and llama.cpp without changing their application code.

# No code changes needed for SME2 acceleration
# KleidiAI provides optimized kernels to existing frameworks

# Example: Running Gemma 3 with llama.cpp on ARM with SME2
# The framework auto-detects SME2 and selects optimized kernels
./llama-cli -m gemma-3-4b-it-Q4_K_M.gguf \
    -p "Explain the attention mechanism" \
    --threads 4

# Under the hood, KleidiAI routes matrix operations through:
# 1. SME2 outer product instructions for weight × activation
# 2. Int8 outer products with Int4 weight unpacking
# 3. ZA tile accumulation with 32-bit precision
Zero Application Changes

Developers do not need to modify their model code or applications to benefit from SME2. KleidiAI is integrated at the framework runtime level, so existing models and applications automatically use SME2-optimized kernels when running on supported hardware. The acceleration is transparent — detected at runtime via CPU feature flags.

Computer Vision and Image Processing

Beyond LLMs, SME accelerates convolutional neural networks (CNNs) for computer vision. Convolution operations can be reformulated as matrix multiplications (im2col + GEMM), which map directly to SME’s outer product accumulation. ARM has demonstrated neural camera denoising with SME2, processing image denoising models in real-time on mobile processors — a use case that previously required a dedicated NPU or cloud offloading.

Hardware Availability

SME is available in Armv9-A processors, with SME2 in Armv9.3-A and later. Key availability milestones:

Key Takeaways

Further Reading