Description
Experimental node for patching the attention mechanism of a diffusion model to use Flash Attention (via the flash_attn library). Unlike ComfyUI's default SDPA fallback, this node enforces flash attention directly. To disable the patch, bypass or disconnect this node. Requires the flash_attn Python package to be installed (run pip install flash-attn).
Inputs
| Name | Type | Required | Default | Description |
|---|---|---|---|---|
model | MODEL | Yes | – | The diffusion model to patch. The attention layers will be replaced with flash attention operations. |
allow_compile | BOOLEAN | No | False | If True, allows the flash attention kernels to be compiled (e.g., using PyTorch's torch.compile) for potential performance gains. This may increase startup time and memory usage. Only enable if you have a compatible environment. |
Outputs
| Name | Type | Description |
|---|---|---|
MODEL | MODEL | The patched model with flash attention enabled. |
Usage Notes
- This node is experimental and may not work with all model architectures or CUDA configurations.
- The
flash_attnlibrary must be installed in your environment. Installation instructions:pip install flash-attn(may require a CUDA-compatible torch build). - If you encounter errors, try bypassing the node to revert to the default attention mechanism.
- The
allow_compileoption is off by default; enable it only if you understand the implications (longer first inference, potential instability).
Comments
Sign in with GitHub to join the discussion.