vipshop/cache-dit

🔥🔥🔥 4x speedup for Z-Image-ControlNet

Open

#611 opened on Dec 25, 2025

View on GitHub
 (0 comments) (2 reactions) (0 assignees)Python (72 forks)github user discovery
good first issue

Repository metrics

Stars
 (1,188 stars)
PR merge metrics
 (PR metrics pending)

Description

4x speedup for Z-Image-ControlNet

Cache-DiT provides inference acceleration for Z-Image and Z-Image-ControlNet via DBCache, Context Parallelism and Tensor Parallelism. It achieves nearly 4x speedup on 4 GPUs with negligible precision loss. Please visit their example for more details.

Z-Image-ControlNet Context Parallel: Ulysses 2 Context Parallel: Ulysses 4 + ControlNet Parallel
Base L20x1: 22s 15.7s 12.7s 🚀7.71s
+ Hybrid Cache + Torch Compile + Async Ulyess CP + FP8 All2All + CUDNN ATTN
🚀6.85s 6.45s 6.38s 🚀6.19s, 5.47s

Reproduce

pip3 install torch==2.9.1 transformers accelerate torchao bitsandbytes torchvision 
pip3 install opencv-python-headless einops imageio-ffmpeg ftfy 
pip3 install git+https://github.com/huggingface/diffusers.git # latest or >= 0.36.0
pip3 install git+https://github.com/vipshop/cache-dit.git # latest

git clone https://github.com/vipshop/cache-dit.git && cd examples

torchrun --nproc_per_node=4 generate.py zimage_controlnet --parallel ulysses \
         --parallel-controlnet --cache --scm fast --rdt 0.6 --compile \
         --compile-controlnet --ulysses-float8 --attn _sdpa_cudnn \
         --warmup 2 --repeat 4     

Contributor guide