DiTango: Cost-Effective Parallel Diffusion Generation with Selective Attention State Reuse
Computer Vision and Pattern Recognition
Summary
The authors looked at how making AI-generated images using Diffusion Transformers can be slow and costly when using many computers. They found that parts of the image close to each other are more important to calculate together, so they created a method to focus computation on these key parts and reuse past results. Their new system, DiTango, smartly plans which parts to compute and which to reuse, speeding up the process significantly without losing image quality. Tests show DiTango is faster than previous methods, especially when using multiple computers.
Authors
Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, Jidong Zhai
Abstract
Recent advances in AI-generated content have driven widespread adoption of Diffusion Transformers (DiTs) for high-resolution, long-duration content generation. While parallelization techniques accelerate diffusion inference, they face significant scalability challenges due to excessive communication overhead in multi-node environments. We observe that sequence partitions in Context Parallelism (CP) exhibit distinct heterogeneity: spatially proximate partitions contribute more significantly to attention computation results. By mapping this heterogeneous pattern to hierarchical communication topology, we can access high-contribution partitions with reduced communication cost. This insight motivates our novel selective attention state mechanism that strategically balances partial attention computation and historical result reuse across denoising steps. We present DiTango, an efficient parallel framework for DiT generation. DiTango features an anchor-guided state selection planner that optimizes computation-reuse decisions for each partition, complemented by a runtime that orchestrates efficient state-centric operations. This design achieves superior system efficiency while preserving generation quality. Experimental evaluation on popular diffusion models demonstrates that DiTango achieves up to 1.9x end-to-end and 3.2x attention speedup with near-linear scaling in multi-node settings, while maintaining generation quality comparable to state-of-the-art approaches.