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/ |
| 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/ |
Operator Dashboard
Legend
- ✅ Transform file available (device target supported)
- — Not yet available
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