JAX-QNN: Qualcomm QNN Backend for JAX

JAX-QNN is an open-source native backend for JAX targeting Qualcomm Snapdragon hardware accelerators (Hexagon NPU, Adreno GPU, and Oryon CPU) via Qualcomm AI Engine Direct (QNN) and the OpenXLA PJRT C API.

It enables standard JAX code to compile StableHLO intermediate representations directly into native Qualcomm execution graphs with hardware acceleration on Windows on ARM, Linux, and Android.

Performance Highlights (Snapdragon® X Elite)

Measured on a 12-Core Snapdragon® X Elite (X1E80100) with 45 TOPS Qualcomm® Hexagon™ NPU:

See docs/benchmarks.md for full benchmark methodology and memory bandwidth analysis.

Hardware Targets

Quickstart

1. Installation

# Clone the repository
git clone https://github.com/carrycooldude/JAX-QNN.git
cd JAX-QNN

# Install in editable mode
pip install -e .

2. Basic Example

import jax
import jax.numpy as jnp
import jax_qnn

# 1. Discover devices
print("Devices:", jax.devices("qnn"))

# 2. Define standard JAX model
@jax.jit
def model(x, w, b):
    return jax.nn.relu(jnp.matmul(x, w) + b)

# 3. Create inputs
x = jnp.ones((128, 512), dtype=jnp.float32)
w = jnp.ones((512, 512), dtype=jnp.float32)
b = jnp.zeros((512,), dtype=jnp.float32)

# 4. Execute on Qualcomm Hardware
output = jax.jit(model, backend="qnn")(x, w, b)
print("Output shape:", output.shape)

Available Examples

Run any example directly:

python examples/matmul_relu.py
python examples/conv2d_relu.py
python examples/transformer_attention.py
python examples/benchmark.py

Contributing Examples

We actively encourage community contributions of new JAX models and benchmarks!

Whether you want to add:

• Vision Models (ResNet, MobileNet, ViT)
• Language Models (Llama, Mistral decoder blocks, RoPE, RMSNorm)
• Audio Models (Whisper encoder, Conformer blocks)
• Quantization Examples (INT8, FP16 GEMM)

Please check our Contributing Examples Guide to get started!

Architecture Overview

                  Python / JAX Code
                         │
           jax.jit(..., backend="qnn")
                         │
                 JAX Core Tracing
                         │
                       JAXPR
                         │
                  MLIR Lowering
                         │
                     StableHLO
                         │
            PJRT_Client_Compile (C API)
                         │
             csrc/stablehlo_to_qnn.cc
                         │
       Direct Qualcomm Native QNN C API:
         - QnnBackend_create()
         - QnnContext_create()
         - QnnGraph_create()
         - QnnGraph_addNode()
         - QnnGraph_finalize()
         - QnnGraph_execute()
                         │
           Qualcomm Hexagon NPU / HTP Hardware

For comprehensive technical specifications, refer to docs/architecture.md and docs/setup.md.

License

Apache-2.0 License. See LICENSE for details.