What kernels should we integrate in Diffusers?

[sayakpaul]

Now that we have an integration 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. 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 would be prioritized. A reference is here: Add kernelize to transformers transformers#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

[dxqb]

Quantization with custom kernels has huge potential for performance improvement.

Nunchaku has shown that to the extreme, but their way is probably impractical for a library like diffusers (they manually fused half the transformer block into 1 kernel).
But even less intrusive methods such as only fusing the activation quantization and quantized matmul of a linear layer into a kernel still has big potential.

[sayakpaul]

#12207 if you want to help a bit with nunchaku. We could at least try to inject the linears from nunchaku and run a quantization step to obtain a checkpoint and load it later. This is still a big win IMO. But we can discuss that on the mentioned PR.

But even less intrusive methods such as only fusing the activation quantization and quantized matmul of a linear layer into a kernel still has big potential.

When using quantization, one can do a horizontal fusion of the QKV projections to increase compute density. Most of the quantization backends supported in Diffusers are compatible with torch.compile. So, when you apply torch.compile to an NF4-quantized model (with QKV fused), for example, opportunities to do fusions are already kinda taken care of.

We have run a number of experiments to gather some numbers: https://github.com/sayakpaul/diffusers-torchao

This issue thread is more relevant with kernels that are existing already as this way, it's easier to quickly prototype something to evaluate if it's worth the integration.

[Akshat-Tripathi]

Hi, are the kernels just for improving performance, or reducing memory usage too? I've seen that fusing operations helps cut down on a lot of wasteful memory allocations that torch.compile doesn't seem to remove.

This might be a bit of a long shot, but could the the Linear -> Activation -> Linear pattern be kernelised? Since the intermediate results are immediately consumed, it could lead to a nice speedup

[sayakpaul]

Hi, are the kernels just for improving performance, or reducing memory usage too? I've seen that fusing operations helps cut down on a lot of wasteful memory allocations that torch.compile doesn't seem to remove.

Memory consumption is also desirable, for sure.

This might be a bit of a long shot, but could the the Linear -> Activation -> Linear pattern be kernelised? Since the intermediate results are immediately consumed, it could lead to a nice speedup

The decision would ultimately boil down to the fact if this is better than using torch.compile, IMO. We could take some of the adaptive layer norm and feed-forward layers and try to see if this works out. I think @MekkCyber has this in the plans.

[hlky]

apply_rotary_emb.zip iirc this gave 10% performance gain in Flux

[ParagEkbote]

I have also begun to integrate kernels-community/flash-attn in this PR: #12318

But there seems to be some issues in this implementation and if anyone has suggestions to get this complete this integration, feel free to reach out.

[bghira]

for what it's worth, qkv fusion doesn't help when memory bandwidth is the bottleneck, from what i've found. it's still a large amount of data being transferred. only things like MLA and quantised attn help with this. fusion helps when you have huge kernel launch overhead that dominates over the compute, but flux is quite large.