Hexagon-MLIR: An AI Compilation Stack For Qualcomm's Neural Processing Units (NPUs)

2026-02-23

AI Summary

This work addresses the memory bandwidth bottleneck and limited flexibility of conventional library-based approaches in deploying AI models on Qualcomm’s Hexagon NPU by introducing the first open-source MLIR compilation stack tailored for this architecture. The proposed stack enables automatic compilation of optimized NPU binaries from both Triton kernels and PyTorch 2.0 subgraphs. Through a structured compilation pipeline, mega-kernel fusion, and explicit optimization of data locality in tightly coupled memory (TCM), it significantly enhances data reuse and alleviates off-chip bandwidth pressure. Notably, this study presents the first end-to-end compilation framework from Triton to Hexagon NPU, establishing a high-performance, scalable, and open-source deployment pathway for edge AI applications.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

In this paper, we present Hexagon-MLIR,an open-source compilation stack that targets Qualcomm Hexagon Neural Processing Unit (NPU) and provides unified support for lowering Triton kernels and PyTorch models . Built using the MLIR framework, our compiler applies a structured sequence of passes to exploit NPU architectural features to accelerate AI workloads. It enables faster deployment of new Triton kernels (hand-written or subgraphs from PyTorch 2.0), for our target by providing automated compilation from kernel to binary. By ingesting Triton kernels, we generate mega-kernels that maximize data locality in the NPU's Tightly Coupled Memory (TCM), reducing the bandwidth bottlenecks inherent in library-based approaches. This initiative complements our commercial toolchains by providing developers with an open-source MLIR-based compilation stack that gives them a path to advance AI compilation capabilities through a more flexible approach. Hexagon-MLIR is a work-in-progress, and we are continuing to add many more optimizations and capabilities in this effort.

Problem

• AI compilation
• Neural Processing Unit
• Triton kernels
• PyTorch models
• bandwidth bottleneck

Innovation

• Hexagon-MLIR
• Neural Processing Unit (NPU)
• Triton kernels
• Tightly Coupled Memory (TCM)
• MLIR