[core] support flash attention through kernels

[sayakpaul]

What does this PR do?

Follow-up of #12236.

Testing code:

import torch
from diffusers import FluxPipeline

model_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'"

with torch._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")

Tip

Works with torch.compile fullgraph compatibility.

I have tested the code on H100 and A100, and it works.

[sayakpaul]

Flash Attention is stable. So, we don't have to mark it private like FA3.

[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 integration 🔥 ! I just left some nits

[MekkCyber]

Why are we fetching both kernels here ?

[sayakpaul]

Because of the way APIs for attention backends are designed and also to support torch.compile with fullgraph traceability (when possible).

We will let it grow a bit and upon feedback, we can revisit how to better deal with this.

[MekkCyber]

is this related to the kernel or it just needs more time to be integrated ?

[sayakpaul]

We don't have models that use varlen.

[bghira]

@sayakpaul qwen image uses varlen. also, native fused qkv+mlp attn requires varlen function.