STEEL: sparsity-aware fused attention reduces energy by 9.17x on AMD's XDNA NPU

Researchers present STEEL, the first open-source implementation of FlashAttention targeting XDNA-like neural processing engines (NPUs), to address the challenge of mapping attention mechanisms onto laptop-class SoCs.

• STEEL introduces a dataflow formulation of prefill attention to exploit spatial parallelism and on-chip memory.
• It uses sparsity-aware pipeline placement to handle load imbalance from causal masks, reducing synchronization overhead.
• Evaluated on the AMD Ryzen AI 9 HX 370 SoC, STEEL reduces energy consumption by an average of 9.17x compared to CPU baselines and 1.75x relative to GPU baselines.
• On XDNA 1 hardware, STEEL achieves a 9.6x latency reduction over the prior state of the art.
• It delivers a 22.8x speedup on average compared to layer-by-layer attention implementations on XDNA 2.

This work enables energy-efficient inference for large language model-based agents on edge devices, mitigating reliability and privacy concerns associated with cloud offloading.