vipshop/cache-dit

解决的问题

Cache-DiT 是一个专为加速 Diffusion Transformers (DiTs) 执行而设计的 PyTorch 原生推理引擎。它解决了用于图像、视频和音频生成的规模化 DiT 模型的高计算成本和内存需求问题，实现了更快的推理速度和更高效的硬件利用率。

工作原理

Cache-DiT 构建在 Hugging Face Diffusers 库之上，采用了多种优化技术来降低延迟和内存占用：

• 混合缓存加速：实现了多种缓存策略（如 DBCache、TaylorSeer、SCM 和 DMD），以避免重复计算。
• 并行化：支持广泛的并行策略，包括上下文并行 (Context Parallelism)、张量并行 (Tensor Parallelism) 以及混合 2D/3D 并行，并对 Text Encoders、VAEs 和 ControlNets 提供专用支持。
• 内存管理：提供量化 (W8A8, W4A4) 和分桶式逐层 CPU 卸载 (bucket-style layerwise CPU offloading)，以最大限度地减少 GPU 内存占用。
• 硬件兼容性：原生运行于 NVIDIA GPU、Ascend NPU 和 AMD GPU，并与 PyTorch 编译兼容。

适用对象

该引擎面向需要为生产环境或高性能计算环境优化其推理流水线的 DiT 基础生成式 AI 模型（如 FLUX、CogVideoX 和 HunyuanVideo）的开发人员和研究人员。

亮点

• 广泛的模型支持：支持来自 Diffusers 的 40 多个 DiT 流水线家族和 120 多个变体。
• 巨大的加速效果：通过结合缓存、上下文并行和编译，能够实现高达 9 倍的加速。
• 深度集成：与 SGLang Diffusion、vLLM-Omni、TensorRT-LLM 和 ComfyUI 等流行框架完全集成。
• 智能体集成：包含模型集成 SKILL，帮助用户使用 GitHub Copilot 或 Claude Code 等编程智能体来集成新的 DiT 流水线。