xdna-engine

A general inference engine for the AMD XDNA2 (Strix) NPU, written in Rust with hand-written AIE kernels. It runs transformer and conv models - ASR, embeddings, small LLMs, and vision - on the NPU under Linux via the open MLIR-AIE / IRON kernel stack, with a host-CPU fallback for ops that are not yet on-device.

ASR was the first target, but the engine is not ASR-specific: the same primitives (resident dataflow, fused decode, KV cache, multi-precision GEMM/GEMV) serve every front through one Frontend / Encoder / Head pipeline.

Why the NPU

This pipeline is data-movement-bound, not compute-bound. The NPU's cores sit mostly idle; the cost is bytes streamed from LPDDR and array shape-reloads. The engine is built around that fact: keep weights and activations on-chip, fuse op sequences into few dispatches, and quantize to cut the bytes moved. The payoff is latency, energy, and freeing the CPU - see docs/data-movement-thesis.md.

What works today

• LLM decode - Qwen3-0.6B generates on the NPU: the whole MLP block and the whole QKV head each compile to one design, and a token costs 170 dispatches against 8 designs. Gemma 3-270M and Gemma 4-12B generate on the same fused-decode rail, as data (designs/decode_fused/llm_decode_spec.py) rather than a second implementation.
• ASR - GigaAM-v3 and Parakeet FastConformer encoders on the NPU; Whisper-small encoder + a full 12-layer decoder fused into a single ELF dispatch.
• Embeddings - BGE-base on the NPU, served over an OpenAI-compatible /v1/embeddings endpoint. ESM-2 (8M/35M) runs the same encoder rails.
• Super-resolution - ESPCN and EDSR on the NPU, gated against a CPU oracle, reached through the xdna-sr CLI and an ffmpeg filter rather than npu serve.
• Precision - selectable bf16 / bfp16 / int8, per-op, gated on WER/accuracy.

Weights convert and match a reference for MiniLM, E5, ModernBERT, ViT, DINOv2, ResNet-18 and opt-125m, but their forward pass still runs on the host. The rails those models need are the ones the LLM decode work is building now, so finishing them is wiring, not research.

Representative measured results (host: AMD Ryzen AI 9 465, XDNA2, Linux):

Result	Number
GigaAM encoder, NPU vs CPU	651 ms vs 890 ms
Parakeet resident engine	4.0 s -> 0.70-0.92 s / clip, WER-lossless
BGE embeddings, NPU vs host	2.5-4x
Qwen3-0.6B decode, on NPU	74.6 -> 54.8 ms/token in one day (13.4 -> 18.3 tok/s)
Qwen3-0.6B dispatch count	366 -> 170 configures/token, 14 -> 8 designs
aiecc kernel build	536 s -> ~7 min cold, < 10 s warm

Build and run

install.sh builds the Rust workspace, installs the npu binary to ~/.local/bin, and writes a systemd --user unit. It starts nothing and touches no device (the NPU is single-tenant and may be in use).

Prerequisites, all checked by the preflight before anything is built:

• Rust (cargo on PATH).
• XRT headers and libs for the amdxdna driver - xrt/xrt_bo.h under XRT_INC_DIR (default /usr/include) and libxrt_coreutil.so* under XRT_LIB_DIR (default /usr/lib).
• An onnx-asr venv with onnx_asr importable. The install both preflights it and copies libonnxruntime.so.* out of it into ~/.local/lib/xdna-engine. Point at yours with ONNX_ASR_VENV=/path/to/venv. Unset, it resolves in order: ./.venv, then ~/.local/share/xdna-engine/onnx-asr-venv. The preflight fails loud, naming the variable, if none of them import onnx_asr.
• Model artifacts under artifacts/, generated by the export scripts in scripts/. The preflight refuses rather than installing an engine whose config points at nothing.

ONNX_ASR_VENV=/path/to/venv ./install.sh    # build + install the service
npu serve                                   # start the engine
npu transcribe audio.wav                    # run ASR
npu embed "some text"                       # run embeddings
npu model ls                                # list loaded models

Weight checkpoints are baked from Hugging Face checkpoints with npu checkpoint bake (see rust/npu-weights). Model export/convert scripts live in scripts/.

Building the AIE kernels themselves is a separate path with its own toolchain: the fork instance pinned by toolchain.lock, brought up by scripts/toolchain_up.sh and gated by scripts/toolchain_smoke.sh. The engine runs against prebuilt xclbins in artifacts/ and does not need it.

Layout

• rust/ - the engine (14 product crates + npu-probes; see ARCHITECTURE.md)
• aie_kernels/ - the kernel library: 48 hand-written AIE kernels (GEMM, GEMV, cascade FFN, MHA, conv, LayerNorm, ...), each with its numpy golden where one exists (index)
• designs/ - the IRON multi-core dataflow graphs the engine actually dispatches, built from those kernels (layout + build model per design)
• experiments/ - one-off studies and A/B probes, not dispatched by the engine
• scripts/ - model export/convert, kernel builds, device probes, eval (index)
• bench/ - latency/energy benchmark harness
• docs/ - engineering deep-dives (data-movement thesis, AIE2P architecture/roofline, benchmark methodology, ...)
• mlir-aie/ - pinned submodule (the open AIE toolchain)

Hardware

AMD Ryzen AI 9 465 (Krackan, XDNA2), Linux with the amdxdna driver and /dev/accel/accel0. The open IRON / MLIR-AIE path is distro-agnostic.

License

Apache-2.0. See LICENSE.