ComfyUI-bleh项目中FlashAttention2集成方案解析

ComfyUI-bleh项目中FlashAttention2集成方案解析

ComfyUI-bleh ComfyUI nodes collection: better TAESD previews (including batch previews), improved HyperTile and Deep Shrink nodes 项目地址: https://gitcode.com/gh_mirrors/co/ComfyUI-bleh

背景介绍

ComfyUI-bleh是一个基于ComfyUI的图像生成扩展项目，该项目提供了多种优化后的注意力机制实现。在深度学习领域，注意力机制是Transformer架构的核心组件，其计算效率直接影响模型推理速度。

技术方案

项目作者blepping提供了一个基于FlashAttention2的实现方案，该方案包含两个主要节点：

1. BlehFlashAttentionSampler：用于局部注意力计算
2. BlehGlobalFlashAttention：用于全局注意力计算

性能对比

作者在多种模型上进行了性能测试，硬件环境为NVIDIA 4060Ti 16GB显卡，测试结果如下：

• Wan(small)模型：

  • SageAttention：6.35秒/迭代
  • FlashAttention2：8.26秒/迭代
  • PyTorch SDP：8.37秒/迭代

• Flux模型：

  • SageAttention：6.82秒/迭代
  • FlashAttention2：7.03秒/迭代
  • PyTorch SDP：7.06秒/迭代

从测试数据可以看出，FlashAttention2在当前硬件配置下性能略优于PyTorch原生实现，但仍不及SageAttention方案。

实现细节

该实现需要将自定义节点文件放置在项目的custom_nodes目录下。值得注意的是，该实现没有包含FlashAttention2缺失时的错误处理逻辑，因此用户需要确保环境中已正确安装FlashAttention2。

部署注意事项

1. 安装要求：FlashAttention2的安装过程较为复杂，需要约2小时编译时间和80GB内存空间
2. 硬件兼容性：不同硬件环境下性能表现可能有所差异
3. 预编译包：建议使用预编译的wheel包以节省安装时间

技术选型建议

对于追求最高性能的用户，在当前硬件环境下仍建议使用SageAttention实现。FlashAttention2方案可能更适合特定硬件配置或作为技术验证使用。

该实现为技术爱好者提供了一个研究不同注意力机制实现的参考方案，用户可根据自身需求和硬件条件选择合适的注意力实现方式。

ComfyUI-bleh ComfyUI nodes collection: better TAESD previews (including batch previews), improved HyperTile and Deep Shrink nodes 项目地址: https://gitcode.com/gh_mirrors/co/ComfyUI-bleh