[core] support sage attention + FA2 through kernels

[sayakpaul]

What does this PR do?

Code to test (SAGE):

fromdiffusersimportDiffusionPipelineimporttorchrepo_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:

FA2:

importtorchfromdiffusersimportFluxPipelinemodel_id="black-forest-labs/FLUX.1-dev"pipe=FluxPipeline.from_pretrained(
model_id, torch_dtype=torch.bfloat16
).to("cuda")
pipe.transformer.set_attention_backend("flash_hub")
pipe.transformer.compile(fullgraph=True)
prompt="A cat holding a sign that says 'hello world'"withtorch._dynamo.config.patch(error_on_recompile=True):
image=pipe(
prompt, num_inference_steps=28, guidance_scale=4.0, generator=torch.manual_seed(0)
).images[0]
image.save("output.png")

---

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

fromdiffusersimportDiffusionPipelineimporttorchrepo_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.

[woct0rdho]

FYI, I've ported SageAttention to Python stable ABI (ABI3) and libtorch stable ABI, which should simplify building for HF Kernels:
woct0rdho/SageAttention@main...abi3_stable

There are also some refactors in my main branch to simplify building. If someone can maintain the build system, then I no longer need to maintain my repo :)

[sayakpaul]

This PR is ready to be reviewed now. As discussed with @MekkCyber over DMs, we're disabling torch.compile support now as the compile branch leads to garbage outputs.

In order for us to support it with torch.compile, some kind of lightweight dispatcher might be needed. d344134 added support for that but I have removed it for now for the above-mentioned purpos. Those changes are still safe in sage-kernels-dispatch branch.

I think we should be good with the PR.

Cc: @MekkCyber@DN6

[sayakpaul]

@DN6 it should be up for another review. I have updated the test suite and ensured that they pass successfully as well. PTAL.

[DN6]

Good to merge. But we need to remove the parallel config check and use supports_context_parallel=False

[DN6]

Use:

@_AttentionBackendRegistry.register(AttentionBackendName.FLASH_HUB,constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],supports_context_parallel=False,)

diffusers/src/diffusers/models/attention_dispatch.py

Line 1307 in 544ba67

supports_context_parallel=True,

It will raise an error when trying to enable parallelism with this backend. This check isn't needed.