Upload transformer/attention.py with huggingface_hub

transformer/attention.py

@@ -36,6 +36,52 @@ try:
 except Exception as e:
     sageattn_func = None
 
+use_hub_kernels = os.getenv("DIFFUSERS_ENABLE_HUB_KERNELS", "false").upper() in [
+    "1",
+    "TRUE",
+]
+
+SAGEATTN_HUB_AVAILABLE = False
+try:
+    if use_hub_kernels and not is_kernels_available():
+        raise EnvironmentError(
+            (
+                "Attempting to use Hub Kernels for Flash Attention 3,"
+                "but the `kernels` library was not found in your environment. "
+                "Please install via `pip install kernels`"
+            )
+        )
+    if os.getenv("DISABLE_SAGEATTENTION", "0") != "0":
+        raise Exception("DISABLE_SAGEATTENTION is set")
+
+    from kernels import get_kernel
+
+    sageattn_hub = get_kernel("kernels-community/sage_attention")
+
+    @torch.library.custom_op(
+        "mylib::sageattn_hub", mutates_args={"q", "k", "v"}, device_types="cuda"
+    )
+    def sageattn_hub_func(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0,
+        is_causal: bool = False,
+    ) -> torch.Tensor:
+        return sageattn_hub(
+            q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal
+        )
+
+    @sageattn_func.register_fake
+    def _sageattn_fake(q, k, v, attn_mask=None, dropout_p=0, is_causal=False):
+        return torch.empty(*q.shape, device=q.device, dtype=q.dtype)
+
+    SAGEATTN_HUB_AVAILABLE = True
+
+except Exception as e:
+    sageattn_hub_func = None
+
 
 def _is_hopper_gpu():
     """Check if the current GPU is a Hopper architecture."""
@@ -44,41 +90,66 @@ def _is_hopper_gpu():
     device_name = torch.cuda.get_device_name(0).lower()
     return "h100" in device_name or "hopper" in device_name
 
+
 FLASH_ATTN_3_AVAILABLE = False
 try:
     import flash_attn_interface
+
     FLASH_ATTN_3_AVAILABLE = _is_hopper_gpu()
-except ModuleNotFoundError:
-    FLASH_ATTN_3_AVAILABLE = False
+except:
+    flash_attn_interface = None
 
 FLASH_ATTN_3_HUB_AVAILABLE = False
 try:
-    use_hub_kernels = os.getenv("DIFFUSERS_ENABLE_HUB_KERNELS", "false").upper() in ["1", "TRUE"]
     if use_hub_kernels and not is_kernels_available():
-        raise EnvironmentError((
-            "Attempting to use Hub Kernels for Flash Attention 3,"
-            "but the `kernels` library was not found in your environment. "
-            "Please install via `pip install kernels`"
-        ))
+        raise EnvironmentError(
+            (
+                "Attempting to use Hub Kernels for Flash Attention 3,"
+                "but the `kernels` library was not found in your environment. "
+                "Please install via `pip install kernels`"
+            )
+        )
 
     from kernels import get_kernel
+
     flash_attn_3_hub = get_kernel("kernels-community/flash-attn3", revision="fake-ops-return-probs")
 
     FLASH_ATTN_3_HUB_AVAILABLE = _is_hopper_gpu()
 
 except:
-    FLASH_ATTN_3_HUB_AVAILABLE = False
+    flash_attn_3_hub = None
 
 FLASH_ATTN_2_AVAILABLE = False
 try:
     import flash_attn
 
     FLASH_ATTN_2_AVAILABLE = True
-except ModuleNotFoundError:
-    FLASH_ATTN_2_AVAILABLE = False
+except:
+    flash_attn = None
+
+
+FLASH_ATTN_2_HUB_AVAILABLE = False
+try:
+    if use_hub_kernels and not is_kernels_available():
+        raise EnvironmentError(
+            (
+                "Attempting to use Hub Kernels for Flash Attention 3,"
+                "but the `kernels` library was not found in your environment. "
+                "Please install via `pip install kernels`"
+            )
+        )
+
+    from kernels import get_kernel
+
+    flash_attn_2_hub = get_kernel("kernels-community/flash-attn2")
+
+    FLASH_ATTN_2_HUB_AVAILABLE = True
+except:
+    flash_attn_2_hub = None
 
 __all__ = ["flash_attention", "attention"]
 
+
 def flash_attention(
     q,
     k,
@@ -107,12 +178,19 @@ def flash_attention(
     deterministic:  bool. If True, slightly slower and uses more memory.
     dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
     """
-    if not FLASH_ATTN_3_AVAILABLE or not FLASH_ATTN_3_HUB_AVAILABLE:
-        return flash_attn.flash_attn_func(
-            q,
-            k,
-            v,
-        )
+    if FLASH_ATTN_3_AVAILABLE and not FLASH_ATTN_3_HUB_AVAILABLE:
+        if FLASH_ATTN_2_HUB_AVAILABLE:
+            return flash_attn_2_hub.flash_attn_func(
+                q,
+                k,
+                v,
+            )
+        else:
+            return flash_attn.flash_attn_func(
+                q,
+                k,
+                v,
+            )
 
     elif FLASH_ATTN_3_HUB_AVAILABLE:
         return flash_attn_3_hub.flash_attn_func(
@@ -182,7 +260,7 @@ def flash_attention(
             deterministic=deterministic,
         ).unflatten(0, (b, lq))
     else:
-        assert FLASH_ATTN_3_AVAILABLE
+        assert FLASH_ATTN_2_AVAILABLE
         x = flash_attn.flash_attn_varlen_func(
             q=q,
             k=k,
@@ -222,9 +300,7 @@ def attention(
     fa_version=None,
     # og_dtype=torch.bfloat16,
 ):
-
-    if SAGEATTN_AVAILABLE:
-        # print("Using sageattention")
+    if SAGEATTN_AVAILABLE or SAGEATTN_HUB_AVAILABLE:
         attn_mask = None
 
         og_dtype = q.dtype
@@ -232,14 +308,19 @@ def attention(
         k = k.transpose(1, 2).to(dtype)
         v = v.transpose(1, 2).to(dtype)
 
-        out = sageattn_func(
-            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
-        )
+        if SAGEATTN_HUB_AVAILABLE:
+            out = sageattn_hub_func(
+                q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+            )
+        else:
+            out = sageattn_func(
+                q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+            )
 
         out = out.transpose(1, 2).contiguous().to(og_dtype)
         return out
 
-    elif FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+    elif FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE or FLASH_ATTN_3_HUB_AVAILABLE:
         return flash_attention(
             q=q,
             k=k,