Technical Summary: SwiftFusion

Problem Statement

Diffusion Transformers (DiTs) are increasingly used for high-quality image and video generation. However, generating high-resolution images or long-duration videos results in long sequence lengths, causing activation tensors to grow linearly with sequence length. This leads to memory constraints and high latency on single GPUs, necessitating distributed inference via Sequence Parallelism (SP).

Existing SP techniques, such as Ring Attention, Ulysses Attention, and their combination Unified Sequence Parallelism (USP), face three primary limitations when scaling across multiple GPU machines:

1. Suboptimal Topology Utilization: Current frameworks often fail to account for the significant bandwidth differences between intra-machine (e.g., NVSwitch) and inter-machine (e.g., EFA/RDMA) networks. For instance, USP uses Ring Attention for inter-machine communication, which incurs high communication volume that does not decrease with more machines, creating a bottleneck.
2. Non-Overlapping Communication: Techniques relying on Ulysses Attention utilize all-to-all operators that are treated as atomic operations. These operations block computation, creating latency bubbles, particularly problematic for inter-machine communication where bandwidth is lower.
3. Synchronization and Computation Overheads: Existing implementations rely on two-sided communication libraries (e.g., NCCL), which require strict sender-receiver synchronization. This introduces implicit barriers and consumes GPU Streaming Multiprocessors (SMs) for kernel-based communication, leading to contention with model computation.

Methodology

The authors propose SwiftFusion, a topology-aware efficient DiT serving engine designed to address these challenges through three key innovations:

1. Topology-Aware Communication Scheduling

SwiftFusion rethinks the assignment of SP techniques based on network topology. Unlike USP, which uses Ring Attention for inter-machine communication, SwiftFusion inverts this strategy:

• Inter-Machine: It employs Ulysses Attention. While Ulysses has high communication volume per machine in some contexts, its all-to-all volume decreases as the number of machines increases, making it more suitable for slower inter-machine networks when scaled.
• Intra-Machine: It employs Ring Attention. Since intra-machine networks (e.g., NVSwitch) offer high bandwidth and low latency, the higher communication volume of Ring Attention is less of a bottleneck here.
  This scheduling minimizes the total inter-machine communication volume, leveraging the high-bandwidth intra-machine links for the more communication-intensive Ring operations.

2. Torus Attention (Communication-Computation Overlap)

To mitigate the latency of the all-to-all operators required by Ulysses Attention (which are non-overlapping in standard implementations), SwiftFusion introduces Torus Attention.

• Chunking: The all-to-all operators are broken down into multiple chunks.
• Pipelining: The system identifies "stationary" elements in the tensors that do not need to be transferred during specific stages of the all-to-all operation. These elements are used immediately for attention computation.
• Overlap: While the stationary elements are being computed, the system initiates the transfer of the remaining chunks in the background. This creates a pipeline where the communication of one chunk overlaps with the computation of the previously received chunk, effectively hiding inter-machine communication latency.

3. One-Sided Communication Implementation

SwiftFusion replaces traditional two-sided libraries (like NCCL) with one-sided communication libraries (e.g., NVSHMEM).

• Mechanism: Instead of send-receive primitives requiring mutual synchronization, it uses push and pull primitives.
• Optimization: The authors design a unified communication algorithm that unifies Torus, Ulysses, and Ring Attention. This approach minimizes synchronization barriers, requiring only intra-machine synchronizations and two specific inter-machine synchronizations (at the start and end of a layer).
• Efficiency: By removing per-transfer synchronization and avoiding CUDA kernel-based communication that competes for SMs, the system reduces both synchronization and computation overheads.

Key Contributions

The paper summarizes its contributions as follows:

1. Topology-Aware Scheduling: A novel strategy that better exploits modern GPU machine topologies by assigning Ulysses Attention to inter-machine and Ring Attention to intra-machine communication.
2. Torus Attention: A new distributed attention algorithm that breaks down all-to-all operators to pipeline communication with computation, hiding inter-machine latency.
3. One-Sided Implementation: A specialized implementation unifying Torus, Ulysses, and Ring Attention using one-sided libraries to minimize synchronization and computation overheads.
4. Empirical Evaluation: Comprehensive evaluation on image (Flux) and video (CogVideoX) generation workloads demonstrating state-of-the-art performance.

Results

The authors evaluated SwiftFusion on AWS p4de.24xlarge instances (8x NVIDIA A100 GPUs per machine, interconnected via NVSwitch and EFA) against the state-of-the-art USP framework.

• Overall Performance: SwiftFusion achieved an average speedup of 1.35× over USP, with a maximum speedup of 1.77× across various distributed configurations and workloads (Flux 3072×3072/4096×4096, CogVideoX 20s/40s).
• Component Analysis:
  • The topology-aware scheduling alone (TAS) provided a 1.27× speedup over USP when using more than two GPU machines, primarily by reducing inter-machine communication volume.
  • The addition of Torus Attention and one-sided communication (SwiftFusion) provided an additional 1.35× speedup over TAS, attributed to the overlapping of communication and computation and reduced synchronization overhead.
• Scalability: SwiftFusion consistently outperformed USP across different numbers of GPU machines (up to 4 machines) and various distributed configurations (varying Ulysses and Ring parallelization degrees).
• Memory: The framework did not introduce higher memory consumption compared to USP, utilizing at most one copy buffer for tensors.
• Workload Generalization: Performance gains were consistent across different sequence lengths, batch sizes, and head dimensions, though gains were slightly more pronounced with larger head dimensions where computation saturation is easier to achieve.

Significance

The paper claims that SwiftFusion provides a new paradigm for efficient and scalable distributed DiT inference. By explicitly addressing the mismatch between network topology and communication patterns, and by introducing mechanisms to hide communication latency and reduce synchronization overhead, SwiftFusion overcomes the communication-bound limitations of current SP techniques. The authors posit that these insights are critical for scaling DiT inference to the high-resolution and long-duration requirements of future generative AI applications.