Zen-Attention: A Compiler Framework for Dynamic Attention Folding on AMD NPUs

Abstract

Transformer-based deep learning models are increasingly deployed on energy, and DRAM bandwidth constrained devices such as laptops and gaming consoles, which presents significant challenges in meeting the latency requirements of the models. The industry is turning to neural processing units (NPUs) for superior performance-per-watt (perf/watt); however, efficiently mapping dynamic attention layers to the NPUs remains a challenging task. For optimizing perf/watt, AMD XDNA NPUs employ software managed caches and share system memory with host. This requires substantial engineering effort to unlock efficient tiling, buffer allocation, and data movement to extract the maximum efficiency from the device. This paper introduces Zen-Attention, a framework that optimizes DRAM bandwidth utilization in the attention layer of models by systematically exploring the complex design space of layer folding, tiling, and data-movement on the interconnect, and the tensor layouts to come up with an optimal solution. Our evaluation includes comparative analysis of end-to-end model latency and specific attention latency in each model. We demonstrate how the framework enhances mapping capabilities by varying input dimensions, which require padding and masking in the attention block. For representative transformer models, the Zen-Attention Framework achieves up to 4x improvement in the latency of the attention block and up to 32% improvement in end-to-end network latency compared to the baseline Unfolded- approaches.