[optimization] help us know which kernels we should integrate in Diffusers

[sayakpaul]

This issue is for knowing which kernels we should integrate into the library through kernels.

Currently, we leverage kernels for different attention backends (FA2, FA3, and SAGE). However, other layers can be optimized as well (RMS Norm, for example), depending on the model size and input payload being used to benchmark that.

I did take a crack at this once, i.e., replacing the norm layers with their optimized counterparts, but didn't realize any noticeable gains. But maybe this is different now.

Resources / notes

• There's a bunch of kernels we maintain in https://huggingface.co/kernels-community, which could be repurposed in this case.
• We can also work together with the community to port impactful kernels and host them through https://huggingface.co/kernels-community and make sure they are kernels compatible.
• transformers gain benefits from this paradigm in terms of latency improvements (@MekkCyber can provide details). Some relevant PRs can be found here: https://github.com/huggingface/transformers/commits?author=MekkCyber (look for PRs with titles starting with "[kernels]")

[ParagEkbote]

Since the integration of the diffusers backend in sglang, perhaps the radix attention kernel could be added to kernels library which could allow users to have speedup in image generation without having to setup a sglang server? WDYT?

[sayakpaul]

perhaps the radix attention kernel could be added to kernels library

Sure, feel free to provide more context and how you envision that to be used in diffusers.

[ParagEkbote]

perhaps the radix attention kernel could be added to kernels library

Sure, feel free to provide more context and how you envision that to be used in diffusers.

The radix attention module could be shipped as a kernel to kernels community and used as follows:

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", torch_dtype=torch.bfloat16, device_map="cuda"
)
pipeline.transformer.set_attention_backend("_radix__hub")

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
pipeline(prompt).images[0]

and since the module is the main innovative portion of the sglang repo, the integration would be less wasteful since it has been tested several times, WDYT?

[ParagEkbote]

Another kernel which I found to be interesting for video diffusion models is the fast-video kernel, which is also used within sglang as a seperate attention backend.

[sayakpaul]

The radix attention module could be shipped as a kernel to kernels community and used as follows:

Does the corresponding module require building from some source? Same applies to fast-video-kernel. How should that be used, i.e., and where?

[ParagEkbote]

For radix attention, I don't think there is any source installation needed similar to Flash-Attention. The kernels would be used similarly like the sage-attention or fa kernels on kernels community, the main speedup can be observed in the UNet/ Transformer layers of the model.Also, sglang would not be a runtime requirement.

[ParagEkbote]

For fast-video-kernel, the installation is a bit more tricker. It uses custom CUDA kernels and kernels from thunderkittens, so installation from source would be required, the speed could be similar to the UNet/ Transformer layers.

For enabling faster operations on other layers like rmsnorm or layernorm as mentioned previously, TurboDiffusion has custom kernels in CUDA which could be experimented with.

WDYT?