利用矩阵扩展提升大模型推理速度，实现跨设备3.94倍加速
SMEPilot: Characterizing and Optimizing LLM Inference with Scalable Matrix Extensions
- 按算子特性动态选择CPU、SME或混合执行方式
- 在手机到服务器上实现最高3.94倍的端到端推理加速
- 适合关注CPU端大模型优化与硬件协同设计的研究者
现代CPU越来越多地集成矩阵扩展（如Arm Scalable Matrix Extension, SME），可在CPU内实现高吞吐量的矩阵运算。然而，对于大语言模型（LLM）推理而言，这些单元并非传统CPU核心的通用替代品：预填充、解码、注意力计算和键值缓存操作表现出不同的计算强度、向量行为和数据布局需求，同时SME与CPU核心仍争夺共享内存带宽。本文通过基于屋顶模型的分析，刻画SME增强型CPU的性能特征，并据此指导算子级执行策略。我们提出SMEPilot，一个能够为每个算子形状选择仅使用CPU、仅使用SME或联合使用SME+CPU的推理引擎。SMEPilot以块粒度划分矩阵任务，在注意力计算中重叠适合SME的矩阵阶段与适合CPU的向量阶段，并保持数据布局状态，避免关键路径上重复构建压缩张量表示。在手机、个人电脑和服务器平台上对Llama-3.2-3B、Qwen3-4B和Qwen3-30BA3B进行测试，SMEPilot实现了最高达3.94倍的端到端推理性能提升。
原文摘要 · Abstract (English)
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$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。