Programming AMD XDNA NPUs with Open-source Compiler Tools: A FlashAttention Case Study

Erwei Wang; Ephrem Wu; Victor J. B. Jung; Jiajie Li; Andre Rosti; Joseph Melber; Samuel Bayliss

cs.ARSep 18, 2026

Programming AMD XDNA NPUs with Open-source Compiler Tools: A FlashAttention Case Study

Authors: Erwei Wang, Ephrem Wu, Victor J. B. Jung, Jiajie Li, Andre Rosti, Joseph Melber, Samuel Bayliss

Organizations: Advanced Micro Devices, Inc., USA. · Integrated Systems Laboratory (IIS), ETH Zürich, Switzerland. · Cornell University, Ithaca, NY, USA.

Abstract

Spatial NPUs such as AMD XDNA place compute tiles beside small local memories and leave data movement between them to software. Mapping a multi-stage workload onto such a device is largely a question of where the intermediate tensors live. We report what we learned making those choices for FlashAttention with the open-source IRON and MLIR-AIR flows. We compare four reference designs on XDNA 1 and XDNA 2: one runs each operator separately, two stream between operators on chip, and one fuses all three attention stages into a single kernel. The fused kernel holds the $\boldsymbol{QK}^{\mathsf T}$ scores in compute-tile local memory and reduces partial results over the cascade interconnect, so the scores never return to shared MemTile memory. On XDNA 2, it reaches 3.62 TFLOP/s over complete end-to-end execution, twice the IRON design, with 5.3 to 7.2 times the energy efficiency of the integrated GPU on the same chip at 2K tokens and above. It covers twelve LLM configurations, from BERT to DeepSeek, up to 128K tokens. Roofline analysis at each memory level explains this result and shows when to stop. XDNA 1 has lower ridge points, so streaming on chip already reaches the compute-bound regime: the same fusion that doubles throughput on XDNA 2 is nearly wasted on XDNA 1. Comparing a mapping's operational intensity against each level's ridge point predicts which case applies before writing any code. Fuse until the mapping clears that ridge point, then stop. We release the reference designs as maintained open source.

Figures & tables

Feature	XDNA 1	XDNA 2
Compute tile array	4 $\times$ 4 (16 tiles)	4 $\times$ 8 (32 tiles)
Clock frequency	1.0 GHz	1.3 GHz 1 1 1 Variable. Fixed at 1.3 GHz in this paper.
MAC data type	bfloat16	bfp16 2 2 2 bfp16 means block floating point. Data I/O uses bfloat16 , and values are vector-cast before computation on XDNA 2.
MAC per clk. ( $m\times k\times n$ )	4 $\times$ 8 $\times$ 4	8 $\times$ 8 $\times$ 8
Exp. intrinsic	LUT approximation	Vector exp2
Compute tile memory	64 KB	64 KB

Table 1. XDNA 1 and XDNA 2 hardware configurations used in this paper. They reflect only the evaluated experimental settings.

Arch.	DMA count	Channel count	Bandwidth (GB/s)	Ridge (FLOP/B)
XDNA 1	$4\times$ Shim tile	4 $\times$ 2	$32.0$	$32.1$
	$4\times$ MemTile	4 $\times$ 6	$96.0$	$10.7$
	$16\times$ Compute tile	16 $\times$ 2	$128.0$	$8.0$
XDNA 2	$8\times$ Shim tile	8 $\times$ 2	$83.2$	$101.0$
	$8\times$ MemTile	8 $\times$ 6	$249.6$	$33.7$
	$32\times$ Compute tile	32 $\times$ 2	$332.8$	$25.2$

Table 2. Bandwidth ceilings at each memory hierarchy level.

	Shim tile		MemTile		Compute tile
