摘要
arXiv:2502.06924v4 宣告类型: replace-cross
摘要:状态空间模型(SSMs)已经成为了序列数据任务中变压器的一种高效替代方案,提供了随序列长度呈线性或近线性扩展的规模性,使其成为自然语言处理、视觉和边缘AI中长序列应用的理想选择,包括实时转录、翻译和上下文搜索。这些应用需要在笔记本电脑和PC等资源受限设备上部署的轻量级高性能模型。为每种新兴的神经网络设计专用加速器既昂贵又不实际;相比之下,优化现有AI PC中的NPUs(神经处理器)性能提供了一种可扩展的解决方案。为此,我们提出XAMBA,这是第一个框架,使其能够在商用现成(COTS)的最新(SOTA)NPUs上启用和优化SSMs。XAMBA遵循三步方法论:(1)在NPUs上启用SSMs,(2)优化性能以满足关键性能指标(KPI)要求,(3)通过牺牲精度以获得额外的性能提升。在NPUs上启用SSMs后,XAMBA使用CumBA和ReduBA缓解关键瓶颈,将顺序的CumSum和ReduceSum操作替换为基于矩阵的计算,显著提高了执行速度和内存效率。此外,ActiBA通过使用分段线性映射来近似昂贵的激活函数(例如,Swish、Softplus),从而减少延迟且损失最小的精度。在Intel Core Ultra Series 2 AI PC上的评估显示,XAMBA相较于基线实现了最多4.8倍的加速。我们的实现可以在https://github.com/arghadippurdue/XAMBA获取。