[core] download hub kernel in AttentionModuleMixin.set_attention_backend

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:

tests/models/test_attention_dispatch.py

@@ -15,16 +15,19 @@
 
 import contextlib
 import os
+import types
 
 import pytest
 import torch
 import torch.distributed as dist
 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