vipshop/cache-dit

A PyTorch-native inference engine with cache, parallelism, quantization and cpu offload for DiTs.

📚 Cache‑DiT 是什么

Cache‑DiT 是一个PyTorch 原生推理引擎，用于加速用于图像、视频、音频和 3D 生成的基于扩散的 Transformer 模型（DiT）。它构建在 Hugging Face 的 Diffusers 库之上，并添加了四项主要性能优化技术：

功能	功能说明	重要性
混合缓存	存储中间键/值张量（DBCache、TaylorSeer、SCM 等），以便后续扩散步骤可以重用它们，而不是重新计算。	生成速度最高提升约 9 倍。
并行化	针对文本编码器、VAE 和 ControlNet 的上下文、张量、2D/3D 及专用并行化。	使大型模型能够在多个 GPU 或 GPU+CPU 配置上运行。
量化	支持 8 位和 4 位权重/激活（W8A8、W4A4）。	在保持质量的同时减少内存和计算量。
CPU 卸载	将模型的部分层以桶状方式按层卸载到 CPU，开销几乎为零。	让更大的 DiT 模型在有限的 GPU 内存上运行。

该引擎兼容编译工具（例如 torch.compile），可在 NVIDIA GPU、AMD GPU 和 Ascend NPU 上运行，并且可以插入 SGLang Diffusion、vLLM‑Omni、TensorRT‑LLM、ComfyUI、Jetson‑Containers 等现有生态系统。

🚀 它能做什么

运行 40 多个 DiT 流水线系列（120 多个具体变体），涵盖 FLUX、StableDiffusion‑3、PixArt‑Sigma、CogVideoX、StableAudio 等流行模型。
使用缓存 + 上下文并行 + torch 编译时，可实现 9 倍加速（如仓库基准图像所示）。
为 GitHub Copilot/Claude-Code 代理提供**“模型集成 SKILL”**，使通过单个命令添加新的 DiT 流水线变得容易。
支持任何 nn.Module 的量化和逐层卸载，这意味着新的 DiT 模型可以自动受益。

🛠️ 如何开始（快速入门）

# 从 PyPI 安装包
pip install cache-dit

# 示例：使用缓存和并行运行 FLUX 模型
python -m cache_dit.run \
    --model flux \
    --cache on \
    --parallelism context \
    --quantize w8a8

README 指向 cache‑dit.io 上的详细文档，包括：

缓存 API – 如何启用/禁用不同的缓存后端。
上下文并行 – 将扩散时间步长分配到多个 GPU 上。
张量并行 – 对 Transformer 权重进行分片。
量化与卸载 – 低精度推理的标志。

🤝 社区与生态系统

Cache‑DiT 已与越来越多的项目集成：

Diffusers（官方 Hugging Face 优化指南）
ComfyUI（用于扩散的可视化节点 UI）
SGLang、vLLM‑Omni、TensorRT‑LLM（大规模服务框架）
Jetson‑Containers（边缘设备部署）
各种开源流水线，如 xDiT、MindIE‑SD、SD.Next、stable‑diffusion.cpp 等。

这些集成大多是添加一个薄适配器层的拉取请求，表明 Cache‑DiT 可以以最小的代码更改集成到现有代码库中。

📄 许可证与引用

根据 Apache‑2.0 许可证（如徽章所示）发布。
提供 BibTeX 格式的引用，以供学术使用。

总结： Cache‑DiT 是一个严肃的、生产级别的扩散 Transformer 模型加速层。它面向已经在使用 Diffusers 并希望在不重写模型的情况下提高速度、降低内存占用或在不同硬件上运行的开发者。