diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 504f4afe3af3..3a7e810697b4 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -159,13 +159,19 @@ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProce return self.processor def set_attention_backend(self, backend: str): - from .attention_dispatch import AttentionBackendName + from .attention_dispatch import ( + AttentionBackendName, + _check_attention_backend_requirements, + _maybe_download_kernel_for_backend, + ) available_backends = {x.value for x in AttentionBackendName.__members__.values()} if backend not in available_backends: raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends)) backend = AttentionBackendName(backend.lower()) + _check_attention_backend_requirements(backend) + _maybe_download_kernel_for_backend(backend) self.processor._attention_backend = backend def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 478fb2853c0f..f4a320652f02 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -15,6 +15,7 @@ import contextlib import os +import types import pytest import torch @@ -22,9 +23,11 @@ import torch.multiprocessing as mp import torch.nn.functional as F +from diffusers.models import attention_dispatch from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig +from diffusers.models.attention_dispatch import AttentionBackendName, dispatch_attention_fn from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx -from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.transformers.transformer_flux import FluxAttention from ..testing_utils import ( is_attention, @@ -231,3 +234,38 @@ def make_qkv(seq_len): ) finally: torch.compiler.reset() + + +@is_attention +class TestAttentionModuleSetAttentionBackend: + @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") + def test_hub_kernel_is_loaded(self, monkeypatch): + import kernels + + def fake_sageattn(*args, **kwargs): + raise NotImplementedError + + fake_kernel = types.SimpleNamespace(__name__="fake_sage_attention", sageattn=fake_sageattn) + monkeypatch.setattr(kernels, "get_kernel", lambda *args, **kwargs: fake_kernel) + monkeypatch.setattr(attention_dispatch, "DIFFUSERS_TRUST_REMOTE_KERNELS", True) + config = attention_dispatch._HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB] + monkeypatch.setattr(config, "kernel_fn", None) + + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend == AttentionBackendName.SAGE_HUB + assert config.kernel_fn is fake_sageattn + + def test_backend_requirements_are_checked(self, monkeypatch): + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + + monkeypatch.setattr(attention_dispatch, "_CAN_USE_FLASH_ATTN", False) + with pytest.raises(RuntimeError, match="Flash Attention backend 'flash' is not usable"): + attn.set_attention_backend("flash") + + monkeypatch.setattr(attention_dispatch, "is_kernels_available", lambda: False) + with pytest.raises(RuntimeError, match="Backend 'sage_hub' is not usable because the `kernels` package"): + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend is None