From 1ba9f2ed87133f9fb3c7c5ec69e3d0f9d78fd949 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 6 Nov 2024 14:54:53 -0800
Subject: [PATCH] Update flex_attention.py

---
 unsloth/kernels/flex_attention.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 887ffca1b7..dfd48504d5 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -40,6 +40,14 @@ pass
 
 if not HAS_FLEX_ATTENTION:
 
+    # Below fails on compiled_autograd, so disable it
+    try:
+        old_compiled_autograd = torch._dynamo.config.compiled_autograd
+        torch._dynamo.config.compiled_autograd = False
+    except:
+        old_compiled_autograd = False
+    pass
+
     # Logit softcapping
     @torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options)
     def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len):
@@ -74,6 +82,13 @@ if not HAS_FLEX_ATTENTION:
         return A
     pass
 
+    # Return compiled_autograd back
+    try:
+        torch._dynamo.config.compiled_autograd = old_compiled_autograd
+    except:
+        pass
+    pass
+
     create_flex_attention_causal_mask = None
     create_flex_attention_sliding_window_mask = None
 else: