diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index fec307b..bb6619a 100644 --- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py +++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py @@ -56,7 +56,21 @@ def attention( scale: float | None = None, gqa: bool = False, ) -> Tensor: - with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + # cuDNN attention is NVIDIA-only, so hardcoding SDPBackend.CUDNN_ATTENTION + # raises "No available kernel" on non-NVIDIA backends (AMD ROCm, Intel XPU, + # Apple MPS). Pass a priority list instead: cuDNN is still preferred on + # NVIDIA, and the dispatcher falls back to flash/efficient/math elsewhere. + # (On ROCm gfx11xx the flash path needs + # TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1; masked attention uses math.) + with sdpa_kernel( + [ + SDPBackend.CUDNN_ATTENTION, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.MATH, + ], + set_priority=True, + ): x = F.scaled_dot_product_attention( q, k, v, attn_mask=mask, scale=scale, enable_gqa=gqa ) diff --git a/version.py b/version.py index c342ff1..7ce0cef 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.24" +VERSION = "0.10.25"