From 679cdccbbc5a2d060449fa722ffee83482d53c9e Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Wed, 14 Oct 2026 13:48:48 +0800 Subject: [PATCH 1/2] Update FlashAttention Hub integration and guard CUDA-only tests Signed-off-by: Wang, Yi --- src/diffusers/models/attention_dispatch.py | 28 +++++++--------------- tests/models/test_attention_dispatch.py | 15 +++++++++--- 2 files changed, 21 insertions(+), 22 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index f45c9e0d1107..18e575d8ebbc 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -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", @@ -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, @@ -1219,7 +1215,7 @@ def _flash_attention_forward_op( window_size[1], softcap, alibi_slopes, - return_lse, + False, ) lse = lse.permute(0, 2, 1) @@ -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, @@ -1342,7 +1335,7 @@ def _flash_attention_hub_forward_op( window_size[1], softcap, alibi_slopes, - return_lse, + False, ) lse = lse.permute(0, 2, 1).contiguous() @@ -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 @@ -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:]) diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 2143707aac40..bfa632ae6632 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -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."), @@ -134,11 +137,17 @@ 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."), + ], ), ], ) From b4ec9bc7378abc8c724e8392950746bcd307a072 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Wed, 14 Oct 2026 15:45:57 +0800 Subject: [PATCH 2/2] fmt Signed-off-by: Wang, Yi --- tests/models/test_attention_dispatch.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index bfa632ae6632..478fb2853c0f 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -146,7 +146,9 @@ class TestContextParallelAttentionBackward: "_flash_3_varlen_hub", 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."), + pytest.mark.skipif( + torch_device != "cuda", reason="CUDA is required for _flash_3_varlen_hub backend." + ), ], ), ],