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

arXiv: https://arxiv.org/abs/2502.06924 · HTML: https://mlsanity.com/paper/2502.06924/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 17 references that are in the map.
- First posted on arXiv 10 Feb 2025.

## Builds on (most cited first, 12 of 12)

- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [Mamba-360: Survey of State Space Models as Transformer Alternative for Long Sequence Modelling: Methods, Applications, and Challenges](https://mlsanity.com/paper/2404.16112/) (2024)
- [Compressing DMA Engine: Leveraging Activation Sparsity for Training Deep Neural Networks](https://mlsanity.com/paper/1705.01626/) (2017)
- [MARCA: Mamba Accelerator with ReConfigurable Architecture](https://mlsanity.com/paper/2409.11440/) (2024)
- [FlexNN: A Dataflow-aware Flexible Deep Learning Accelerator for Energy-Efficient Edge Devices](https://mlsanity.com/paper/2403.09026/) (2024)
- [Flex-SFU: Accelerating DNN Activation Functions by Non-Uniform Piecewise Approximation](https://mlsanity.com/paper/2305.04546/) (2023)
- [Design Space Exploration of Neural Network Activation Function Circuits](https://mlsanity.com/paper/1810.08650/) (2018)
- [NITRO: LLM Inference on Intel Laptop NPUs](https://mlsanity.com/paper/2412.11053/) (2024)

## Built on by (most cited first, 3 of 3)

- [A KL Lens on Quantization: Fast, Forward-Only Sensitivity for Mixed-Precision SSM-Transformer Models](https://mlsanity.com/paper/2604.13440/) (2026)
- [Context-Driven Performance Modeling for Causal Inference Operators on Neural Processing Units](https://mlsanity.com/paper/2509.25155/) (2025)
- [KATANA: A Fast, Low-Power Mapping of Kalman Filters onto Edge NPUs for Real-Time Tracking](https://mlsanity.com/paper/2606.14992/) (2026)

## Abstract (arXiv)

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.