Skip to content

vipshop/cache-dit

Repository files navigation

🤗🎉A PyTorch-native Inference Engine with Hybrid
Cache Acceleration and Massive Parallelism for DiTs
Featured|HelloGitHub

🤗Why Cache-DiT❓❓Cache-DiT is built on top of the 🤗Diffusers library and now supports nearly ALL DiTs from Diffusers. It provides hybrid cache acceleration (DBCache, TaylorSeer, SCM, etc.) and comprehensive parallelism optimizations, including Context Parallelism, Tensor Parallelism, hybrid 2D or 3D parallelism, and dedicated extra parallelism support for Text Encoder, VAE, and ControlNet.

Cache-DiT is compatible with compilation, CPU Offloading, and quantization, fully integrates with SGLang Diffusion, vLLM-Omni, ComfyUI, and runs natively on NVIDIA GPUs, Ascend NPUs and AMD GPUs. Cache-DiT is fast, easy to use, and flexible for various DiTs (online docs at 📘readthedocs.io).

⚡️9x speedup by Cache-DiT with Cache, Context Parallelism and Compilation

🚀Quick Start: Cache, Parallelism and Quantization

First, you can install the cache-dit from PyPI or install from source:

uv pip install -U cache-dit # or, uv pip install git+https://github.com/vipshop/cache-dit.git

Then, try to accelerate your DiTs with just ♥️one line♥️ of code ~

>>> import cache_dit
>>> from diffusers import DiffusionPipeline
>>> pipe = DiffusionPipeline.from_pretrained(...).to("cuda")
>>> cache_dit.enable_cache(pipe) # Cache Acceleration with One-line code.
>>> from cache_dit import DBCacheConfig, ParallelismConfig
>>> cache_dit.enable_cache( # Or, Hybrid Cache Acceleration + Parallelism.
...   pipe, cache_config=DBCacheConfig(), # w/ default
...   parallelism_config=ParallelismConfig(ulysses_size=2))
>>> from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig
>>> cache_dit.enable_cache( # Or, Hybrid Cache + Parallelism + Quantization.
...   pipe, cache_config=DBCacheConfig(), # w/ default
...   parallelism_config=ParallelismConfig(ulysses_size=2),
...   quantize_config=QuantizeConfig(quant_type="float8_per_row"))
>>> output = pipe(...) # Then, just call the pipe as normal.

For more advanced features, please refer to our online documentation at 📘Documentation.

🚀Quick Start: SVDQuant (W4A4) PTQ workflow

First, build Cache-DiT from source with SVDQuant support (Experimental):

git clone https://github.com/vipshop/cache-dit.git && cd cache-dit
CACHE_DIT_BUILD_SVDQUANT=1 uv pip install -e ".[quantization]" --no-build-isolation

Then, try to quantize your model with just ♥️a few lines♥️ of code ~

>>> from cache_dit import QuantizeConfig
>>> pipe = DiffusionPipeline.from_pretrained(...).to("cuda")
>>> # 0. Define the calibration function for PTQ.
>>> def calibrate_fn(**_: object) -> None:
...     with torch.inference_mode():
...         for prompt in calibration_prompts:
...             _ = pipe(prompt=prompt, ...)
>>> # 1. Build the QuantizeConfig for SVDQuant PTQ.
>>> quant_config = QuantizeConfig(
...     quant_type="svdq_int4_r32", # _r{rank}, e.g., r16, r32, r64, r128, etc.
...     calibrate_fn=calibrate_fn,
...     serialize_to=..., 
... )
>>> # 2. Apply quantization with `cache_dit.quantize(...)` API.
>>> pipe.transformer = cache_dit.quantize(pipe.transformer, quant_config) 
>>> output = pipe(...) # 3. Use the quantized model for inference.
>>> # 4. Or, reload the quantized model from disk for inference.
>>> pipe.transformer = cache_dit.load(pipe.transformer, ...)

For more advanced features, please refer to our online documentation at 📘Low-bits Quantization.

🌐Community Integration

©️Acknowledgements

Special thanks to vipshop's Computer Vision AI Team for supporting document, testing and deployment of this project. We learned the design and reused code from the following projects: Diffusers, SGLang, vLLM, vLLM-Omni, ParaAttention, xDiT and TaylorSeer.

©️Citations

@misc{cache-dit@2025,
  title={Cache-DiT: A PyTorch-native Inference Engine with Hybrid Cache Acceleration and Massive Parallelism for DiTs.},
  url={https://github.com/vipshop/cache-dit.git},
  note={Open-source software available at https://github.com/vipshop/cache-dit.git},
  author={DefTruth, vipshop.com, etc.},
  year={2025}
}

About

A PyTorch-native inference engine with hybrid cache acceleration and massive parallelism for DiTs.

Topics

Resources

License

Stars

Watchers

Forks

Packages

 
 
 

Contributors