Consider using pre-built Flash Attention kernels via `kernels`
Hey,
I am Sayak from the Kernels team at Hugging Face. The GUI exposes a "Use Flash Attention" option for Lumina, Anima and HunyuanImage (kohya_gui/class_lumina.py, class_anima.py, class_hunyuan_image.py), with the hint "Requires flash-attn package". Since flash-attn isn't in requirements.txt, ticking that box today means a source build, which is a notoriously long one, plus the nvcc/CUDA-version matching that goes with it. We ship pre-built binaries (which provide bit-exact outputs as the upstream) and thereby, we make it easy to use.
Using FA3 on a supported machine is as easy as:
# make sure `kernels` is installed: `pip install -U kernels`
from kernels import get_kernel
kernel_module = get_kernel("kernels-community/flash-attn3")
flash_attn_func = kernel_module.flash_attn_func
flash_attn_func(...)I realise the actual import and dispatch live upstream in the sd-scripts submodule (library/attention.py, the attn_mode == "flash" branch), so the code change would need to happen there — I'm raising it here first since this repo owns the install story and the checkbox that points users at flash-attn. Happy to take it upstream as well if you'd prefer.
Importantly, this would slot in as another optional source rather than a replacement. The existing behaviour would stay exactly as it is: if flash_attn is importable it wins, otherwise kernels is tried, and if neither is available things fall through to the current xformers / torch (SDPA) paths just like they do today. Machines that aren't covered by the pre-built binaries, and users who already have a working source build, are unaffected — kernels stays a soft import behind the same try/except ImportError pattern already used there.
Let us know if you'd be interested in this and we'd be happy to provide a draft of how it would look in your repo.
Source: bmaltais/kohya_ss