[core] support sage attention through kernels

[sayakpaul]

What does this PR do?

Code to test:

from diffusers import DiffusionPipeline 
import torch 

repo_id = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")
pipe.transformer.set_attention_backend("sage_hub")

image = pipe(
    prompt="a dog sitting by the sea, waiting for its companion to come",
    guidance_scale=3.5,
    num_inference_steps=30,
    max_sequence_length=512,
    generator=torch.manual_seed(0)
).images[0]
image.save("sage_flux.png")

Result:

---

Notes

1. It would be nice to get torch.compile support when using sage attention like we have for flash and flash 3. Currently, this fails.

Code to test

from diffusers import DiffusionPipeline 
import torch 

repo_id = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")
pipe.transformer.set_attention_backend("sage_hub")
pipe.transformer.compile_repeated_blocks(fullgraph=True)

with (
    torch._inductor.utils.fresh_inductor_cache(),
    torch._dynamo.config.patch(error_on_recompile=True),
):
    image = pipe(
        prompt="a dog sitting by the sea, waiting for its companion to come",
        guidance_scale=3.5,
        num_inference_steps=30,
        max_sequence_length=512,
        generator=torch.manual_seed(0)
    ).images[0]
image.save("sage_flux.png")

Error: https://pastebin.com/3HS6HNzR

2. We have other sageattn variants (see here), which would be cool to expose from the Hub kernel.

Cc: @MekkCyber

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[MekkCyber]

Very cool ! I will try to look into the torch compile compatibility, but for the other variants, they are the same as sageattn, what i mean is sageattn is just a wrapper that dispatches to the correct kernel depending on the hardware used : https://github.com/thu-ml/SageAttention/blob/main/sageattention/core.py#L140

[sayakpaul]

they are the same as sageattn, what i mean is sageattn is just a wrapper that dispatches to the correct kernel depending on the hardware used :

So, you mean we shouldn't have to have different dispatched functions like this?

diffusers/src/diffusers/models/attention_dispatch.py

Line 194 in ce90f9b

_SAGE_QK_INT8_PV_FP8_CUDA = "_sage_qk_int8_pv_fp8_cuda"

[MekkCyber]

Yes I think we don't need that because it depends on the hardware. For example if a user chooses : _sage_qk_int8_pv_fp8_cuda on A100 (8.0) it will fail, because this function is only supported and compiled for 8.9 gpus

[sayakpaul]

I don't see their usage, hence removed.