DiTFastAttnV2: Head-wise Attention Compression for Multi-Modality Diffusion Transformers
This work addresses efficiency and scalability issues in text-to-image generation models, offering an incremental improvement through compression techniques.
The paper tackles the computational bottleneck in attention mechanisms of Multimodal Diffusion Transformers (MMDiT) for text-to-image generation by introducing DiTFastAttnV2, a post-training compression method that achieves a 68% reduction in attention FLOPs and 1.5x end-to-end speedup on 2K image generation while maintaining visual fidelity.
Text-to-image generation models, especially Multimodal Diffusion Transformers (MMDiT), have shown remarkable progress in generating high-quality images. However, these models often face significant computational bottlenecks, particularly in attention mechanisms, which hinder their scalability and efficiency. In this paper, we introduce DiTFastAttnV2, a post-training compression method designed to accelerate attention in MMDiT. Through an in-depth analysis of MMDiT's attention patterns, we identify key differences from prior DiT-based methods and propose head-wise arrow attention and caching mechanisms to dynamically adjust attention heads, effectively bridging this gap. We also design an Efficient Fused Kernel for further acceleration. By leveraging local metric methods and optimization techniques, our approach significantly reduces the search time for optimal compression schemes to just minutes while maintaining generation quality. Furthermore, with the customized kernel, DiTFastAttnV2 achieves a 68% reduction in attention FLOPs and 1.5x end-to-end speedup on 2K image generation without compromising visual fidelity.