PipeFusion: Patch-level Pipeline Parallelism for Diffusion Transformers Inference

TL;DR

PipeFusion introduces patch-level pipeline parallelism for Diffusion Transformers to curb latency and memory costs during high-resolution image generation. By partitioning both patches and DiT layers across GPUs and reusing stale features from previous diffusion steps, it reduces inter-device communication while maintaining generation quality. Empirical results on PixArt, Stable-Diffusion 3, and Flux.1 show state-of-the-art latency and favorable memory profiles, with strong scalability on multi-GPU clusters. The work suggests that exploiting input temporal redundancy enables practical scaling and invites hybrid parallel configurations for even larger deployments.

Abstract

This paper presents PipeFusion, an innovative parallel methodology to tackle the high latency issues associated with generating high-resolution images using diffusion transformers (DiTs) models. PipeFusion partitions images into patches and the model layers across multiple GPUs. It employs a patch-level pipeline parallel strategy to orchestrate communication and computation efficiently. By capitalizing on the high similarity between inputs from successive diffusion steps, PipeFusion reuses one-step stale feature maps to provide context for the current pipeline step. This approach notably reduces communication costs compared to existing DiTs inference parallelism, including tensor parallel, sequence parallel and DistriFusion. PipeFusion enhances memory efficiency through parameter distribution across devices, ideal for large DiTs like Flux.1. Experimental results demonstrate that PipeFusion achieves state-of-the-art performance on 8$\times$L40 PCIe GPUs for Pixart, Stable-Diffusion 3, and Flux.1 models. Our source code is available at https://github.com/xdit-project/xDiT.

Figures (11)

• Figure 1: Workflow of DiTs inference.
• Figure 2: Comparison of DistriFusion with Sequence Parallelism methods (DeepSpeed-Ulysses and Ring-Attention) for attention layers.
• Figure 3: Above: partitioning strategy for Input and DiTs backbone network. Below: Workflow of the PipeFusion as patch-level pipelined parallelism.
• Figure 4: The fresh part of activations during diffusion timestep $T$ of Figure \ref{['fig:pipefusion']}. The dark gray represents fresh data and the light gray represents stable data.
• Figure 5: Latency on Pixart of various parallel approaches on two image generation tasks with the 20-Step DPMSolverMultistepScheduler.