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

2025-02-10

AI Summary

Deploying state space models (SSMs) on resource-constrained AI PCs—such as those equipped with Intel Core Ultra NPUs—faces severe bottlenecks due to hardware inefficiency in executing sequence-dependent operations (e.g., cumsum, reduce_sum). Method: We propose a holistic optimization framework: (i) CumBA/ReduBA matrix-ify cumulative and reduction operations; (ii) ActiBA accelerates Swish/Softplus activations via piecewise-linear approximations; and (iii) architecture-aware customization, kernel-level NPU adaptation, and precision-performance co-optimization. Contribution/Results: This work presents the first efficient SSM inference implementation on commercial NPUs. Experiments on real-world AI PCs demonstrate a 2.6× speedup in inference latency, significantly improved memory efficiency, and substantial performance gains for long-sequence tasks—including real-time transcription and translation. Our open-source implementation is publicly available.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Learning & Optimization for NLP
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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 2.6X speed-up over the baseline. Our implementation is available at https://github.com/arghadippurdue/XAMBA.

Problem

• Optimizing State-Space Models for NPUs
• Enhancing performance on resource-constrained devices
• Reducing latency with minimal accuracy loss

Innovation

• Optimizes SSMs on COTS NPUs
• Replaces sequential operations with matrices
• Approximates activation functions for efficiency