💻 computer science

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

DiTango is a cost-effective parallel diffusion framework that accelerates multi-node DiT generation by leveraging the heterogeneity of context partitions to strategically reuse attention states, achieving up to 3.2x speedup with near-linear scaling while preserving generation quality.

Original authors: Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, Jidong Zhai

Published 2026-07-20

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, Jidong Zhai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers can dream up entire movies, not just still pictures, but flowing, high-definition stories that last for minutes. This is the magic of "Diffusion Transformers," a type of artificial intelligence that builds images and videos step-by-step, slowly turning static noise into a clear, beautiful scene. Think of it like a sculptor chipping away at a block of marble; the AI starts with a chaotic cloud of pixels and, over many rounds of refinement, reveals a perfect video.

However, there's a catch: these digital sculptors are incredibly slow. Creating a single five-second video can take over an hour on a powerful computer. To speed things up, engineers often try to get many computers to work together, splitting the long video sequence into chunks and handing them out like a relay race. But here's the problem: when these computers try to share their work, they get stuck in traffic jams. The time spent waiting to pass data back and forth often cancels out the speed gained by having more workers. This creates a frustrating bottleneck where adding more computers doesn't actually make the video finish any faster, and sometimes even makes it slower.

Enter DiTango, a new system designed to fix this traffic jam. The researchers behind DiTango noticed something fascinating about how these AI models "pay attention" to different parts of the video. They found that not all parts of the video are equally important at every moment. Just like how you focus intently on the person speaking right in front of you but barely notice the person standing three rooms away, the AI's attention is strongest for nearby parts of the video and much weaker for distant ones.

DiTango uses this insight to play a clever game of "skip and reuse." Instead of forcing every computer to constantly fetch and calculate every single piece of data from every other computer (which causes the traffic jam), DiTango's smart planner decides which parts are crucial to calculate fresh and which parts are so unimportant that the computer can just use an old, cached result from a few seconds ago. It's like a group of chefs in a massive kitchen: instead of everyone running to the pantry to grab the same spice every time, they realize that for the garnish on the far side of the table, the spice jar sitting on the counter nearby is "good enough" and doesn't need a fresh trip.

By being selective, DiTango dramatically cuts down on the time computers spend talking to each other. In tests using powerful video models, this approach made the generation process nearly twice as fast end-to-end and sped up the core "attention" calculations by more than three times when using 32 computers together. Crucially, the researchers found that this speedup didn't ruin the quality of the video; the final movies looked just as sharp and coherent as those made by slower, traditional methods. DiTango proves that by understanding where the AI actually needs to look, we can stop it from wasting energy on things it barely notices, making the dream of instant, high-quality AI video generation a little closer to reality.

Technical Summary: DiTango

Problem Statement

Diffusion Transformers (DiTs) have become the dominant architecture for high-resolution, long-duration AI-generated content (AIGC). However, their inference latency is prohibitive, particularly for long sequences (e.g., 720p videos), because they require full attention computation over the entire spatiotemporal sequence at every denoising step. Unlike autoregressive LLMs, DiTs cannot amortize attention costs via standard KV caching across steps.

While parallel inference offers a path to acceleration, existing strategies face significant scalability bottlenecks in multi-node environments:

Tensor Parallelism (TP) and Context Parallelism (CP) rely heavily on communication. TP requires frequent collective operations for sharded weights, while CP variants (Ring-CP, Head-CP) necessitate extensive Key-Value (KV) partition exchanges (All-to-All or P2P) to access global sequences.
Scalability-Quality Dilemma: In multi-node settings, limited inter-node bandwidth causes severe performance degradation. To mitigate this, existing frameworks often employ aggressive feature reuse (skipping computations), which inevitably leads to substantial generation quality degradation. Conversely, preserving quality requires full communication, resulting in poor scaling efficiency.

The paper identifies that current approaches treat parallelization and feature reuse as independent optimizations, failing to leverage the specific structural properties of DiT attention.

Methodology

DiTango addresses these challenges by exploiting a key observation: spatial locality in attention contributions. The authors demonstrate that attention contributions from different sequence partitions are highly heterogeneous; partitions spatially proximate to the query contribute significantly more to the final result than distant partitions.

DiTango introduces a Selective Attention State Reuse mechanism that aligns this computational heterogeneity with the hierarchical communication topology of distributed systems. The system consists of two primary components:

1. Anchor-Guided Selection Planner

This component determines which partitions should be computed fresh and which can reuse historical results.

