Triton-XDNA Examples

These examples demonstrate how to write Triton kernels that compile and run on AMD XDNA™ NPUs via the MLIR-AIR compilation flow.

End-to-End Models

Complete models—not just single operators—built by composing Triton kernels and running across the iGPU and NPU.

Model Description Datatype(s) Example
GPT-2 End-to-end GPT-2 inference (all four sizes: small/medium/large/xl) composed from Triton kernels, running across iGPU and NPU. bf16 gpt2/
Qwen2.5 End-to-end Qwen2.5-Instruct inference (0.5B/1.5B) composed from Triton kernels, with KV-cached autoregressive generation running across iGPU and NPU. bf16 qwen2_5/
Llama-3.2-1B (Q4NX) End-to-end Llama-3.2-1B on the NPU: a Triton prefill built from this repo’s kernels, feeding mlir-air’s fused Q4NX decode. bf16 prefill, q4nx decode llama32_1b_q4nx/
Llama-3.2-3B (Q4NX) The Llama-3.2-3B configuration of the same Triton-prefill / AIR-decode split, sharing the llm_q4nx harness. bf16 prefill, q4nx decode llama32_3b_q4nx/
Llama-3.1-8B (Q4NX) The Llama-3.1-8B configuration of the same split. The same block as the 1B – SwiGLU, one norm pair, no qk-norm – but 32 layers of 4096 against the 1B’s 16 of 2048, and an untied LM head. bf16 prefill, q4nx decode llama31_8b_q4nx/
Qwen3-4B (Q4NX) Qwen3-4B on the same Triton-prefill / AIR-decode split, and the first Q4NX model here that is not Llama-shaped: each head is RMS-normalized between the QKV projection and RoPE. bf16 prefill, q4nx decode qwen3_4b_q4nx/
Qwen3-8B (Q4NX) Qwen3-4B’s block at 4096, reusing its forward unchanged. What it adds is size: an untied LM head, and a decode whose weights no longer fit one buffer. bf16 prefill, q4nx decode qwen3_8b_q4nx/
Qwen2.5-7B (Q4NX) Llama-shaped but for a bias on the q, k and v projections, which nothing else here has. Its weights are not a bundle either: an upstream checkpoint is quantized on load. bf16 prefill, q4nx decode qwen25_7b_q4nx/
Qwen2.5-3B (Q4NX) Qwen2.5-7B’s block at 2048, on the same biased-QKV forward. The only model here whose decode runs with the dual-channel weight feed off. bf16 prefill, q4nx decode qwen25_3b_q4nx/
Phi-4-mini (Q4NX) Llama-shaped but for the rotation: RoPE covers 96 of each head’s 128 lanes, on frequencies read from a LongRoPE table in the bundle. The first model here whose delta is in the operator. bf16 prefill, q4nx decode phi4_mini_q4nx/
Gemma3-4B (Q4NX) Gemma3-4B on the same split, and the furthest from the Llama block: a four-norm sandwich, two RoPE thetas chosen per layer, a 1024-token sliding window, and a GELU-tanh GLU. bf16 prefill, q4nx decode gemma3_4b_q4nx/
Gemma4-E2B (Q4NX) Gemma3’s block plus per-layer embeddings, layers that differ in shape from each other, and twenty that carry no KV cache of their own. Its decode runs on mlir-air’s per-layer-embedding engine rather than the shared one. bf16 prefill, q4nx decode gemma4_e2b_q4nx/
Gemma4-E4B (Q4NX) Gemma4-E2B’s architecture at D 2560 and 42 layers, with two KV heads. mlir-air’s fused NPU prefill and per-layer-embedding decode. bf16 prefill, q4nx decode gemma4_e4b_q4nx/
Gemma4-E2B flash (image) An image turn on Gemma4-E2B: the vision tower encodes the image’s patches into soft tokens that take the place of the image placeholders in the prompt. The tower, the prefill and the decode each run on the NPU or the iGPU. bf16 vision + prefill, q4nx decode gemma4_e2b_flash/
Gemma4-E4B flash (image) Gemma4-E2B flash’s image turn on the Gemma4-E4B checkpoint, whose soft tokens are projected to D 2560. bf16 vision + prefill, q4nx decode gemma4_e4b_flash/

Operator Dashboard

Category Operation Datatype(s) AIE2 AIE2P Example
Matrix Matrix Multiplication (BF16) bf16 ✅ ✅ matmul_bf16_m64_n64_k64/
Matrix Padded Matrix Multiplication (F32, A Transposed) f32 (bf16 emulation) — ✅ matmul_f32_m64_n32_k16_padded_atransposed/
Matrix Matrix Multiplication (INT8) i8 — ✅ matmul_i8_m64_n64_k64/
Matrix Matrix Multiplication (INT8, Large Tile) i8 — ✅ matmul_i8_m128_n64_k64/
Matrix Matrix Multiplication (BF16, K Loop) bf16 — ✅ matmul_bf16_k_loop/
Matrix Matrix Multiplication (BF16, Elementwise Epilogue) bf16 — ✅ matmul_bf16_epilogue/
Matrix Matrix Multiplication (Autotune) bf16 ✅ — autotune-matmul/
Element-wise ReLU bf16 ✅ ✅ relu/
Element-wise Sigmoid bf16 ✅ ✅ sigmoid/
Element-wise SiLU bf16 ✅ ✅ silu/
Element-wise GELU bf16 — ✅ gelu/
Element-wise Erf / exact GELU bf16 — ✅ erf/
Element-wise Leaky ReLU bf16 ✅ ✅ leaky_relu/
Element-wise SwiGLU bf16 ✅ ✅ swiglu/
Element-wise AXPY bf16 ✅ ✅ axpy/
Element-wise Vector Add bf16 ✅ ✅ vec-add/
Normalization RMS Normalization bf16 — ✅ rms_norm/
Normalization Weighted RMS Normalization bf16 — ✅ weighted_rms_norm/
Normalization Softmax bf16 ✅ ✅ test_softmax/
Normalization Layer Normalization f32 ✅ ✅ test_layernorm/
Pooling Average Pool bf16 ✅ ✅ average_pool/
Special 2D Block Load f32 — — load_2d_block/
Special Multi-Driver bf16 ✅ ✅ multi_drivers/

Legend

AIE2 = AMD Ryzen™ AI (Phoenix, NPU1)    AIE2P = AMD Ryzen™ AI (Strix, NPU2)

Runtime Features

Backend capabilities rather than single operators, demonstrated end to end.

Feature Description Datatype(s) Example
Zero-Copy iGPU/NPU Buffers Buffers the iGPU and the NPU can both address, so a tensor handed from one to the other never moves – with a benchmark of the hand-off with and without copies, under both NPU runtimes. f32 zero_copy/

Running Examples

Make sure XRT is sourced and a virtual environment with triton-xdna is active (see top-level README):

source /opt/xilinx/xrt/setup.sh

# Run an example on AIE2 (NPU1):
cd matmul_bf16_m64_n64_k64
AIR_TRANSFORM_TILING_SCRIPT=transform_aie2.mlir python matmul_bf16_m64_n64_k64.py

# Run on AIE2P (NPU2):
AIR_TRANSFORM_TILING_SCRIPT=transform_aie2p.mlir python matmul_bf16_m64_n64_k64.py

Running All Tests

python scripts/run_tests.py --device aie2 --verbose
python scripts/run_tests.py --device aie2p --verbose