Strategy	Matrices spilled	FLOP/B	Matrices spilled	FLOP/B	Matrices spilled	FLOP/B
Layer by layer	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$ , scale	$<\nicefrac{{d_{k}}}{{2}}$	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$ , scale	$<\nicefrac{{d_{k}}}{{2}}$	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$ , scale	$<\nicefrac{{d_{k}}}{{2}}$
DATO	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$	$\propto$ Seq. Len.	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$ , scale	$<\nicefrac{{d_{k}}}{{2}}$	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$ , scale	$<\nicefrac{{d_{k}}}{{2}}$
IRON	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$	$\propto$ Seq. Len.	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$	$\leq\nicefrac{{d_{k}}}{{2}}$	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V},\boldsymbol{QK}^{\mathsf{T}},\boldsymbol{G}$	$\leq\nicefrac{{d_{k}}}{{2}}$
MLIR-AIR	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$	$\propto$ Seq. Len.	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$	$\leq d_{k}$	$\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$	$\leq d_{k}$

Table 3. Asymptotic OI bounds and memory traffic at each NPU memory hierarchy. Tensors are in bfloat16 .

Arch.	Strategy	Shim tile			MemTile			Compute tile
Arch.	Strategy	Pred. OI (FLOP/B)	Ridge (FLOP/B)	Pred. Regime	Pred. OI (FLOP/B)	Ridge (FLOP/B)	Pred. Regime	Pred. OI (FLOP/B)	Ridge (FLOP/B)	Pred. Regime
XDNA 1	Layer by layer	$<32$	$32.1$	NR	$<32$	$10.7$	CB	$<32$	$8.0$	CB
	DATO	$\propto$ Seq. Len.	$32.1$	CB	$<32$	$10.7$	CB	$<32$	$8.0$	CB
	IRON	$\propto$ Seq. Len.	$32.1$	CB	$\leq 32$	$10.7$	CB	$\leq 32$	$8.0$	CB
	MLIR-AIR	$\propto$ Seq. Len.	$32.1$	CB	$\leq 64$	$10.7$	CB	$\leq 64$	$8.0$	CB
XDNA 2	Layer by layer	$<32$	$101.0$	MB	$<32$	$33.7$	NR	$<32$	$25.2$	NR

Table 4. Predicted execution regime at

d_{k}=64

, derived by comparing the OI bounds in Table 3 with the ridge points in Table 2 . The prediction assumes the long-sequence regime,

\mathrm{Seq.\ Len.}\gg d_{k}

. MB denotes memory-bound, CB denotes compute-bound, and NR denotes near ridge point. MB and NR are highlighted because they are the most likely performance-limiting regimes in practice.

Platform	Seq. Len.	Layer by layer		DATO		IRON		MLIR-AIR
Platform	Seq. Len.	Latency (ms)	GFLOP GEMM /s	Latency (ms)	GFLOP GEMM /s	Latency (ms)	GFLOP GEMM /s	Latency (ms)	GFLOP GEMM /s	$\Delta$ vs. IRON
XDNA 1	$512$	$22.0$	$37$	$200.2$	$4$	$9.8$	$82$	$7.0$	$114$	1.39 $\times$
	$1024$	$42.1$	$77$	$225.0$	$14$	$30.8$	$105$	$20.1$	$160$	1.53 $\times$
	$2048$	$112.8$	$114$	$499.7$	$26$	$83.4$	$154$	$67.7$	$190$	1.23 $\times$
	$4096$	$378.0$	$136$	$1220.1$	$42$	$264.6$	$195$	$250.0$	$206$	1.06 $\times$
	$8192$	$3073.8$	$67$	$4854.1$	$42$	$1054.0$	$196$	$955.2$	$216$	1.10 $\times$

Table 5. GPT-2 Small attention performance of the four mapping strategies on XDNA 1 and XDNA 2. DATO only reports XDNA 1 results.

Figure 2. Roofline analysis of the four mapping strategies at three memory hierarchy levels on XDNA 1 and 2.

Panel of Figure 3 — Figure 3. Energy efficiency of the MLIR-AIR fused kernel on XDNA 2 against CPU and iGPU.