Error Modeling: The system models the accumulated error of reusing an attention state ($AS$) from a previous timestep. It defines the error as a function of partition importance ( $w$ ) and temporal error growth ( $\delta$ ).
Anchor Steps: To prevent unbounded error accumulation, the system enforces periodic "anchor steps" where all partitions are computed fresh. These steps reset error accumulation and update partition importance weights.
Online Prediction: Between anchor steps, the planner predicts error growth using local state changes (which are always computed) as a proxy for remote state evolution. This allows for precise, low-overhead decision-making.
Group-wise Strategy: To manage the complexity of selecting individual partitions, the system groups spatially adjacent partitions. Decisions are made at the group level to balance accuracy and efficiency, ensuring that high-contribution partitions are computed while low-contribution, high-cost partitions are reused.

2. State-Centric Parallel Runtime

This component orchestrates the execution of heterogeneous compute/reuse decisions to maximize overlap between communication and computation.

State Composition: DiTango utilizes the composability of attention states (output and log-sum-exp). Instead of caching raw KV pairs (which is memory-intensive), it caches composed attention states. These states can be combined associatively and commutatively.
Cross-Group KV Transfer: When a group requires fresh computation, the runtime performs symmetric P2P exchanges to relocate necessary KV partitions to the devices closest to where they will be consumed, optimizing the use of intra-node (NVLink) vs. inter-node (InfiniBand) bandwidth.
Group State Computation: Within a group, devices perform ring-based composed attention. This overlaps $(g-1)$ KV transfers with attention computation, producing a full composed state for the group.
Dynamic Group Compose: To manage memory pressure, the runtime can recursively merge cached attention states into coarser-grained groups, exponentially reducing memory overhead without recomputation.
Pipeline Orchestration: The runtime decouples operations into asynchronous computation and communication streams. It strategically schedules lightweight "reuse" operations (loading cached states) to fill computation bubbles caused by cross-node communication latency.

Key Contributions

Identification of Spatial Locality: The paper reveals that DiT attention contributions exhibit strong spatial locality, where nearby partitions dominate computation. This pattern is intrinsic to diffusion attention and persists across models and prompts.
DiTango Framework: The design and implementation of a high-performance generation system that leverages theoretical error modeling to make precise reuse decisions. It introduces an anchor-guided planner and a state-centric runtime to manage the trade-off between communication cost and generation quality.
Scalability and Quality Preservation: Extensive evaluations demonstrate that DiTango achieves near-linear scaling across 32 GPUs while maintaining generation quality comparable to state-of-the-art approaches, effectively breaking the scalability-quality trade-off inherent in previous methods.

Experimental Results

The authors evaluated DiTango on popular open-source models (Wan2.1-14B, Wan2.1-1.3B, HunyuanVideo) using 32 NVIDIA H20 GPUs.

Speedup: DiTango achieves up to 1.9× end-to-end speedup and 3.2× core attention speedup compared to state-of-the-art baselines (including TP, CP, xDiT, VideoSys, and SGLang-Diffusion) in multi-node settings.
Scaling: The system demonstrates near-linear scaling up to 32 GPUs, whereas baselines often suffer from communication bottlenecks or quality degradation in multi-node scenarios.
Quality: DiTango maintains generation quality comparable to full-attention baselines. In terms of VBench scores, PSNR, and SSIM, DiTango outperforms lossy acceleration frameworks (like VideoSys and SGLD) which suffer from quality degradation due to coarse-grained skipping.
Efficiency: The system achieves high FLOP/s in multi-node settings by minimizing cross-node communication, whereas traditional parallel approaches are often communication-bound.

Significance and Claims

The paper claims that DiTango represents a significant step forward in distributed diffusion inference by integrating algorithmic insights (spatial locality) with system-level orchestration.

Overcoming the Trade-off: The primary significance is the ability to achieve high scalability in distributed environments without sacrificing generation quality, a limitation that has plagued previous parallel diffusion systems.
System-Algorithm Alignment: The work highlights the importance of aligning computational patterns (spatial locality) with system topology (hierarchical bandwidth) to unlock optimization spaces previously unreachable.
Practical Deployment: By demonstrating near-linear scaling and second-level latency potential for long video generation, DiTango is positioned as a viable solution for real-time and interactive generative applications on large-scale GPU clusters.

The authors acknowledge limitations, noting that in high-bandwidth single-node environments, pure computation reduction might yield higher throughput, and that memory overhead for state caching, while manageable, remains a constraint for production environments. Future work is suggested to focus on cache compression and CPU offloading.