huggingface / huggingface/diffusers
What kernels should we integrate in Diffusers?
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
Now that we have an [integration](https://github.com/huggingface/diffusers/pull/12236) with the `kernels` lib to use Flash Attention 3 (FA3), it'd be nice to gather community interest about which kernels we should try to incorporate in the library through the [`kernels` lib](https://github.com/huggingface/kernels/). FA3 delivers a significant speedup on Hopper GPUs.
I have done some work in the `kernelize` branch to see if replacing `GELU`, `SiLU`, and `RMSNorm` with their optimized kernels would have any speedups on Flux. So far, it hasn't had any. Benchmarking script: https://gist.github.com/sayakpaul/35236dd96e15d9f7d658a7ad11918411. One can compare the changes here: https://github.com/huggingface/diffusers/compare/kernelize?expand=1.
> [!NOTE]
> The changes in the `kernelize` branch are quite hacky as we're still evaluating things.
Please use this issue to let us know which kernels we should try to support in Diffusers. Some notes to keep in mind:
* Layers where the `forward()` method is easily replaceable with the `kernelize()` [mechanism](https://github.com/huggingface/kernels/blob/main/docs/source/layers.md#kernelizing-a-model) would be prioritized. A reference is here: https://github.com/huggingface/transformers/pull/38205.
* Even if a kernel isn't directly compatible with `kernels`, we can try to make it so, like we have for https://huggingface.co/kernels-community/flash-attn3.
* Not all kernels contribute non-trivial gains in terms of speedup. So, please bear that in mind when proposing a kernel.
Cc: @MekkCyber
Contributor guide
Research direction
Start with the linked kernelize branch, benchmarking script, and kernels library's layer-kernelizing documentation. Compare candidate kernels using the provided Flux benchmarking setup and review the existing Flash Attention 3 integration. The issue has no named implementation target or acceptance criterion; completion would require an agreed kernel and a concrete integration plan.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100