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

TL;DR

Zen-Attention is a framework that optimizes the deployment of transformer attention layers on AMD NPUs, significantly improving latency and efficiency by exploring layer folding, tiling, and data movement strategies.

Contribution

It introduces a systematic framework for optimizing dynamic attention layer mapping on AMD NPUs, addressing complex design space challenges for better performance and energy efficiency.

Findings

Up to 4x reduction in attention layer latency

Up to 32% improvement in end-to-end network latency

Enhanced mapping capabilities for varying input dimensions

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.