DiTango: Cost-Effective Parallel Diffusion Generation with Selective Attention State Reuse

2026-07-17

AI Summary

This work addresses the challenge of inefficient scaling in multi-node diffusion model inference, which is hindered by substantial communication overhead. To mitigate this bottleneck, the authors propose a spatial proximity–aware selective attention state reuse mechanism that dynamically balances computation and reuse of historical states during denoising steps to reduce inter-node communication. The core innovation lies in the co-design of a novel anchor-guided state selection planner and a state-centric runtime system, enabling efficient computation-communication scheduling. Experimental results demonstrate that the proposed approach achieves up to 1.9× end-to-end speedup and 3.2× acceleration in attention computation across multiple nodes, exhibiting near-linear scalability while preserving generation quality.

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.

Problem

• Diffusion Transformers
• parallel generation
• communication overhead
• scalability
• multi-node

Innovation

• Selective Attention State Reuse
• Context Parallelism
• Diffusion Transformers
• Multi-node Scaling
• Anchor-guided Planning