Repository: huggingface/diffusers
File: src/diffusers/models/attention_dispatch.py (Python)

Code before:

        return decorator

    @classmethod
    def get_active_backend(cls):
        return cls._active_backend, cls._backends[cls._active_backend]

    @classmethod
    def list_backends(cls):
        return list(cls._backends.keys())


@lru_cache(maxsize=None)
def _load_fa3_hub():
    from ..utils.kernels_utils import _get_fa3_from_hub

    fa3_hub = _get_fa3_from_hub()  # won't re-download if already present
    if fa3_hub is None:
        raise RuntimeError(
            "Failed to load FlashAttention-3 kernels from the Hub. Please ensure the wheel is available for your platform."
        )
    return fa3_hub


def flash_attn_3_hub_func(*args, **kwargs):
    return _load_fa3_hub().flash_attn_func(*args, **kwargs)


@contextlib.contextmanager
def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE):
    """
    Context manager to set the active attention backend.
    """
    if backend not in _AttentionBackendRegistry._backends:
        raise ValueError(f"Backend {backend} is not registered.")

    backend = AttentionBackendName(backend)
    _check_attention_backend_requirements(backend)

    old_backend = _AttentionBackendRegistry._active_backend
    _AttentionBackendRegistry._active_backend = backend

    try:
        yield
    finally:
        _AttentionBackendRegistry._active_backend = old_backend


def dispatch_attention_fn(
    query: torch.Tensor,
    key: torch.Tensor,

Review comment (suggestion):
Rather than loading the kernel every time we invoke flash attention, it would be better to import the function at the top of the file, similar to the other FA backends. 

https://github.com/huggingface/diffusers/blob/e58711e73cba70c1c02aaa67f80945a1458901b8/src/diffusers/models/attention_dispatch.py#L63-L68

Code after:
    _checks_enabled = DIFFUSERS_ATTN_CHECKS

    @classmethod
    def register(cls, backend: AttentionBackendName, constraints: Optional[List[Callable]] = None):
        logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}")

        def decorator(func):
            cls._backends[backend] = func
            cls._constraints[backend] = constraints or []
            cls._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys())
            return func

        return decorator

    @classmethod
    def get_active_backend(cls):
        return cls._active_backend, cls._backends[cls._active_backend]

    @classmethod
    def list_backends(cls):
        return list(cls._backends.keys())


@contextlib.contextmanager
def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE):
    """
    Context manager to set the active attention backend.
    """
    if backend not in _AttentionBackendRegistry._backends:
        raise ValueError(f"Backend {backend} is not registered.")

    backend = AttentionBackendName(backend)
    _check_attention_backend_requirements(backend)

    old_backend = _AttentionBackendRegistry._active_backend
    _AttentionBackendRegistry._active_backend = backend

    try:
        yield
    finally:
        _AttentionBackendRegistry._active_backend = old_backend


def dispatch_attention_fn(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_mask: Optional[torch.Tensor] = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: Optional[float] = None,