PaperlayerView exact revision on arXiv

Source-linked AI summary

xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism

Jiarui Fang, Jinzhe Pan, Xibo Sun, Aoyu Li, Jiannan Wang

arXiv:2411.01738v1cs.DCcs.AI

TL;DR

Long-sequence DiTs create substantial inference demands, while scaling them across varied hardware remains challenging. xDiT combines intra-image and inter-image parallel strategies, and its evaluation reports communication and memory advantages for PipeFusion alongside efficient scaling across large clusters.

  • Problem

    DiT attention costs grow with sequence length, and scaling inference across heterogeneous networks remains an open challenge for large-scale deployment.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    xDiT combines Sequence Parallel and PipeFusion for intra-image parallelism with CFG parallelism for inter-image generation, supporting hybrid configurations.

    Source: §4 · p. 4

  • Results

    xDiT efficiently scales inference across four image-generation and one video-generation DiT on 16×L40 and 8×A100 configurations, with PipeFusion showing communication and memory advantages over other methods.

    Source: §6 · p. 13

  • Takeaways & Limitations

    The evaluation demonstrates xDiT's ability to scale DiT inference to large-scale clusters.

    Source: §6 · p. 13

    For the 4096px task, DistriFusion encounters out-of-memory issues as its KV Buffer grows linearly with sequence length, so it is not used as the foundational parallel method.

    Source: §5.2.1 · p. 8

Abstract

from arXiv · show

Diffusion models are pivotal for generating high-quality images and videos. Inspired by the success of OpenAI's Sora, the backbone of diffusion models is evolving from U-Net to Transformer, known as Diffusion Transformers (DiTs). However, generating high-quality content necessitates longer sequence lengths, exponentially increasing the computation required for the attention mechanism, and escalating DiTs inference latency. Parallel inference is essential for real-time DiTs deployments, but relying on a single parallel method is impractical due to poor scalability at large scales. This paper introduces xDiT, a comprehensive parallel inference engine for DiTs. After thoroughly investigating existing DiTs parallel approaches, xDiT chooses Sequence Parallel (SP) and PipeFusion, a novel Patch-level Pipeline Parallel method, as intra-image parallel strategies, alongside CFG parallel for inter-image parallelism. xDiT can flexibly combine these parallel approaches in a hybrid manner, offering a robust and scalable solution. Experimental results on two 8xL40 GPUs (PCIe) nodes interconnected by Ethernet and an 8xA100 (NVLink) node showcase xDiT's exceptional scalability across five state-of-the-art DiTs. Notably, we are the first to demonstrate DiTs scalability on Ethernet-connected GPU clusters. xDiT is available at https://github.com/xdit-project/xDiT.

1 Introduction

DiTs face growing inference demands because attention cost scales quadratically with sequence length and generation requires multiple denoising steps. The paper introduces xDiT, which combines parallel strategies for different hardware and model settings.

1 Introduction

  • DiT attention computation grows quadratically with sequence length, while all denoising steps must finish before the output is produced.A Sora-like video generation API takes over four minutes to generate a video of a few seconds.Source: §1 · p. 1
  • Existing parallel methods can be mismatched to device interconnects: collective communication suits NVLink, while peer-to-peer methods suit PCIe or Ethernet.Using these methods in isolation does not adapt to heterogeneous interconnects.Source: §1 · p. 1
  • DiT architectures vary in instruction conditioning, block connections, and Transformer operator layouts, making direct application of LLM parallel methods potentially unsuitable.The paper argues a single parallel method is unlikely to scale effectively across model architectures and network hardware.Source: §1 · pp. 1–2
  • xDiT successfully scales DiT inference to 16 GPUs across five image and video generation models.The evaluated models are Pixart, Stable-Diffusion3, Flux.1, HunyuanDiT, and CogVideoX.Source: §1 · p. 2
  • xDiT selects PipeFusion and sequence parallelism for intra-image work and CFG parallelism for inter-image work, then combines them in a hybrid system.The authors report systematically investigating existing DiT parallel methods before selecting these approaches.Source: §1 · p. 2

2 Background & Related Works

This section outlines diffusion denoising and the Transformer architecture, then reviews temporal redundancy and parallel methods used to accelerate inference. It contrasts their communication and memory trade-offs.

