add flash attention support using triton

Author: lshqqytiger

modules/rocm_triton_windows.py

@@ -1,7 +1,9 @@
 import sys
-from functools import wraps
 import torch
-from modules import shared, devices
+from modules import devices
+
+
+is_available = False
 
 
 if sys.platform == "win32":
@@ -94,34 +96,7 @@ def triton_runtime_driver_active_utils_get_device_properties(device):
                 return props
             triton.runtime.driver.active.utils.get_device_properties = triton_runtime_driver_active_utils_get_device_properties
 
-            if 'Flash attention' in shared.opts.sdp_options:
-                from modules.flash_attn_triton_amd import interface_fa
-                sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
-                @wraps(sdpa_pre_flash_atten)
-                def sdpa_flash_atten(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
-                    if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
-                        if scale is None:
-                            scale = query.shape[-1] ** (-0.5)
-                        head_size_og = query.size(3)
-                        if head_size_og % 8 != 0:
-                            query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
-                            key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
-                            value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
-                        query = query.transpose(1, 2)
-                        out_padded = torch.zeros_like(query)
-                        interface_fa.fwd(
-                            query,
-                            key.transpose(1, 2),
-                            value.transpose(1, 2),
-                            out_padded,
-                            dropout_p,
-                            scale,
-                            is_causal,
-                        )
-                        return out_padded[..., :head_size_og].transpose(1, 2)
-                    else:
-                        return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
-                torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
-                print('Torch attention: type="triton flash attention"')
+            global is_available
+            is_available = True
         except Exception:
             pass

modules/sd_hijack_optimizations.py

@@ -2,14 +2,15 @@
 import math
 import psutil
 import platform
+from functools import wraps
 
 import torch
 from torch import einsum
 
 from ldm.util import default
 from einops import rearrange
 
-from modules import shared, errors, devices, sub_quadratic_attention
+from modules import shared, errors, devices, sub_quadratic_attention, rocm_triton_windows
 from modules.hypernetworks import hypernetwork
 
 import ldm.modules.attention
@@ -34,7 +35,7 @@ def title(self):
 
         return f"{self.name} - {self.label}"
 
-    def is_available(self):
+    def is_available(self) -> bool:
         return True
 
     def apply(self):
@@ -143,8 +144,62 @@ def apply(self):
         sgm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
 
 
+class SdOptimizationTritonFlashAttention(SdOptimization):
+    name = "Flash attention"
+    cmd_opt = "flash_attn"
+    priority = 100
+
+    def __init__(self):
+        super().__init__()
+        self.sdpa_pre_flash_atten = None
+
+    def is_available(self):
+        return hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(torch.nn.functional.scaled_dot_product_attention) and devices.has_zluda() and rocm_triton_windows.is_available
+
+    def apply(self):
+        if self.sdpa_pre_flash_atten is None:
+            from modules.flash_attn_triton_amd import interface_fa
+            self.sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
+            @wraps(self.sdpa_pre_flash_atten)
+            def sdpa_flash_atten(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
+                if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
+                    if scale is None:
+                        scale = query.shape[-1] ** (-0.5)
+                    head_size_og = query.size(3)
+                    if head_size_og % 8 != 0:
+                        query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
+                        key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
+                        value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
+                    query = query.transpose(1, 2)
+                    out_padded = torch.zeros_like(query)
+                    interface_fa.fwd(
+                        query,
+                        key.transpose(1, 2),
+                        value.transpose(1, 2),
+                        out_padded,
+                        dropout_p,
+                        scale,
+                        is_causal,
+                    )
+                    return out_padded[..., :head_size_og].transpose(1, 2)
+                else:
+                    return self.sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
+            torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
+
+        ldm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
+        ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
+        sgm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
+        sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
+
+    def undo(self):
+        super().undo()
+        torch.nn.functional.scaled_dot_product_attention = self.sdpa_pre_flash_atten
+        self.sdpa_pre_flash_atten = None
+
+
 def list_optimizers(res):
     res.extend([
+        SdOptimizationTritonFlashAttention(),
         SdOptimizationXformers(),
         SdOptimizationSdpNoMem(),
         SdOptimizationSdp(),