From 2a251ec5948ea44d7332ab8b053b83276ae237ca Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jul 2024 18:18:09 -0700
Subject: [PATCH] Update flex_attention.py

---
 unsloth/kernels/flex_attention.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 1eb2486998..a992a02382 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -25,18 +25,23 @@ torch_compile_options = {
 }
 
 # Flex Attention supported from torch 2.5 onwards only
-import torch.nn.attention
-if hasattr(torch.nn.attention, "flex_attention"):
-    import torch.nn.attention.flex_attention
-    from torch.nn.attention.flex_attention import flex_attention
-    from torch.nn.attention.flex_attention import create_block_mask
-    FLEX_ATTENTION_PADDING = getattr(
-        torch.nn.attention.flex_attention,
-        "_DEFAULT_SPARSE_BLOCK_SIZE",
-        1,
-    )
-    flex_attention = torch.compile(flex_attention, dynamic = False)
-    HAS_FLEX_ATTENTION = True
+import torch.nn
+if hasattr(torch.nn, "attention"):
+    import torch.nn.attention
+    if hasattr(torch.nn.attention, "flex_attention"):
+        import torch.nn.attention.flex_attention
+        from torch.nn.attention.flex_attention import flex_attention
+        from torch.nn.attention.flex_attention import create_block_mask
+        FLEX_ATTENTION_PADDING = getattr(
+            torch.nn.attention.flex_attention,
+            "_DEFAULT_SPARSE_BLOCK_SIZE",
+            1,
+        )
+        flex_attention = torch.compile(flex_attention, dynamic = False)
+        HAS_FLEX_ATTENTION = True
+    else:
+        HAS_FLEX_ATTENTION = False
+    pass
 else:
     HAS_FLEX_ATTENTION = False
 pass