SMEPilot: Characterizing and Optimizing LLM Inference with Scalable Matrix Extensions

2026-06-15

AI Summary

This work addresses the challenge of efficiently leveraging matrix extension units in modern CPUs—such as Arm Scalable Matrix Extension (SME)—for large language model (LLM) inference, which is hindered by mismatches in arithmetic intensity, vector behavior, memory layout requirements, and bandwidth contention with conventional cores. The paper presents the first LLM inference engine tailored for SME architectures, employing roofline-model-guided operator-level performance analysis to dynamically select among CPU-only, SME-only, or cooperative execution strategies. Key innovations include operator-level heterogeneous scheduling, tile-granularity workload partitioning, overlapping matrix and vector phases in attention computation, and tensor layout reuse along critical paths. Evaluated on Llama-3.2-3B, Qwen3-4B, and Qwen3-30B across mobile, desktop, and server platforms, the approach achieves up to 3.94× end-to-end speedup.

Abstract

Modern CPUs increasingly integrate matrix extensions, such as Arm Scalable Matrix Extension (SME), that provide high-throughput matrix execution within the CPU. For LLM inference, however, these units are not a universal replacement for conventional CPU cores: prefill, decode, attention, and KV-cache operations expose different arithmetic intensities, vector behavior, and layout requirements, while SME units and CPU cores still compete for shared memory bandwidth. This paper studies this mismatch through a roofline-based characterization of SME-enabled CPUs and uses the resulting model to guide operator-level execution choices. We present SMEPilot, an LLM inference engine that selects CPU-only, SME-only, or cooperative SME+CPU execution for each operator shape. SMEPilot partitions matrix work across SME and CPU cores at tile granularity, overlaps SME-suitable matrix stages with CPU-suitable vector stages in attention, and maintains layout state so packed tensor representations are reused rather than repeatedly rebuilt on critical paths. Across Llama-3.2-3B, Qwen3-4B, and Qwen3-30BA3B on phone, PC, and server platforms, SMEPilot improves end-to-end inference performance by up to 3.94$\times$.

Problem

• LLM inference
• matrix extensions
• CPU cores
• memory bandwidth
• arithmetic intensity

Innovation

• Scalable Matrix Extension
• LLM inference optimization
• heterogeneous execution
• roofline modeling
• tensor layout management