Context-Driven Performance Modeling for Causal Inference Operators on Neural Processing Units
This work addresses the challenge of deploying large language models on resource-constrained edge devices by providing insights for hardware-aware model co-design, though it is incremental as it focuses on benchmarking existing methods.
The paper analyzed the performance of causal inference operators, including quadratic and sub-quadratic attention models, on a Neural Processing Unit (NPU), finding that quadratic attention becomes memory-bound with over 95% pipeline stalls at long contexts, while sub-quadratic models can be compute-bound.
The proliferation of large language models (LLMs) has driven demand for long context inference on resource constrained edge devices. However, deploying these models on Neural Processing Units (NPUs) presents significant challenges due to the architectural mismatch: quadratic complexity of standard attention mechanisms conflicts with memory and compute patterns of edge accelerators. This paper presents a comprehensive performance analysis of various causal inference operators on a modern NPU. We benchmark standard quadratic attention against several sub-quadratic alternatives, including structured state-space and linear attention models. Our analysis reveals that while sub-quadratic methods offer superior scalability, they introduce distinct computational bottlenecks on the NPU's specialized execution units. We identify that quadratic attention becomes severely memory-bound, suffering from cache inefficiency and pipeline stalls exceeding 95% at long contexts. In contrast, sub-quadratic models can become compute-bound on programmable vector cores. These findings provide critical insights for the co-design of hardware-aware models and optimization strategies to enable on-device AI inference with long-contexts.