xDiT: A Scalable Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism
xDiT: A Scalable Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism
Diffusion Transformers (DiTs) are driving advancements in high-quality image and video generation. With the escalating input context length in DiTs, the computational demand of the Attention mechanism grows quadratically! Consequently, multi-GPU and multi-machine deployments are essential to meet the real-time requirements in online services.
To meet real-time demand for DiTs applications, parallel inference is a must. xDiT is an inference engine designed for the parallel deployment of DiTs on a large scale. xDiT provides a suite of efficient parallel approaches for Diffusion Models, as well as computation accelerations.
The overview of xDiT is shown as follows.
Sequence Parallelism, USP is a unified sequence parallel approach proposed by us combining DeepSpeed-Ulysses, Ring-Attention.
PipeFusion, a sequence-level pipeline parallelism, similar to TeraPipe but takes advantage of the input temporal redundancy characteristics of diffusion models.
Data Parallel: Processes multiple prompts or generates multiple images from a single prompt in parallel across images.
CFG Parallel, also known as Split Batch: Activates when using classifier-free guidance (CFG) with a constant parallelism of 2.
The four parallel methods in xDiT can be configured in a hybrid manner, optimizing communication patterns to best suit the underlying network hardware.
As shown in the following picture, xDiT offers a set of APIs to adapt DiT models in huggingface/diffusers to hybrid parallel implementation through simple wrappers. If the model you require is not available in the model zoo, developing it by yourself is not so difficult; please refer to our Dev Guide.
We also have implemented the following parallel strategies for reference:
Cache method, including TeaCache, First-Block-Cache and DiTFastAttn, which exploits computational redundancies between different steps of the Diffusion Model to accelerate inference on a single GPU.
Optimization is orthogonal to parallel and focuses on accelerating performance on a single GPU.
First, xDiT employs a series of kernel acceleration methods. In addition to utilizing well-known Attention optimization libraries, we leverage compilation acceleration technologies such as torch.compile and onediff.
The following open-sourced DiT Models are released with xDiT in day 1.
ComfyUI, is the most popular web-based Diffusion Model interface optimized for workflow. It provides users with a UI platform for image generation, supporting plugins like LoRA, ControlNet, and IPAdaptor. Yet, its design for native single-GPU usage leaves it struggling with the demands of today's large DiTs, resulting in unacceptably high latency for users like Flux.1.
Using our commercial project TACO-DiT, a close-sourced ComfyUI variant built with xDiT, we've successfully implemented a multi-GPU parallel processing workflow within ComfyUI, effectively addressing Flux.1's performance challenges. Below is an example of using TACO-DiT to accelerate a Flux workflow with LoRA:
By using TACO-DiT, you could significantly reduce your ComfyUI workflow inference latency, and boosting the throughput with Multi-GPUs. Now it is compatible with multiple Plug-ins, including ControlNet and LoRAs.
More features and details can be found in our Intro Video:
The blog article is also available: Supercharge Your AIGC Experience: Leverage xDiT for Multiple GPU Parallel in ComfyUI Flux.1 Workflow.
ComfyUI plugin for xDiT is now available: xdit-comfyui-private
About diffusers version:
pyproject.toml, newer models may require later versions or even need to be installed from main branch.flash_attn is an optional library that can be installed with xDiT. More supported attention backends can be seen here.
pip install xfuser # Basic installation
pip install "xfuser[flash-attn]" # With flash attentionpip install -e .
# Or optionally, with flash attention
pip install -e ".[flash-attn]"Note that we use two self-maintained packages:
The flash_attn used for yunchang should be >= 2.6.0
We provide a docker image for developers to develop with xDiT. The docker image is thufeifeibear/xdit-dev. For running with AMD GPUs (MI300X or newer), a monthly image with validated support for select models is available as well: rocm/pytorch-xdit
The xDiT Model Runner provides a single entry point for running most supported diffusion models with proper benchmarking and profiling support. To use it, simply run:
xdit --model FLUX.1-dev \
--prompt "A cat running in a garden" \
--ulysses_degree 8The runner does not support all older models. For those we have the example scripts below. More information on how to run the model runner is available here.
We provide examples demonstrating how to run models with xDiT in the ./examples/ directory. You can easily modify the model type, model directory, and parallel options in the examples/run.sh within the script to run some already supported DiT models.
bash examples/run.shHybridizing multiple parallelism techniques together is essential for efficiently scaling. It's important that the product of all parallel degrees matches the number of devices. Note use_cfg_parallel means cfg_parallel=2. For instance, you can combine CFG, PipeFusion, and sequence parallelism with the command below to generate an image of a cute dog through hybrid parallelism. Here ulysses_degree * pipefusion_parallel_degree * cfg_degree(use_cfg_parallel) == number of devices == 8.
torchrun --nproc_per_node=8 \
examples/pixartalpha_example.py \
--model models/PixArt-XL-2-1024-MS \
--pipefusion_parallel_degree 2 \
--ulysses_degree 2 \
--num_inference_steps 20 \
--warmup_steps 0 \
--prompt "A cute dog" \
--use_cfg_parallel⚠️ Applying PipeFusion requires setting warmup_steps, also required in DistriFusion, typically set to a small number compared with num_inference_steps.
The warmup step impacts the efficiency of PipeFusion as it cannot be executed in parallel, thus degrading to a serial execution.
We observed that a warmup of 0 had no effect on the PixArt model.
Users can tune this value according to their specific tasks.
You can also launch an HTTP service to generate images with xDiT.
Launching a Text-to-Image Http Service
When initializing the runtime, xDiT checks which attention backends are installed and available and chooses the fastest one automatically.
This behaviour can be overriden via command line argument --attention-backend <backend cli name>.
Several different attention backends are supported:
| Backend name | CLI name |
|---|---|
| SDPA | sdpa |
| SDPA - Math | sdpa_math |
| SDPA - Memory Efficient | sdpa_efficient |
| SDPA - Flash | sdpa_flash |
| cuDNN | cudnn |
| FAv2 | flash |
| FAv3 | flash_3 |
| FAv3 FP8 | flash_3_fp8 |
| Transformer Engine FP8 | nvte_fp8 |
| FAv4 | flash_4 |
| FAv4 FP4 | flash_4_fp4 |
| SAGE | sage |
| AITER | aiter |
| AITER MHA v4 (dense) | aiter_bf16, aiter_bf16fp8, aiter_i8fp8, aiter_fp8, aiter_mxfp8, aiter_f8f6, aiter_mxfp6, aiter_f6f4, aiter_mxfp4, aiter_f4f4 |
| AITER MHA v4 Sparge | aiter_i8fp8_sparge, aiter_fp8_sparge, aiter_mxfp8_sparge, aiter_f8f6_sparge, aiter_mxfp6_sparge, aiter_f6f4_sparge, aiter_mxfp4_sparge, aiter_f4f4_sparge |
| AITER Sage | aiter_sage |
| AITER Sage V2 | aiter_sage_v2 |
| AITER Sparse Sage | aiter_sparse_sage |
| AITER Sparse Sage V2 | aiter_sparse_sage_v2 |
| AITER Sparge | aiter_sparge |
| AITER Sparge V2 | aiter_sparge_v2 |
| [Flex Block Sparge](https://github.com/xdit-project/xDi |
No open issues yet, or sync has not completed.