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

2026-07-10

AI Summary

This work addresses the inefficiencies of large language model inference on laptop-class SoCs, where limited NPU support for attention mechanisms leads to poor energy efficiency, high data movement overhead, and load imbalance caused by causal masking. To overcome these challenges, the authors propose STEEL—the first open-source FlashAttention implementation tailored for XDNA-class NPUs. STEEL introduces a novel sparsity-aware pipelined scheduling strategy to mitigate load imbalance and designs a fused attention dataflow optimized for the XDNA architecture, effectively leveraging on-chip memory and spatial parallelism. Evaluated on the AMD Ryzen AI 9 HX 370 (XDNA 1/2), STEEL achieves average energy savings of 9.17× and 1.75× over CPU and GPU baselines, respectively, reduces latency by 9.6× compared to the previous best solution, and accelerates inference by 22.8× over a layer-wise implementation.

Abstract

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.

Problem

• energy-efficient inference
• attention mechanism
• neural processing unit
• load imbalance
• long-sequence

Innovation

• Sparsity-Aware Attention
• Fused Attention
• NPU Dataflow
• Energy-Efficient Inference
• Long-Sequence Processing