DiTango通过选择性复用注意力状态,让扩散模型在多节点下更快生成高质量内容。
DiTango: Cost-Effective Parallel Diffusion Generation with Selective Attention State Reuse

- 按空间远近选择性复用注意力状态,减少通信开销。
- 多节点环境下实现3.2倍注意力计算加速,接近线性扩展。
- 适合需要高效生成高分辨率图像或视频的科研与工业用户。
近年来,AI生成内容推动了扩散变换器(DiTs)在高分辨率、长时长内容生成中的广泛应用。尽管并行化技术能加速扩散推理,但在多节点环境中仍面临通信开销过大的可扩展性挑战。我们观察到上下文并行(CP)中序列分块存在显著异质性:空间邻近的分块对注意力计算贡献更大。将这一异质模式映射到分层通信拓扑,可更低成本地访问高贡献分块。据此提出一种选择性注意力状态机制,在去噪步骤间策略性平衡部分计算与历史结果复用。我们提出DiTango,一种高效的DiT生成并行框架。DiTango包含锚点引导的状态选择规划器,优化每个分块的计算-复用决策,并由运行时协调高效的状态中心操作。该设计在保持生成质量的同时显著提升系统效率。在主流扩散模型上的实验表明,DiTango在多节点设置下实现最高1.9倍端到端加速和3.2倍注意力速度提升,且接近线性扩展,生成质量与当前最优方法相当。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。