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 )