vipshop/cache-dit

📚 What Cache‑DiT Is

Cache‑DiT is a PyTorch‑native inference engine that speeds up diffusion‑based transformer models (DiTs) used for image, video, audio and 3‑D generation. It sits on top of Hugging Face’s Diffusers library and adds four main performance tricks:

Feature	What it does	Why it matters
Hybrid cache	Stores intermediate key/value tensors (DBCache, TaylorSeer, SCM, …) so later diffusion steps can reuse them instead of recomputing.	Up to ~9× faster generation.
Parallelism	Context, tensor, 2‑D/3‑D and specialised parallelism for the text encoder, VAE and ControlNet.	Makes large models run on multiple GPUs or GPU+CPU setups.
Quantization	Supports 8‑bit and 4‑bit weight/activation (W8A8, W4A4).	Cuts memory and compute while keeping quality.
CPU off‑load	Bucket‑style layer‑wise off‑loading of parts of the model to CPU with near‑zero overhead.	Lets you fit bigger DiTs on limited GPU memory.

The engine is compatible with compilation tools (e.g., torch.compile), works on NVIDIA GPUs, AMD GPUs and Ascend NPUs, and can be plugged into existing ecosystems such as SGLang Diffusion, vLLM‑Omni, TensorRT‑LLM, ComfyUI, Jetson‑Containers and many others.

---

🚀 What It Can Do

• Run >40 families of DiT pipelines (120+ concrete variants) covering popular models like FLUX, StableDiffusion‑3, PixArt‑Sigma, CogVideoX, StableAudio, etc.
• Provide 9× speed‑up when using cache + context parallelism + torch compilation (as shown in the repo’s benchmark image).
• Offer a “model‑integration SKILL” for GitHub Copilot/Claude‑Code agents, making it easy to add new DiT pipelines with a single command.
• Support quantization and layer‑wise off‑load for any nn.Module, meaning new DiTs can benefit automatically.

---

🛠️ How to Get Started (Quick‑Start)

# Install the package from PyPI
pip install cache-dit

# Example: run a FLUX model with cache and parallelism
python -m cache_dit.run \
    --model flux \
    --cache on \
    --parallelism context \
    --quantize w8a8

The README points to detailed docs at cache‑dit.io for:

• Cache API – how to enable/disable different cache back‑ends.
• Context Parallelism – splitting the diffusion timesteps across GPUs.
• Tensor Parallelism – sharding the transformer weights.
• Quantization & Off‑load – flags for low‑precision inference.

---

🤝 Community & Ecosystem

Cache‑DiT is already integrated with a growing list of projects:

• Diffusers (official Hugging Face optimisation guide)
• ComfyUI (visual node‑based UI for diffusion)
• SGLang, vLLM‑Omni, TensorRT‑LLM (large‑scale serving frameworks)
• Jetson‑Containers (edge‑device deployment)
• Various open‑source pipelines like xDiT, MindIE‑SD, SD.Next, stable‑diffusion.cpp, etc.

These integrations are mostly pull‑requests that add a thin adapter layer, showing that Cache‑DiT can be dropped into existing codebases with minimal changes.

---

📄 License & Citation

• Licensed under the Apache‑2.0 license (as indicated by the badge).
• Citation provided in BibTeX format for academic use.

---

Bottom line: Cache‑DiT is a serious, production‑grade acceleration layer for diffusion transformer models. It targets developers who already use Diffusers and want to squeeze more speed, lower memory, or run on heterogeneous hardware without rewriting their models.