change to Hub.

Author: sayakpaul

src/diffusers/models/attention_dispatch.py

@@ -26,6 +26,7 @@
     is_flash_attn_3_available,
     is_flash_attn_available,
     is_flash_attn_version,
+    is_kernels_available,
     is_sageattention_available,
     is_sageattention_version,
     is_torch_npu_available,
@@ -54,7 +55,6 @@
 _CAN_USE_XLA_ATTN = is_torch_xla_available() and is_torch_xla_version(">=", _REQUIRED_XLA_VERSION)
 _CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _REQUIRED_XFORMERS_VERSION)
 
-_DEFAULT_HUB_ID_FA3 = "kernels-community/vllm-flash-attn3"
 
 if _CAN_USE_FLASH_ATTN:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -67,16 +67,22 @@
     from flash_attn_interface import flash_attn_func as flash_attn_3_func
     from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
 else:
-    try:
-        from kernels import get_kernel
+    flash_attn_3_func = None
+    flash_attn_3_varlen_func = None
 
-        vllm_flash_attn3 = get_kernel(_DEFAULT_HUB_ID_FA3)
-        flash_attn_3_func = vllm_flash_attn3.flash_attn_func
-        flash_attn_3_varlen_func = vllm_flash_attn3.flash_attn_varlen_func
-        logger.debug(f"Using Flash Attention 3 from {_DEFAULT_HUB_ID_FA3} using the `kernels` lib.")
-    except ImportError:
-        flash_attn_3_func = None
-        flash_attn_3_varlen_func = None
+if is_kernels_available():
+    from ..utils.kernels_utils import _get_fa3_from_hub
+
+    flash_attn_interface_hub = _get_fa3_from_hub()
+    if flash_attn_interface_hub is not None:
+        flash_attn_3_hub_func = flash_attn_interface_hub.flash_attn_func
+        flash_attn_3_varlen_hub_func = flash_attn_interface_hub.flash_attn_varlen_func
+    else:
+        flash_attn_3_hub_func = None
+        flash_attn_3_varlen_hub_func = None
+else:
+    flash_attn_3_hub_func = None
+    flash_attn_3_varlen_hub_func = None
 
 
 if _CAN_USE_SAGE_ATTN:
@@ -162,6 +168,8 @@ class AttentionBackendName(str, Enum):
     FLASH_VARLEN = "flash_varlen"
     _FLASH_3 = "_flash_3"
     _FLASH_VARLEN_3 = "_flash_varlen_3"
+    _FLASH_3_HUB = "_flash_3_hub"
+    _FLASH_VARLEN_3_HUB = "_flash_varlen_3_hub"  # not supported yet.
 
     # PyTorch native
     FLEX = "flex"
@@ -355,11 +363,20 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
             )
 
     elif backend in [AttentionBackendName._FLASH_3, AttentionBackendName._FLASH_VARLEN_3]:
-        if not _CAN_USE_FLASH_ATTN_3 and (flash_attn_3_func is None and flash_attn_3_varlen_func is None):
+        if not _CAN_USE_FLASH_ATTN_3:
             raise RuntimeError(
                 f"Flash Attention 3 backend '{backend.value}' is not usable because of missing package or the version is too old. Please build FA3 beta release from source."
             )
 
+    # TODO: add support Hub variant of FA3 varlen later
+    elif backend in [AttentionBackendName._FLASH_3_HUB]:
+        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`."
+            )
+    elif backend in [AttentionBackendName._FLASH_VARLEN_3_HUB]:
+        raise NotImplementedError
+
     elif backend in [
         AttentionBackendName.SAGE,
         AttentionBackendName.SAGE_VARLEN,
@@ -523,6 +540,22 @@ def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> Tuple[torc
     return torch.empty_like(query), query.new_empty(lse_shape)
 
 
+@_custom_op("flash_attn_3_hub_func", mutates_args=(), device_types="cuda")
+def _wrapped_flash_attn_3_hub(
+    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    out, lse = flash_attn_3_hub_func(query, key, value)
+    lse = lse.permute(0, 2, 1)
+    return out, lse
+
+
+@_register_fake("flash_attn_3_hub_func")
+def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    batch_size, seq_len, num_heads, head_dim = query.shape
+    lse_shape = (batch_size, seq_len, num_heads)
+    return torch.empty_like(query), query.new_empty(lse_shape)
+
+
 # ===== Attention backends =====
 
 
@@ -645,34 +678,59 @@ def _flash_attention_3(
     deterministic: bool = False,
     return_attn_probs: bool = False,
 ) -> torch.Tensor:
-    sig = inspect.signature(flash_attn_3_func)
-    accepted = set(sig.parameters)
-    params = {
-        "q": query,
-        "k": key,
-        "v": value,
-        "softmax_scale": scale,
-        "causal": is_causal,
-        "qv": None,
-        "q_descale": None,
-        "k_descale": None,
-        "v_descale": None,
-        "window_size": window_size,
-        "attention_chunk": 0,
-        "softcap": softcap,
-        "num_splits": 1,
-        "pack_gqa": None,
-        "deterministic": deterministic,
-        "sm_margin": 0,
-    }
-    kwargs = {}
-    for name, value in params.items():
-        if name not in accepted:
-            logger.debug(f"{name} is not accepted by the `flash_attn_3_func` method, so it will be discarded.")
-        else:
-            kwargs[name] = value
+    out, lse, *_ = flash_attn_3_func(
+        q=query,
+        k=key,
+        v=value,
+        softmax_scale=scale,
+        causal=is_causal,
+        qv=None,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        window_size=window_size,
+        attention_chunk=0,
+        softcap=softcap,
+        num_splits=1,
+        pack_gqa=None,
+        deterministic=deterministic,
+        sm_margin=0,
+    )
+    return (out, lse) if return_attn_probs else out
+
 
-    out, lse, *_ = flash_attn_3_func(**kwargs)
+@_AttentionBackendRegistry.register(
+    AttentionBackendName._FLASH_3_HUB,
+    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+)
+def _flash_attention_3_hub(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size: Tuple[int, int] = (-1, -1),
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    out, lse, *_ = flash_attn_3_hub_func(
+        q=query,
+        k=key,
+        v=value,
+        softmax_scale=scale,
+        causal=is_causal,
+        qv=None,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        window_size=window_size,
+        softcap=softcap,
+        num_splits=1,
+        pack_gqa=None,
+        deterministic=deterministic,
+        sm_margin=0,
+    )
     return (out, lse) if return_attn_probs else out
 
 

src/diffusers/utils/kernels_utils.py

@@ -0,0 +1,17 @@
+from .import_utils import is_kernels_available
+
+
+_DEFAULT_HUB_ID_FA3 = "kernels-community/vllm-flash-attn3"
+
+
+def _get_fa3_from_hub():
+    if not is_kernels_available():
+        return None
+    else:
+        from kernels import get_kernel
+
+        try:
+            vllm_flash_attn3 = get_kernel(_DEFAULT_HUB_ID_FA3)
+            return vllm_flash_attn3
+        except Exception as e:
+            raise e