From 2e8d5c0e7f59b26f310ee30b0e5d8d79f23f769d Mon Sep 17 00:00:00 2001 From: Alex-Wengg <36247722+Alex-Wengg@users.noreply.github.com> Date: Sun, 21 Jun 2026 04:30:37 +0000 Subject: [PATCH] Fix flash-attn2 hub-kernel attribute names in attention dispatch The FLASH_HUB and FLASH_VARLEN_HUB configs in _HUB_KERNELS_REGISTRY point at flash_attn_interface._wrapped_flash_attn_{,varlen_}{forward,backward}, which do not exist in kernels-community/flash-attn2 (it exposes the non-_wrapped_ names). _resolve_kernel_attr therefore raises AttributeError when the flash-attn2 hub backend is used. Use the actual attribute names, matching the FLASH_3_HUB config. Fixes #14012 Co-authored-by: Claude --- src/diffusers/models/attention_dispatch.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index d9920a877112..e4c383c44fec 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -345,15 +345,15 @@ 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", + wrapped_forward_attr="flash_attn_interface._flash_attn_forward", + wrapped_backward_attr="flash_attn_interface._flash_attn_backward", version=1, ), 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", + wrapped_forward_attr="flash_attn_interface._flash_attn_varlen_forward", + wrapped_backward_attr="flash_attn_interface._flash_attn_varlen_backward", version=1, ), AttentionBackendName.SAGE_HUB: _HubKernelConfig(