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 64-bit bfp16 block-exponent register declared 576 bits wide, so once register pressure forced a spill, the reload came back corrupted (no crash, no verifier complaint, just wrong numbers on one element in eight). I root-caused it from the reporter's own harness — 682 wrong elements out of 3968, down to zero — fixed it at the backend, and upstreamed the TableGen check that would have caught it at build time 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.