XAMBA让高效模型在低资源设备上跑得更快,适合边缘AI部署。
XAMBA: Enabling Efficient State Space Models on Resource-Constrained Neural Processing Units
- 将SSM模型适配到现成NPU,用矩阵运算替代慢操作
- 实测速度比基线快4.8倍,延迟大幅降低
- 适合做实时语音、翻译等边缘计算任务的开发者
状态空间模型(SSMs)作为序列任务中Transformer的高效替代方案,具有线性或近似线性的序列长度扩展能力,适用于自然语言处理、视觉及边缘AI中的长序列应用,如实时转录、翻译和上下文搜索。这些场景需要轻量且高性能的模型部署于笔记本电脑等资源受限设备。为避免为每类新网络设计专用加速器,我们提出XAMBA,首个针对商用主流神经处理器单元(NPU)优化SSMs的框架。XAMBA采用三步法:(1) 在NPU上启用SSMs,(2) 通过CumBA与ReduBA优化性能,将原本串行的CumSum与ReduceSum替换为基于矩阵的计算,显著提升执行速度与内存效率;(3) 通过ActiBA近似复杂激活函数(如Swish、Softplus),使用分段线性映射降低延迟,仅带来轻微精度损失。在Intel Core Ultra Series 2 AI PC上的评估显示,XAMBA相较基线实现最高4.8倍加速。代码已开源:https://github.com/arghadippurdue/XAMBA。
原文摘要 · Abstract (English)
State-Space Models (SSMs) have emerged as efficient alternatives to transformers for sequential data tasks, offering linear or near-linear scalability with sequence length, making them ideal for long-sequence applications in NLP, vision, and edge AI, including real-time transcription, translation, and contextual search. These applications require lightweight, high-performance models for deployment on resource-constrained devices like laptops and PCs. Designing specialized accelerators for every emerging neural network is costly and impractical; instead, optimizing models for existing NPUs in AI PCs provides a scalable solution. To this end, we propose XAMBA, the first framework to enable and optimize SSMs on commercial off-the-shelf (COTS) state-of-the-art (SOTA) NPUs. XAMBA follows a three-step methodology: (1) enabling SSMs on NPUs, (2) optimizing performance to meet KPI requirements, and (3) trading accuracy for additional performance gains. After enabling SSMs on NPUs, XAMBA mitigates key bottlenecks using CumBA and ReduBA, replacing sequential CumSum and ReduceSum operations with matrix-based computations, significantly improving execution speed and memory efficiency. Additionally, ActiBA enhances performance by approximating expensive activation functions (e.g., Swish, Softplus) using piecewise linear mappings, reducing latency with minimal accuracy loss. Evaluations on an Intel Core Ultra Series 2 AI PC show that XAMBA achieves up to 4.8X speed-up over the baseline. Our implementation is available at https://github.com/arghadippurdue/XAMBA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。