Model	$d_{\mathrm{model}}$	$H/H_{\mathrm{kv}}$	$d_{k}$	Type	Mask	Seq. Len.	Latency ( $\mu$ s)	Throughput (GFLOP/s)
BERT-Base	768	12/12	64	MHA	Non-causal	$2048$	$5775$	2231.15
BERT-Large	1024	16/16	64	MHA	Non-causal	$2048$	$7526$	2282.74
GPT-2 Small / OPT-125M	768	12/12	64	MHA	Causal	$2048$	$5810$	2217.71
GPT-2 Medium	1024	16/16	64	MHA	Causal	$2048$	$7535$	2280.01
GPT-2 Large	1280	20/20	64	MHA	Causal	$2048$	$9196$	2335.24
Llama-2 7B	4096	32/32	128	MHA	Causal	$2048$	$34\,201$	2009.28

Table 6. Supported LLM configurations with MLIR-AIR fused kernel on XDNA 2.

Panel of Figure 4 — Figure 4. Numerical stability versus input standard deviation for GPT-2 Small ( $d_{k}=64$ , $H=12$ , $L_{Q}=L_{K}=2048$ ). Inputs satisfy $Q,K,V\sim\mathcal{N}(0,\sigma^{2})$ . Reference: FP32 scaled dot-product attention.

Roofline Modeling

See all 9 papers

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU
Jul 10, 2026Victor J. B. Jung, Gagandeep Singh, Joseph Melber +3Efficient Transformer Inference Energy-Efficient ML
Integrated Systems Laboratory (IIS), ETH Zürich, Switzerland. · AMD Research and Advanced Development (RAD). · Department of Electrical, Electronic and Information Engineering (DEI), University of Bologna, Italy.
TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs
Jun 9, 2026Wesley Pang, Gregory Hyegang Jun, Feiyang Liu +1Mixed-Precision Quantization LLM Quantization
University of Illinois Urbana-Champaign Department of Electrical and Computer Engineering USA
AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
Sep 11, 2026Varun Kumar Dasoju, Tian ZhaoEfficient Attention GPU Kernel Optimization
Department of Computer Science University of Wisconsin–Milwaukee

Jul 10, 2026J/K move · Enter open · S save

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU

Victor J. B. Jung, Gagandeep Singh, Joseph Melber, Kristof Denolf+2

Integrated Systems Laboratory (IIS), ETH Zürich, Switzerland. · AMD Research and Advanced Development (RAD). · Department of Electrical, Electronic and Information Engineering (DEI), University of Bologna, Italy.

Efficient Transformer Inference Energy-Efficient ML Efficient Attention On-Device Inference

Open paper AskPDF arXiv

The growing adoption of large language model-based agents within operating system workflows has increased the importance of energy-efficient inference on laptop-class systems-on-chip (SoCs). While cloud offloading remains common, it introduces reliability and privacy concerns that are particularly problematic for agentic workloads. Recent laptop SoCs, therefore, incorporate neural processing engines (NPUs) optimized for energy efficiency; however, effectively mapping attention mechanisms onto NPUs remains challenging due to architectural diversity and explicit data-movement programming models. In this work, we present STEEL, the first open-source implementation of FlashAttention targeting XDNA-like NPUs. STEEL introduces a dataflow formulation of prefill attention, enabling efficient exploitation of spatial parallelism and on-chip memory. Furthermore, STEEL addresses the load imbalance induced by the causal mask by leveraging a sparsity-aware pipeline placement onto the NPU array, reducing synchronization overhead and improving utilization. We evaluate STEEL on the AMD Ryzen AI 9 HX 370 SoC and compare its performance against optimized CPU and GPU implementations. Experimental results show that STEEL reduces energy consumption by an average of 9.17x and 1.75x relative to CPU and GPU baselines, respectively. On XDNA 1, STEEL achieves an average 9.6x latency reduction over the prior state of the art, and delivers a 22.8x speedup on average compared to a layer-by-layer attention implementation on XDNA 2.

Abstract

Roofline Modeling

Explore similar work

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU

TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels