Skip to content

Update FlashAttention Hub integration and guard CUDA-only tests #15024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 3 commits into from
Oct 14, 2026
Merged

Update FlashAttention Hub integration and guard CUDA-only tests #15024

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
679cdcc
Update FlashAttention Hub integration and guard CUDA-only tests
sywangyi Oct 14, 2026
669c0b4
Merge branch 'main' into flashattn_fix
sywangyi Oct 14, 2026
b4ec9bc
fmt
sywangyi Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 9 additions & 19 deletions src/diffusers/models/attention_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,16 +342,16 @@ class _HubKernelConfig:
AttentionBackendName.FLASH_HUB: _HubKernelConfig(
repo_id="kernels-community/flash-attn2",
function_attr="flash_attn_func",
wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_forward",
wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_backward",
version=1,
wrapped_forward_attr="flash_attn_interface._flash_attn_forward",
wrapped_backward_attr="flash_attn_interface._flash_attn_backward",
version=3,
),
AttentionBackendName.FLASH_VARLEN_HUB: _HubKernelConfig(
repo_id="kernels-community/flash-attn2",
function_attr="flash_attn_varlen_func",
wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_varlen_forward",
wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_varlen_backward",
version=1,
wrapped_forward_attr="flash_attn_interface._flash_attn_varlen_forward",
wrapped_backward_attr="flash_attn_interface._flash_attn_varlen_backward",
version=3,
),
AttentionBackendName.SAGE_HUB: _HubKernelConfig(
repo_id="SageAttention/sage-attention",
Expand Down Expand Up @@ -1203,10 +1203,6 @@ def _flash_attention_forward_op(
if scale is None:
scale = query.shape[-1] ** (-0.5)

# flash-attn only returns LSE if dropout_p > 0. So, we need to workaround.
if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1):
dropout_p = dropout_p if dropout_p > 0 else 1e-30

with torch.set_grad_enabled(grad_enabled):
out, lse, S_dmask, rng_state = _wrapped_flash_attn_forward(
query,
Expand All @@ -1219,7 +1215,7 @@ def _flash_attention_forward_op(
window_size[1],
softcap,
alibi_slopes,
return_lse,
False,
)
lse = lse.permute(0, 2, 1)

Expand Down Expand Up @@ -1327,9 +1323,6 @@ def _flash_attention_hub_forward_op(
deterministic = False
grad_enabled = any(x.requires_grad for x in (query, key, value))

if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1):
dropout_p = dropout_p if dropout_p > 0 else 1e-30

with torch.set_grad_enabled(grad_enabled):
out, lse, S_dmask, rng_state = wrapped_forward_fn(
query,
Expand All @@ -1342,7 +1335,7 @@ def _flash_attention_hub_forward_op(
window_size[1],
softcap,
alibi_slopes,
return_lse,
False,

@sayakpaul sayakpaul Oct 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

So this means return_softmax is always False?

@sywangyi sywangyi Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

see https://huggingface.co/kernels/kernels-community/flash-attn2/blob/v3/build/torch-stable-abi210-cu132-x86_64-linux/flash_attn_interface.py#L142, return_softmax will impact the third output(dmask) which is not used in the code.

@sayakpaul sayakpaul Oct 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

My question is what parameter does that False value correspond to? Is it not return_softmax?

@sywangyi sywangyi Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

it's return_softmax in v3, and default is False

)
lse = lse.permute(0, 2, 1).contiguous()

Expand Down Expand Up @@ -1452,9 +1445,6 @@ def _flash_varlen_attention_hub_forward_op(
deterministic = False
grad_enabled = any(x.requires_grad for x in (query, key, value))

if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1):
dropout_p = dropout_p if dropout_p > 0 else 1e-30

batch_size, seq_len_q, num_heads, _ = query.shape
_, seq_len_kv, _, _ = key.shape

Expand Down Expand Up @@ -1493,7 +1483,7 @@ def _flash_varlen_attention_hub_forward_op(
window_size[1],
softcap,
alibi_slopes,
return_lse,
False,
)

out = out_packed.view(batch_size, seq_len_q, *out_packed.shape[1:])
Expand Down
17 changes: 14 additions & 3 deletions tests/models/test_attention_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,10 @@ class TestContextParallelAttentionBackward:
"attention_backend",
[
"_native_flash",
"_native_cudnn",
pytest.param(
"_native_cudnn",
marks=pytest.mark.skipif(torch_device != "cuda", reason="CUDA is required for _native_cudnn backend."),
),
pytest.param(
"flash_hub",
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
Expand All @@ -134,11 +137,19 @@ class TestContextParallelAttentionBackward:
),
pytest.param(
"_flash_3_hub",
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
marks=[
pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
pytest.mark.skipif(torch_device != "cuda", reason="CUDA is required for _flash_3_hub backend."),
],
),
pytest.param(
"_flash_3_varlen_hub",
marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
marks=[
pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
pytest.mark.skipif(
torch_device != "cuda", reason="CUDA is required for _flash_3_varlen_hub backend."
),
],
),
],
)
Expand Down
Loading