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.AR · cs.LG
Quick verdict
Reports four FlashAttention mappings on AMD XDNA NPUs; fused kernel hits 3.62 TFLOP/s (2× IRON, 6.5× iGPU energy); convincing internally, but measurement rigor and baselines are thin.
Summary
This vendor-led case study compares four mappings of FlashAttention onto AMD XDNA 1 and XDNA 2 NPUs: layer-by-layer operator execution, two streaming dataflow designs (DATO, IRON), and a fully fused MLIR-AIR kernel that holds QK^T scores in compute-tile local memory and reduces partials over the cascade interconnect. On XDNA 2 the fused kernel sustains 3.62 TFLOP/s of effective GEMM throughput over end-to-end execution at 32k context, twice the IRON design, and reports 6.5x better energy efficiency than the same chip's integrated GPU (31.5x vs CPU); twelve LLM configurations from BERT to DeepSeek run up to 128k tokens. Three-level roofline analysis (shim, MemTile, compute-tile DMA) predicts which mappings are compute-bound before implementation, and measurements match the predictions; on XDNA 1 the lower ridge points make streaming sufficient, so extra fusion is nearly wasted. Limitations: the ablation uses GPT-2 Small dimensions only, DATO has no XDNA 2 numbers, energy measurement methodology is underspecified, latency is reported as minimum over 20 iterations with no variance, code links are anonymized for review, and numerical stability is validated only on synthetic Gaussian inputs.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- A fused MLIR-AIR FlashAttention kernel that keeps QK^T scores in compute-tile local memory and reduces partials over the cascade interconnect reaches 3.62 TFLOP/s of end-to-end effective GEMM throughput on XDNA 2, twice the IRON design (Table 5: 3622 vs 1805 GFLOP/s at 32768 context).
- The fused kernel achieves 6.5x the energy efficiency of the same chip's iGPU (and 31.5x that of the CPU) at 32k sequence length on the Ryzen AI 9 HX 370 (Section 6.4, Figure 3).
- Multi-level roofline analysis comparing a mapping's operational intensity against each memory level's ridge point predicts, before implementation, which mappings are compute-bound; measurements on both generations match these predictions (Tables 2-4, Section 6.3).
- Because XDNA 1's ridge points are lower, on-chip streaming (DATO/IRON) already reaches the compute-bound regime there, so the fusion that doubles throughput on XDNA 2 yields little benefit on XDNA 1 (Table 4, Sections 5.3 and 7.2).
- The fused mapping generalizes across twelve LLM configurations from BERT to DeepSeek, up to 131072 tokens, sustaining 2.0-3.7 TFLOP/s on XDNA 2 (Table 6).
Strengths
Systematic four-point ablation on two NPU generations
Section 4 defines layer-by-layer, DATO, IRON, and fused MLIR-AIR mappings spanning three distinct intermediate-placement levels (DDR spill, MemTile spill, compute-tile local); Section 6 measures all of them on XDNA 1 and XDNA 2, with IRON-versus-fused ratios of 1.95-2.01x across 1k-32k context on XDNA 2 (Table 5).
Predictive three-level roofline analysis, validated against measurements
Tables 2-4 derive per-level bandwidths, ridge points, and predicted regimes (e.g., XDNA 2 MemTile ridge 33.7 FLOP/B vs IRON OI <=32 predicting near-ridge); Section 6.3 and Figure 2 check predictions and explicitly flag where they are only approximately met, which strengthens credibility.
Headline numbers precisely qualified
Actionable design guidance tied to numbers
Broad generality sweep for the final design
Weaknesses
No dispersion statistics anywhere
Missing ablation arm: DATO on XDNA 2
Energy-efficiency measurement underspecified
External baselines and reproducibility thin
Compute ceiling derived from instruction counting
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
This is a vendor-authored case study (AMD with ETH Zurich and Cornell collaborators) on mapping scaled dot-product attention onto spatial NPUs, positioned explicitly as an experience report rather than a new compiler technique. The central thesis (Sections 4-5): where the QK^T intermediate lives — DDR, MemTile, or compute-tile local memory — is the first-order determinant of throughput, and a three-level roofline analysis predicts which mappings are compute-bound before any code is written. Headline claims: 3.62 TFLOP/s end-to-end effective GEMM throughput on XDNA 2 (2x the IRON design at 32k context, Table 5), 6.5x iGPU energy efficiency (Figure 3), and generality across twelve LLM configurations up to 131,072 tokens (Table 6).
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
How much faster is the fused FlashAttention kernel on AMD XDNA 2 than previous NPU implementations?
The paper reports the fused MLIR-AIR kernel at 3.62 TFLOP/s of effective GEMM throughput measured over complete end-to-end attention execution on XDNA 2, about twice the IRON streamed-dataflow design at every evaluated sequence length from 1024 to 32768 (ratios 1.95-2.01x), and up to 3.91x at 256 context.
How energy efficient is the XDNA 2 NPU for attention compared with the integrated GPU?
At 32768 sequence length, the paper reports 6.5 times better energy efficiency than the same Ryzen AI 9 HX 370 chip's integrated GPU running a Triton-based fused FlashAttention (ROCm 7.2), and 31.5 times better than the CPU. The measurement method (power boundary, tool) is described only as average device power from a profiler, which a reviewer would want specified.
Which LLM models and sequence lengths does the XDNA 2 fused attention kernel support?
Table 6 lists twelve configurations spanning BERT-Base/Large, GPT-2 Small/Medium/Large (also OPT-125M), Llama-2 7B, Llama-3 8B, Qwen1.5-0.5B (including a 131072-token run), Qwen2.5-0.5B-Instruct, and DeepSeek 7B/67B, covering MHA and GQA, causal and non-causal masking, d_k of 64 and 128, sustaining roughly 2.0-3.7 TFLOP/s.
When is fusing an attention kernel into one NPU kernel worth it?
The paper's rule: compute each candidate mapping's operational intensity and compare it with the ridge point (bandwidth divided by compute ceiling) at each memory level. Fuse only until the mapping clears the ridge point of the limiting level. On XDNA 2 the MemTile ridge is high enough that scores must stay in compute-tile memory (fusion pays); on XDNA 1 the ridge points are low, so streaming through MemTile already reaches the compute-bound regime and further fusion is nearly wasted.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures (1-4) are present only as extracted captions and axis labels; roofline placements and energy-efficiency curves could not be verified from the text.
- Code repository links are anonymized ('blocked for double-blind review'); the open-source release claim could not be verified.
- Minor text-extraction artifacts exist (footnote markers inside 'TFLOPGEMM/s', two duplicated reference pairs), but section structure and tables are complete.
- Energy-efficiency comparisons rely on Figure 3's plotted values; underlying per-length numbers are not tabulated.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.