refactor according to Dhruv's ideas.

Co-authored-by: Dhruv Nair <dhruv@huggingface.co>

Author: sayakpaul

src/diffusers/models/attention_dispatch.py

@@ -16,8 +16,9 @@
 import functools
 import inspect
 import math
+from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
 
@@ -40,7 +41,7 @@
     is_xformers_available,
     is_xformers_version,
 )
-from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS, DIFFUSERS_ENABLE_HUB_KERNELS
+from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS
 
 
 if TYPE_CHECKING:
@@ -78,9 +79,6 @@
     flash_attn_3_func = None
     flash_attn_3_varlen_func = None
 
-_BACKEND_HANDLES: Dict["AttentionBackendName", Callable] = {}
-_PREPARED_BACKENDS: Set["AttentionBackendName"] = set()
-
 if _CAN_USE_SAGE_ATTN:
     from sageattention import (
         sageattn,
@@ -222,9 +220,7 @@ def decorator(func):
 
     @classmethod
     def get_active_backend(cls):
-        backend = cls._active_backend
-        _ensure_attention_backend_ready(backend)
-        return backend, cls._backends[backend]
+        return cls._active_backend, cls._backends[cls._active_backend]
 
     @classmethod
     def list_backends(cls):
@@ -242,6 +238,25 @@ def _is_context_parallel_enabled(
         return supports_context_parallel and is_degree_greater_than_1
 
 
+@dataclass
+class _HubKernelConfig:
+    """Configuration for downloading and using a hub-based attention kernel."""
+
+    repo_id: str
+    function_attr: str
+    revision: Optional[str] = None
+    kernel_fn: Optional[Callable] = None
+
+
+# Registry for hub-based attention kernels
+_HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = {
+    # TODO: temporary revision for now. Remove when merged upstream into `main`.
+    AttentionBackendName._FLASH_3_HUB: _HubKernelConfig(
+        repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs"
+    )
+}
+
+
 @contextlib.contextmanager
 def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE):
     """
@@ -251,7 +266,7 @@ def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBacke
         raise ValueError(f"Backend {backend} is not registered.")
 
     backend = AttentionBackendName(backend)
-    _ensure_attention_backend_ready(backend)
+    _check_attention_backend_requirements(backend)
 
     old_backend = _AttentionBackendRegistry._active_backend
     _AttentionBackendRegistry._active_backend = backend
@@ -398,13 +413,9 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
 
     # TODO: add support Hub variant of FA3 varlen later
     elif backend in [AttentionBackendName._FLASH_3_HUB]:
-        if not DIFFUSERS_ENABLE_HUB_KERNELS:
-            raise RuntimeError(
-                f"Flash Attention 3 Hub backend '{backend.value}' is not usable because the `DIFFUSERS_ENABLE_HUB_KERNELS` env var isn't set. Please set it like `export DIFFUSERS_ENABLE_HUB_KERNELS=yes`."
-            )
         if not is_kernels_available():
             raise RuntimeError(
-                f"Flash Attention 3 Hub backend '{backend.value}' is not usable because the `kernels` package isn't available. Please install it with `pip install kernels`."
+                f"Backend '{backend.value}' is not usable because the `kernels` package isn't available. Please install it with `pip install kernels`."
             )
 
     elif backend in [
@@ -445,39 +456,6 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
             )
 
 
-def _ensure_flash_attn_3_func_hub_loaded():
-    cached = _BACKEND_HANDLES.get(AttentionBackendName._FLASH_3_HUB)
-    if cached is not None:
-        return cached
-
-    from ..utils.kernels_utils import _get_fa3_from_hub
-
-    flash_attn_interface_hub = _get_fa3_from_hub()
-    func = flash_attn_interface_hub.flash_attn_func
-    _BACKEND_HANDLES[AttentionBackendName._FLASH_3_HUB] = func
-
-    return func
-
-
-_BACKEND_PREPARERS: Dict[AttentionBackendName, Callable[[], None]] = {
-    AttentionBackendName._FLASH_3_HUB: _ensure_flash_attn_3_func_hub_loaded,
-}
-
-
-def _prepare_attention_backend(backend: AttentionBackendName) -> None:
-    preparer = _BACKEND_PREPARERS.get(backend)
-    if preparer is not None:
-        preparer()
-
-
-def _ensure_attention_backend_ready(backend: AttentionBackendName) -> None:
-    if backend in _PREPARED_BACKENDS:
-        return
-    _check_attention_backend_requirements(backend)
-    _prepare_attention_backend(backend)
-    _PREPARED_BACKENDS.add(backend)
-
-
 @functools.lru_cache(maxsize=128)
 def _prepare_for_flash_attn_or_sage_varlen_without_mask(
     batch_size: int,
@@ -581,6 +559,29 @@ def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
     return q_idx >= kv_idx
 
 
+# ===== Helpers for downloading kernels =====
+def _maybe_download_kernel_for_backend(backend: AttentionBackendName) -> None:
+    if backend not in _HUB_KERNELS_REGISTRY:
+        return
+    config = _HUB_KERNELS_REGISTRY[backend]
+
+    if config._kernel_fn is not None:
+        return
+
+    try:
+        from kernels import get_kernel
+
+        kernel_module = get_kernel(config.repo_id, revision=config.revision)
+        kernel_func = getattr(kernel_module, config.function_attr)
+
+        # Cache the downloaded kernel function in the config object
+        config._kernel_fn = kernel_func
+
+    except Exception as e:
+        logger.error(f"An error occurred while fetching kernel '{config.repo_id}' from the Hub: {e}")
+        raise
+
+
 # ===== torch op registrations =====
 # Registrations are required for fullgraph tracing compatibility
 # TODO: this is only required because the beta release FA3 does not have it. There is a PR adding
@@ -1348,9 +1349,7 @@ def _flash_attention_3_hub(
     return_attn_probs: bool = False,
     _parallel_config: Optional["ParallelConfig"] = None,
 ) -> torch.Tensor:
-    func = _BACKEND_HANDLES.get(AttentionBackendName._FLASH_3_HUB)
-    if func is None:
-        func = _ensure_flash_attn_3_func_hub_loaded()
+    func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB]._kernel_fn
     out = func(
         q=query,
         k=key,

src/diffusers/models/modeling_utils.py

@@ -595,7 +595,11 @@ def set_attention_backend(self, backend: str) -> None:
                 attention as backend.
         """
         from .attention import AttentionModuleMixin
-        from .attention_dispatch import AttentionBackendName, _ensure_attention_backend_ready
+        from .attention_dispatch import (
+            AttentionBackendName,
+            _check_attention_backend_requirements,
+            _maybe_download_kernel_for_backend,
+        )
 
         # TODO: the following will not be required when everything is refactored to AttentionModuleMixin
         from .attention_processor import Attention, MochiAttention
@@ -606,8 +610,10 @@ def set_attention_backend(self, backend: str) -> None:
         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)
