On-device NPU inference engine

The bar before me

On Linux, a general open engine running real model graphs (transformers, ASR, vision, protein models) natively on this NPU, end to end, with kernels you can read and change, wasn't a thing that existed. The open building blocks were isolated example kernels. The turnkey runtimes ship their NPU kernels as closed binaries. So I wrote one from scratch, kernels included.

What I built

An inference engine on the open MLIR-AIE / IRON stack. One parity-gated loader compiles 16 models across 11 architectures (BERT · Whisper / Parakeet / GigaAM · ViT / DINOv2 / ResNet-18 / CLIP · ESM-2) to a bf16 arena at ~4e-3 against the reference implementations, with no Python at runtime. ASR is fully end-to-end on the NPU: GigaAM / Parakeet encoders and Whisper decode with on-device lm-head + argmax. Embeddings (BGE and other BERT encoders, 2.5–4× vs host) and vision (ViT / DINOv2 / ResNet-18) also run on the NPU; the remaining families, including ESM-2, are still in progress. It ships as a live local service a dictation app uses daily. Public, Apache-2.0: github.com/atassis/xdna-engine↗.

I mapped the pipeline by device trace and showed it is data-movement-bound: ~93% of on-device time is context-reload and DMA, not compute. The array actually computes only ~7% of the time, and when it does, it reaches ~31% of the chip's bf16 peak. So the wins came from eliminating movement, not faster arithmetic. The sharpest one collapses a 12-layer transformer decode into a single resident NPU dispatch (from ~33 host dispatches per token down to one), with weights and the KV cache resident on the device, so each token isn't paying host round-trip and re-upload costs.

When the toolchain broke

I built this as a consumer of AMD's open NPU stack. At transformer scale, it broke at every layer. I didn't file blocker bugs and wait, and I didn't downscope to route around them: I went down to the layer each defect lived on and fixed it there, then kept building.

The compiler (mlir-aie) couldn't finish compiling large designs. I cut the O(n²) passes to linear (several of them, not just aiecc): a full kernel build went from ~20 min to under a minute, with byte-identical output. The operator library (IRON) was missing the GEMM-fusion and batched transpose those designs need, so I added them. In the runtime (mlir-air), I fixed a real single-trip cascade re-entrancy bug: an abort on the second dispatch that broke any re-run of a single-trip (M=1 decode) cascade kernel. In the driver (xdna-driver), I found a race where a device sync could return before the last DMA line had landed: a stale read that matched the golden value until it didn't (13–15% of 500 runs). I proved it with a 3200-run A/B (43 failures → 0, p ≈ 1e-19), and it merged the same day. And an on-hardware test caught AMD's LLVM backend silently computing wrong answers: a load scheduled into a soft-float call's delay slots on the in-order AIE core. I root-caused it, pinned CI to unblock the project now, fixed it at the backend, and merged a register-overflow diagnostic upstream into LLVM proper itself, the compiler infrastructure Clang, Rust and Swift are built on.

At equal accuracy (WER 0.117 on speech), full-NPU decode draws measurably lower package power: roughly −29% energy versus the same workload on the CPU, at zero CPU cores busy. (Measured at the package via RAPL; the NPU-only draw is inferred from that, not isolated by an external meter.) It's an energy-and-offload win, not a latency one. That's the honest point of an NPU here.