Upload x_transformer_1_23_2.py

x_transformer_1_23_2.py

@@ -26,10 +26,16 @@
 from functools import partial
 from typing import Optional, Tuple
 
+import os
+os.environ['USE_FLASH_ATTENTION'] = '1'
+
 import torch
 from torch import nn, einsum, Tensor
 import torch.nn.functional as F
+
+# Flash attention
 from torch.nn.attention import SDPBackend, sdpa_kernel
+torch.backends.cuda.enable_flash_sdp(True)
 
 from collections import namedtuple
 from functools import wraps
@@ -259,11 +265,14 @@ class Attend(nn.Module):
         
         # Legacy code...
         # with torch.backends.cuda.sdp_kernel(enable_math=True, enable_mem_efficient=True):
+        # with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
 
-        # New SDP kernel code...
-        # with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
-        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
+        # PyTorch 2.3-2.4 SDPA backend code...
+        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION]):
 
+        # New PyTorch 2.5 SDPA backend code:
+        # with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            
             out = F.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask = mask,
@@ -508,7 +517,7 @@ class AutoregressiveWrapper(Module):
         # whether to add router z-loss
         self.add_attn_z_loss = add_attn_z_loss
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def generate(
         self,
@@ -2462,4 +2471,4 @@ class XTransformer(nn.Module):
             enc, mask = dropout_seq(enc, mask, self.cross_attn_tokens_dropout)
 
         out = self.decoder(tgt, context = enc, context_mask = mask)
-        return out
+        return out