-        _ensure_attention_backend_ready(backend)
+        _check_attention_backend_requirements(backend)
+        _maybe_download_kernel_for_backend(backend)
 
         attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
         for module in self.modules():

src/diffusers/utils/constants.py

@@ -46,7 +46,6 @@
 DEFAULT_HF_PARALLEL_LOADING_WORKERS = 8
 HF_ENABLE_PARALLEL_LOADING = os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES
 DIFFUSERS_DISABLE_REMOTE_CODE = os.getenv("DIFFUSERS_DISABLE_REMOTE_CODE", "false").upper() in ENV_VARS_TRUE_VALUES
-DIFFUSERS_ENABLE_HUB_KERNELS = os.environ.get("DIFFUSERS_ENABLE_HUB_KERNELS", "").upper() in ENV_VARS_TRUE_VALUES
 
 # Below should be `True` if the current version of `peft` and `transformers` are compatible with
 # PEFT backend. Will automatically fall back to PEFT backend if the correct versions of the libraries are

tests/others/test_attention_backends.py

@@ -7,7 +7,6 @@
 
 ```bash
 export RUN_ATTENTION_BACKEND_TESTS=yes
-export DIFFUSERS_ENABLE_HUB_KERNELS=yes
 
 pytest tests/others/test_attention_backends.py
 ```