[core] support flash attention through kernels

[sayakpaul]

What does this PR do?

Follow-up of #12236.

Testing code:

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")

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.

[sayakpaul]

@DN6 a gentle ping on this one.

[DN6]

This is a very thin wrapper. I would just call _get_from_hub("fa3") directly in attention_dispatch.

[DN6]

Suggested change

	def_get_from_hub(key: str):
	def_get_kernel_from_hub(key: str):

[sayakpaul]

Closing for #12439.