使用Torch编译缓存提升AI推理速度

Torch编译缓存，提升推理速度

很多模型，特别是FLUX系列的模型，会应用各种torch.compile技术/技巧来提高推理速度。

对编译函数的第一次调用会跟踪并编译代码，这会增加开销。随后的调用则运行优化后的代码，速度会显著加快。

提示 在我们对black-forest-labs/flux-kontext-dev模型进行推理速度测试时，编译版本比未编译版本的运行速度快了30%以上。

性能提升

通过在模型容器生命周期之间缓存编译产物，我们看到了冷启动时间的显著改善：

• black-forest-labs/flux-kontext-dev: ~120秒 → ~60秒（提速50%）
• prunaai/flux-schnell: ~150秒 → ~70秒（提速53%）
• prunaai/flux.1-dev-lora: ~400秒 → ~150秒（提速62%）

该缓存还改善了所有使用torch.compile的模型从容器的启动到首次成功预测所需的时间。

工作原理

该缓存系统的工作原理与许多CI/CD缓存系统类似：

• 当模型容器启动时，它会寻找缓存的编译产物。
• 如果找到，Torch会复用它们，而不是从头开始重新编译。
• 当容器正常关闭时，如果需要，它们会更新缓存。
• 缓存文件以模型版本为键进行存储，并放置在靠近GPU节点的位置。

了解更多

要了解更多关于如何使用torch.compile的信息，请查阅相关的技术文档和官方PyTorch torch.compile教程。FINISHED