KJNodes/experimentalgenerated
Patch Flash Attention KJ(PatchFlashAttentionKJ)
Experimental node for patching attention to use flash attention, without the silent SDPA fallback the ComfyUI default does. Patches the attention of the model passing through this node; to disable, bypass or disconnect this node. Requires the flash_attn library to be installed.
Patch Flash Attention KJ
model
MODEL
allow_compile
KJNodes
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.