XAMBA: Enabling Efficient State Space Models on Resource-Constrained Neural Processing Units

TL;DR

The paper tackles the cost and inefficiency of deploying State-Space Models (SSMs) on resource-constrained NPUs, which limits their real-world edge applications. It introduces XAMBA, a three-step framework that enables SSMs on COTS NPUs and then optimizes their execution via CumBA (CumSum as MatMul), ReduBA (ReduceSum as matrix-vector multiply), and ActiBA (PLU-based activations) to achieve significant latency reductions with minimal accuracy loss. Key contributions include transforming sequential DSP bottlenecks into parallel, memory-efficient computations and validating on Intel Core Ultra Series 2 NPUs, achieving up to 4.8x speed-ups and decoding rates surpassing practical KPI targets. The work demonstrates that SSMs can be effectively deployed on existing hardware without bespoke accelerators, enabling energy-efficient, real-time edge AI across NLP, vision, and streaming tasks.

Abstract

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.

Figures (5)

• Figure 1: Execution bottlenecks for Mamba and Mamba-2 on Intel® Core™ Ultra Series 2 NPU. Mamba is limited by sequential DSP execution of Swish (a.k.a. SiLU) and SoftPlus, while Mamba-2 faces CumSum and ReduceSum bottlenecks. Refer \ref{['subsec:mamba1_vs_2']} for models' architectural details.
• Figure 2: XAMBA: (a) NPU architecture (b) Sequential CumSum and ReduceSum computation on a DSP. (c) CumBA and ReduBA masks for optimized computations. (d) Sequential execution of activation functions (Swish/SiLU and SoftPlus) on DSP. (e) ActiBA: Efficient execution of SoftPlus and Swish activations using C-LUT in PLU.
• Figure 3: CumBA: Enhancing memory, bandwidth, and compute efficiency by exploiting CumBA mask sparsity and two-sided sparsity acceleration in the NPU datapath.
• Figure 4: Latency reduction for Mamba and Mamba-2 130M models on Intel® Core™ Ultra Series 2 lnl NPU with XAMBA optimizations (CumBA, ReduBA, and ActiBA).
• Figure 5: Mamba mamba vs. Mamba-2 mamba2 showcasing structural differences and operational trade-offs. Mamba-2 simplifies the design but suffers from decreased performance on Intel® AI PCs lnlmtl due to less hardware-efficient operations.