Skip to content

[core] support flash attention through kernels#12387

Closed
sayakpaul wants to merge 7 commits intomainfrom
fa-hub
Closed

[core] support flash attention through kernels#12387
sayakpaul wants to merge 7 commits intomainfrom
fa-hub

Conversation

@sayakpaul
Copy link
Member

@sayakpaul sayakpaul commented Sep 25, 2025

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 sayakpaul requested a review from DN6 September 25, 2025 08:01
@sayakpaul sayakpaul added the performance Anything related to performance improvements, profiling and benchmarking label Sep 25, 2025
@sayakpaul sayakpaul mentioned this pull request Sep 25, 2025
Closed
6 tasks
sayakpaul
sayakpaul commented Sep 25, 2025
View reviewed changes
src/diffusers/models/attention_dispatch.py
# `flash-attn`
FLASH = "flash"
FLASH_VARLEN = "flash_varlen"
FLASH_HUB = "flash_hub"
Copy link
Member Author

@sayakpaul sayakpaul Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Sep 25, 2025

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
MekkCyber reviewed Sep 25, 2025
View reviewed changes
Copy link

@MekkCyber MekkCyber left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very cool integration 🔥 ! I just left some nits

src/diffusers/models/attention_dispatch.py
Comment on lines +85 to +88
fa3_interface_hub = _get_fa3_from_hub()
flash_attn_3_func_hub = fa3_interface_hub.flash_attn_func
fa_interface_hub = _get_fa_from_hub()
flash_attn_func_hub = fa_interface_hub.flash_attn_func
Copy link

@MekkCyber MekkCyber Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we fetching both kernels here ?

Copy link
Member Author

@sayakpaul sayakpaul Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

src/diffusers/models/attention_dispatch.py
FLASH = "flash"
FLASH_VARLEN = "flash_varlen"
FLASH_HUB = "flash_hub"
# FLASH_VARLEN_HUB = "flash_varlen_hub" # not supported yet.
Copy link

@MekkCyber MekkCyber Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

@sayakpaul sayakpaul Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have models that use varlen.

Copy link
Contributor

@bghira bghira Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@sayakpaul sayakpaul mentioned this pull request Sep 25, 2025
Merged
@sayakpaul
Copy link
Member Author

sayakpaul commented Oct 6, 2025

@DN6 a gentle ping on this one.

This was referenced Oct 13, 2025
Closed
Merged
DN6
DN6 reviewed Oct 22, 2025
View reviewed changes
src/diffusers/utils/kernels_utils.py
raise


def _get_fa3_from_hub():
Copy link
Collaborator

@DN6 DN6 Oct 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

src/diffusers/utils/kernels_utils.py


def _get_fa3_from_hub():
def _get_from_hub(key: str):
Copy link
Collaborator

@DN6 DN6 Oct 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def _get_from_hub(key: str):
def _get_kernel_from_hub(key: str):

This was referenced Oct 22, 2025
Merged
Merged
@sayakpaul
Copy link
Member Author

sayakpaul commented Nov 20, 2025

Closing for #12439.

@sayakpaul sayakpaul closed this Nov 20, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@DN6 DN6 DN6 left review comments

+2 more reviewers

@bghira bghira bghira left review comments

@MekkCyber MekkCyber MekkCyber left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

performance Anything related to performance improvements, profiling and benchmarking

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@sayakpaul @HuggingFaceDocBuilderDev @DN6 @bghira @MekkCyber