2 Background & Related Works

  • Diffusion models iteratively predict noise from a noisy input and condition, then use a sampler-specific update to obtain the preceding denoising step.After repeated steps, the latent output is decoded to pixel space, and forward propagation through the noise-prediction model is the predominant inference-latency contributor.Source: §2 · p. 2
  • DiTs encode inputs as latent-space tokens and use Transformer self-attention to model relationships among patches, with attention computation growing quadratically as sequences lengthen.Conditioning methods include adaptive layer norm, cross-attention, and extra input tokens.Source: §2 · p. 2
  • Temporal redundancy across diffusion steps motivates methods that cache or reuse activations, including DiT-specific approaches such as Δ-DiT and PAB.Other examples include DeepCache and Tgate for U-Nets and DiT-FastAttn, which compresses attention using spatial, temporal, and conditional redundancies.Source: §2 · p. 2
  • Tensor parallelism partitions model parameters but incurs sequence-proportional AllReduce communication, whereas sequence parallelism partitions image tokens but stores full model parameters on each device.PipeFusion uses temporal redundancy for patch-level pipeline parallelism and has lower communication overhead and parameter-memory usage than SP and DistriFusion.Source: §2 · pp. 2–3

3 Challenges in Parallel DiTs Inference

DiT inference presents scaling challenges from rapidly growing computation, heterogeneous network connections, high memory demand, and diverse model architectures. The section describes why one parallel approach may not meet these varied requirements.

3 Challenges in Parallel DiTs Inference

  • DiT attention scales quadratically with sequence length, and increasing model size and sequence length for higher-quality generation yields cubic growth in deployment computation.The section gives examples ranging from 262 thousand tokens for Flux.1 at 1024px to 4.2 million tokens at 4096px.Source: §3 · p. 3
  • DiTs denoise extremely long sequences over multiple steps, unlike autoregressive LLMs, so LLM inference systems cannot be applied directly.DiTs and LLM prefill are compute-bound, while LLM decoding is memory-bound.Source: §3 · p. 3
  • A single parallel method is ill-suited to both RDMA+NVLink and Ethernet+PCIe settings, motivating hybrid parallelism for large-scale inference.The cited gap includes scaling across multiple GPU nodes while accommodating high- and low-end networks.Source: §3 · p. 3
  • Long sequences and large models create memory demands that sequence parallelism alone cannot address by distributing model parameters, while VAE decoding also requires substantial activation and operator-temporary memory.The section notes that the 12-billion-parameter open-source model reached at the time was expected to grow further.Source: §3 · p. 3
  • Condition-injection choices change partitioning needs, and MM-DiT blocks make tensor parallelism challenging by combining text and image latents in QKV projections.AdaLN-Zero and cross-attention require splitting image computation, while in-context conditioning also requires splitting caption-Transformer computation.Source: §3 · p. 3
  • Skip-connected DiT blocks such as U-ViT and HunyuanDiT complicate communication patterns for pipeline parallelism.Source: §3 · p. 3

4 xDiT System

xDiT combines intra-image Sequence Parallelism and PipeFusion with inter-image CFG parallelism, and supports hybrid configurations plus patch-parallel VAE decoding. Its methods address both DiT attention workloads and VAE memory demands.

4 xDiT System

  • xDiT combines intra-image Sequence Parallelism and PipeFusion with inter-image CFG parallelism; these methods can be hybridized, while VAE decoding uses patch parallelism.The inference pipeline comprises a Text-Encoder, Transformers Model, and VAE, with most computation in the Transformers Model.Source: §4 · p. 4

4.1 Intra-Image Parallelism for Transformer Model

  • For In-Context Conditioning, xDiT splits both condition and image sequences into corresponding shards, enabling local computation that matches the serial result.Standard sequence parallelism partitions visual tokens into patches; the tailored method also parallelizes condition encoding.Source: §4.1.1 · pp. 4–5
  • PipeFusion assigns consecutive transformer layers to GPUs and divides the latent image into patches, using stale activations and asynchronous P2P transfers to overlap computation and communication.The pipeline uses temporal redundancy so devices can begin before receiving all spatial activations for the current timestep.Source: §4.1.2 · p. 5
  • Skip-connection structures can impair PipeFusion’s communication overlap; the paper reports poor scalability on 2048px tasks on 8×A100.Such structures can require communication between non-adjacent devices as well as neighboring ones.Source: §4.1.2 · p. 5
  • PipeFusion has the lowest communication cost when N < 2L and overlaps communication with computation; its KV-buffer memory decreases with device count, unlike DistriFusion’s.Both methods maintain KV buffers, which can impose substantial activation-memory overhead for long sequences.Source: Table 1 · p. 5, §4.1.3 · p. 6
  • Hybrid SP+PipeFusion retains intermediate K,V results from SP attention so devices in an SP group use consistent values, with only minor SP-algorithm modifications and no added overhead.Naive hybridization can leave parts of each device’s K,V buffer stale.Source: §4.1.4 · p. 7

