Search papers, labs, and topics across Lattice.
This paper introduces APT, a software-hardware co-designed accelerator that optimizes high-resolution Diffusion Transformers (DiTs) by leveraging attention probabilities for fine-grained pruning and adaptive precision scaling. The proposed Attention Probability-guided Adaptive Dual Thresholding (APDT) method enhances computational efficiency while maintaining model performance, addressing the significant latency issues associated with self-attention in DiTs. Experimental results demonstrate that APT achieves up to 8.16脳 speedup and 14.98脳 higher energy efficiency compared to NVIDIA A100, marking a substantial improvement in the operational efficiency of state-of-the-art diffusion models.
APT accelerates high-resolution diffusion models by up to 8.16脳 while enhancing energy efficiency, revolutionizing the feasibility of real-time generative AI applications.
Recent advances in generative AI have significantly increased the demand for high-resolution image and video generation, positioning diffusion models as a core technology. Among them, Diffusion Transformers (DiTs) have emerged as the state-of-the-art (SOTA) models due to their scalability and output quality. However, self-attention in DiTs incurs significant computational overhead, leading to excessively long latency as the complexity grows with the fourth power of the output resolution. While prior works have attempted to mitigate this cost using sparsity and quantization techniques, they fall short of effectively reducing the computational cost in high-resolution DiTs. In this paper, we present APT, a software-hardware co-designed accelerator for high-resolution DiTs. APT leverages attention probabilities as a unified importance metric to jointly optimize computation through fine-grained pruning and adaptive precision scaling. At the algorithm level, we propose Attention Probability-guided Adaptive Dual Thresholding (APDT), which dynamically performs element selection and precision assignment using dual thresholds. To ensure compatibility with memory-efficient FlashAttention, we introduce Timestep-Aware FlashAttention (TAFA), which predicts attention probabilities across timesteps by exploiting temporal similarity. At the architecture level, we co-design a specialized accelerator that efficiently supports irregular sparsity and dual-precision execution, featuring dynamic mask management, address translation, dual-precision compute units, and a tile-based dataflow. Finally, we evaluate APT on SOTA DiT models, including PixArt-$\alpha$, Stable Diffusion 3, and FLUX. APT achieves up to 8.16$\times$ speedup and 14.98$\times$ higher energy efficiency over NVIDIA A100, and up to 3.01$\times$ speedup and 2.04$\times$ higher energy efficiency over EXION, a SOTA diffusion model accelerator.
