Cache Diffusion技术解析：Model Optimizer免训练加速扩散模型的秘密

Cache Diffusion技术解析Model Optimizer免训练加速扩散模型的秘密【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址: https://gitcode.com/GitHub_Trending/te/Model-OptimizerCache Diffusion是 NVIDIA 开源项目 Model Optimizer 提供的一项免训练training-free扩散模型加速技术它通过复用扩散去噪过程中相邻步骤的缓存输出替代重复计算让 Stable Diffusion XL、SD3 等模型在不损失画质、不重新训练的前提下获得可观的推理加速。本文将从原理、实现到实测效果带你快速掌握这项技术的秘密。为什么扩散模型推理慢扩散模型如 Stable Diffusion生成一张图需要执行几十次去噪步骤diffusion steps。每一步都要让整张 UNet 或 DiT 网络完整前向一次因此步骤多常见 20~50 步每步都是全模型计算相邻步骤高度相似研究发现相邻两步中网络各层block的输出变化很小大量计算其实是在重复劳动。朴素的做法是直接减少步数但这会明显损害画质。Cache Diffusion 的思路是步数不减但在特定步骤跳过部分 block 的计算直接复用缓存的输出。核心原理跳过计算复用缓存下图展示了 Model Optimizer 中 SDXL 默认的 Cache Diffusion 计算图绿色实块表示该步骤真实计算的 block虚线块表示被跳过、直接复用缓存结果的 block。可以看到在 T1 步中大部分 block 被跳过从而获得显著加速。其实现非常轻量核心逻辑见 module.py用CachedModule包裹原始 block内部维护当前步数cur_step与缓存结果cached_results每次前向时先判断当前步是否需要计算if_cache需要则执行原始 block 并更新缓存否则直接返回上一次的缓存输出提供enable/disable开关方便随时对比基线与加速模式。模块替换与模型适配逻辑在 cachify.py 中目前已支持 UNetSDXL、视频模型 UNet3D与 DiTPixArt、SD3等主流结构。三步开启 Cache Diffusion 加速使用方式极其简单无需任何训练或校准数据只需三行 API完整可运行演示见 example.ipynbfrom cache_diffusion import cachify from cache_diffusion.utils import SDXL_DEFAULT_CONFIG cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) # 1. 注入缓存模块 cachify.enable(pipe) # 2. 开启缓存 with cachify.infer(pipe): # 3. 正常推理即可 img pipe(promptprompt, num_inference_steps20).images[0]其中 SDXL 的默认配置定义在 utils.py对up_blocks.2之外的所有 block在奇数步复用缓存结果。自定义缓存策略两个关键参数如果你希望针对自己的模型调节速度-画质平衡只需调整两个参数参数作用说明wildcard_or_filter_func决定哪些模块参与缓存支持通配符字符串如*up_blocks*或过滤函数select_cache_step_func决定哪些步骤复用缓存返回 True 的步跳过计算直接返回缓存例如把step % 2 ! 0改成step % 3 ! 0即可进一步提高缓存频率获得更大加速对 PixArt、SD3、视频模型utils.py中也内置了对应的默认配置PIXART_DEFAULT_CONFIG、SD3_DEFAULT_CONFIG、SVD_DEFAULT_CONFIG可直接使用。实测效果提速 1.63 倍画质几乎无损官方演示在单卡 RTX 6000 Ada 上用 PyTorch 完成。下图对比了三组 SDXL 生成结果20 步基线、20 步 Cache Diffusion1.63x 加速、11 步直接减步1.62x 加速。可以看到在相近耗时下Cache Diffusion 的画质明显优于暴力减步甚至非常接近 20 步基线。与量化、TensorRT 组合使用Cache Diffusion 是正交的优化手段可以和 Model Optimizer 的其他能力叠加与 PTQ 量化组合先用 INT8/FP8 量化压缩精度再叠加 Cache Diffusion 跳过冗余计算双重提速TensorRT 加速将 UNet 拆分成多个 TensorRT engine缓存逻辑按 PyTorch 模块的方式在 engine 之间组合执行部署时同样生效模型支持Stable Diffusion 3、SDXL 均已列入支持矩阵详见 examples/diffusers/README.md。 小提示扩散管线的随机采样会导致每次校准/生成存在轻微差异建议多跑几次选择最佳结果固定随机种子可提升结果可复现性。小结一张清单看懂 Cache Diffusion✅免训练无需数据集、无需梯度更新纯推理期优化✅易集成3 行 API 即可注入任意 diffusers 管线✅可定制模块筛选 步数策略双参数灵活调节速度与画质✅可叠加与量化、TensorRT 部署无缝组合✅效果实在SDXL 实测 1.63 倍加速画质接近全量 20 步基线。 相关源码与资料索引核心注入逻辑cachify.py缓存模块实现module.py各模型默认配置utils.py完整可运行示例example.ipynbDiffusers 优化总览examples/diffusers/README.md【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址: https://gitcode.com/GitHub_Trending/te/Model-Optimizer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考