4.2 Inter-Image Parallelism

  • CFG parallelism computes the unconditional and text-guided latents separately, then all-gathers latent results after each diffusion step and before scheduler execution.Its communication overhead is much smaller than that of PipeFusion and sequence parallelism.Source: §4.2 · p. 7

4.3 Parallel VAE

  • 7168px images—more than 12.25 times larger than the naive VAE approach on 8×L40 (48GB)—are supported by combining patch parallelism with staged convolution execution.Patch parallelism reduces peak intermediate-activation memory; staged execution targets convolution temporary-memory spikes.Source: §4.3 · pp. 7–8

5 Experiments

Across five DiTs and two GPU-cluster settings, xDiT’s experiments show that hybrid parallelism adapts to model and hardware differences, while PipeFusion offers strong latency and memory efficiency. The VAE and image-quality evaluations show separate practical trade-offs and checks for parallel decoding.

5.2 Backbone Performance

  • 13.29× speedup on Pixart’s 4096px task let xDiT reduce latency from 245 seconds to 17 seconds on 16 GPUs.The best hybrid configuration for this task was ulysses=8, cfg=2.Source: §5.2.1 · p. 8
  • For CogVideoX, xDiT achieved a 6.0× speedup across 12 L40s connected via Ethernet, reducing video-generation latency from over five minutes to just 52 seconds.Source: §5.2.1 · pp. 9–10
  • On A100s, Pixart’s 4096px task achieved 8.4×, 4.1×, and 2.0× speedups over the 1 GPU baseline; 8-GPU latency fell from 186 seconds to 22 seconds.For Pixart at 1024px, the optimal hybrid configurations allocated CFG parallelism and the remaining degree to PipeFusion.Source: §5.2.2 · p. 10
  • On 8×A100 GPUs, SD3’s lowest-latency single method is PipeFusion at 1024px, and at 2048px it is PipeFusion on 2 GPUs and SP-Ulysses on 4 and 8; Flux’s is SP-Ulysses at both sizes.Source: §5.2.2 · p. 10
  • For Pixart, DistriFusion’s memory overhead grows with image size; for SD3 and Flux, PipeFusion uses 32% and 36% of SP memory on Flux.1 at 1024px and 2048px.PipeFusion stores only 1/pipefusion_degree of the parameters, but its stage-local KV buffers raise “others” memory.Source: §5.2.3 · p. 11
  • Across five DiTs tested on 16×L40 Ethernet and 8×A100 NVLink, PipeFusion showed strong efficiency, while hybrid parallelism was needed for hardware- and model-specific scaling.At 16×L40, single-method approaches scaled poorly; CFG parallelism reduced inter-node communication to one latent-image AllGather per diffusion step.Source: §5.2.1 · p. 8, §5.1 · p. 8, §5.2.4 · p. 11

5.3 Parallel VAE Performance Results

  • At 2048px, the 1-GPU VAE baseline supports decoding, while 8-GPU parallel VAE raises the maximum supported resolution to 7168px on L40 and 8192px on A100.Parallel VAE generally increases latency for a fixed input, so the authors recommend it to address OOM rather than accelerate decoding.Source: Table 3 · p. 12, §5.3 · pp. 12–13

5.4 Quality Results

  • Across Pixart’s tested parallel settings, generated images were visually indistinguishable from baseline to human observers, and hybrid PipeFusion+USP did not worsen FID.The evaluation used 1024px Pixart-alpha images and one warmup iteration; the authors caution that FID may not reflect visual quality.Source: Figure 19 · p. 12, §5.4 · p. 13

6 Conclusion

xDiT combines adapted Sequence Parallel, PipeFusion, and CFG parallelism in a hybrid inference engine, with patch parallelism for the VAE. Evaluations across five image and video DiTs show PipeFusion’s communication and memory advantages and xDiT’s ability to scale inference to large clusters.

6 Conclusion

  • xDiT adapts Sequence Parallel and PipeFusion across DiT architectures and flexibly combines them with CFG parallelism for hybrid inference.The paper identifies Sequence Parallel and PipeFusion as intra-image strategies and CFG parallelism as an inter-image strategy.Source: §6 · p. 13
  • Patch parallelism for the VAE module is implemented to prevent out-of-memory issues.Source: §6 · p. 13
  • Evaluations on 16×L40 and 8×A100 configurations show PipeFusion’s communication and memory advantages and xDiT’s efficient scaling to large clusters.The evaluation covers four image-generation models and one video-generation model.Source: §6 · p. 13

Related papers

Loading 2411.01738v1…