On-device NPU inference engine

The bar before me

On Linux, the open-source high-water mark for this NPU was a single matmul at ~0.6% utilization, and AMD's own Linux stack offloads zero ops to it. Running real model graphs — transformers, ASR, vision, protein models — natively on the part, end to end, was simply not a thing that existed.

What I built

An inference engine on the open MLIR-AIE / IRON stack that compiles and runs 16 models across 11 architectures natively on the NPU — BERT · Whisper / Parakeet / GigaAM · ViT / DINOv2 / ResNet-18 / CLIP · ESM-2 — at numerical parity ~4e-3 against the reference implementations.

The core trick is collapsing a 12-layer transformer decode into a single NPU dispatch — 72 dispatches per token down to 1 — 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 — and 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 its O(n²) passes to linear (byte-identical output, ~58% off aiecc, ~11× less peak memory). The operator library (IRON) was missing the GEMM-fusion and batched transpose those designs need — I added them. 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, and fixed it at the backend, upstreaming a register-overflow diagnostic all the way into LLVM proper — the compiler infrastructure Clang, Rust and Swift are built on.

At equal accuracy (WER 0.117 on speech), the engine runs at −29% energy and roughly half the package power of the same workload on the CPU — the point of an NPU, finally realized on this Linux stack.