Upgrade sdpa_kernel

Author: palonso

src/nets/common_former.py

@@ -2,9 +2,7 @@
 
 import torch
 import torch.nn as nn
-import gin.torch
-
-from .net import Net
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 
 class MHAPyTorchScaledDotProduct(nn.Module):
@@ -38,9 +36,8 @@ def forward(self, x):
         queries, keys, values = qkv
 
         use_dropout = 0.0 if not self.training else self.dropout
-        with torch.backends.cuda.sdp_kernel(
-            enable_flash=True, enable_math=False, enable_mem_efficient=False
-        ):
+
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
             context_vec = nn.functional.scaled_dot_product_attention(
                 queries,
                 keys,

src/nets/conformer.py

@@ -4,6 +4,7 @@
 import gin
 import torch
 import torch.nn as nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 from .common_former import DeepNorm
 from .rope import RotaryEmbedding
 
@@ -74,9 +75,7 @@ def forward(self, x):
         use_dropout = 0.0 if not self.training else self.dropout
 
         if self.use_flash_attention:
-            with torch.backends.cuda.sdp_kernel(
-                enable_flash=True, enable_math=False, enable_mem_efficient=False
-            ):
+            with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                 context_vec = nn.functional.scaled_dot_product_attention(
                     queries,
                     keys,