NEXUS: Speeding Up the Next Wave of AI Workloads, GNN and SSM Optimization on NPUs

Abstract

The rise of on-device AI has transformed modern client PCs into powerful platforms for executing diverse AI workloads locally, enabling low-latency, energy-efficient, and privacy-preserving applications. Traditional on-device AI applications have predominantly focused on computer vision tasks such as image classification, segmentation, and object detection, enabled by Convolutional Neural Networks (CNNs) and widely explored on edge devices. However, the recent surge in Generative AI and Large Language Models (LLMs) has expanded the application landscape to include personal assistants, contextual search, knowledge graph reasoning, and real-time transcription. To support these emerging workloads, the community has largely adopted transformer-based architectures due to their versatility and strong modeling capabilities across modalities. Despite their success, transformers exhibit quadratic complexity in sequence length, making them inefficient for long-context tasks. As a result, State Space Models (SSMs) have gained traction as efficient alternatives, offering linear or near-linear scaling with competitive accuracy for applications such as transcription, translation, and long-sequence language modeling. Meanwhile, Graph Neural Networks (GNNs) have emerged as state-of-the-art models for tasks involving structured and relational data, including knowledge graph reasoning, personalized recommendations, and event-driven visual analytics, achieving superior performance compared to transformers in certain domains. While CNN and transformer workloads are already well-supported on edge platforms, the enablement and optimization of newer model classes remain underexplored. To address this, we present NEXUS, a unified framework that establishes a systematic methodology for accelerating emerging AI workloads on commercial-off-the-shelf (COTS) Neural Processing Units (NPUs). NEXUS introduces a general systematic three-step approach: (1) Enablement, which adapts model execution to NPUs’ static input shapes, partitioned compute, and memory constraints; (2) Optimization, which remaps control-heavy and sequential operations to hardware-efficient regular data-parallel kernels such as matrix multiplications; and (3) Quality-Efficiency Tradeoffs, which apply quantization and approximations to improve performance and energy efficiency with minimal accuracy loss. While this work focuses on Graph Neural Networks (GNNs) and State Space Models (SSMs) as representative case studies, the NEXUS methodology is model-agnostic and can be readily extended to future emerging workloads and NPU platforms. Experimental evaluations on Intel® Core™ Ultra Series 1 and 2 AI PCs demonstrate that NEXUS achieves up to 10.8 × speedup and 8.6 × energy efficiency improvement for GNNs over CPU baselines, and up to 3.8 × latency reduction for SSMs. Individual techniques like CumBA, ReduBA, ActiBA, and QuBA contribute 2.7 ×, 1.2 ×, 3.7 ×, and 1.7 × reductions, respectively, all while maintaining model accuracy within 1.5% of full-precision baselines. Our code implementation is available at https://github.com/arghadippurdue/NEXUS.