From cb7157c4bf68e826ec040a7afa7037e6bd39acb1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 26 Mar 2025 04:11:27 -0700
Subject: [PATCH] check SDPA for Mistral 3, Pixtral

---
 unsloth/models/_utils.py |  8 +++++---
 unsloth/models/loader.py |  4 ++--
 unsloth/models/vision.py | 25 +++++++------------------
 3 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 0044c7e761..223e0f51fd 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1176,9 +1176,10 @@ def unsloth_compile_transformers(
             "so turning off some optimizations!"
         )
         return
-    if disable: return
-
     model_types = list(dict().fromkeys(model_types).keys())
+    if disable: return model_types, False
+
+    supports_sdpa = [True]
     for model_type in model_types:
         _unsloth_compile_transformers(
             model_type,
@@ -1206,12 +1207,13 @@ def unsloth_compile_transformers(
             import_from_cache      = import_from_cache,
             disable                = disable,
             return_logits          = return_logits,
+            supports_sdpa          = supports_sdpa,
         )
     pass
     # Redo patches which override compiler
     for temporary_patch in TEMPORARY_PATCHES:
         temporary_patch()
-    return model_types
+    return model_types, supports_sdpa[0]
 pass
 
 # We need an empty logits flag to warn people logits will not be returned anymore unless asked ie
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index c2bf51c791..cac5acd838 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -642,7 +642,6 @@ class FastModel(FastBaseModel):
             trust_remote_code = trust_remote_code,
         )
         model_types = ["siglip"] + model_types
-        print("model_types", model_types)
 
         # Set forced float32 env flag
         os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
@@ -664,7 +663,7 @@ class FastModel(FastBaseModel):
 
         with redirector:
             patch_loss_functions(torch_compile = False)
-            model_types = unsloth_compile_transformers(
+            model_types, supports_sdpa = unsloth_compile_transformers(
                 dtype                   = dtype,
                 model_name              = model_name,
                 model_types             = model_types,
@@ -727,6 +726,7 @@ class FastModel(FastBaseModel):
             tokenizer_name    = tokenizer_name,
             auto_model        = auto_model,
             use_gradient_checkpointing = use_gradient_checkpointing,
+            supports_sdpa     = supports_sdpa,
             *args, **kwargs,
         )
 
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 6244a6146d..4e9e5c5a4e 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -66,11 +66,6 @@ __all__ = [
     "FastBaseModel",
 ]
 
-global FORCE_EAGER_ATTENTION
-FORCE_EAGER_ATTENTION = [
-    "pixtral",    # Pixtral SDPA not implemented
-]
-
 global NUM_LOGITS_TO_KEEP
 NUM_LOGITS_TO_KEEP = dict()
 global PROMPT_LOOPKUP
@@ -240,6 +235,7 @@ class FastBaseModel:
         tokenizer_name    = None,
         auto_model        = AutoModelForVision2Seq,
         use_gradient_checkpointing = "unsloth",
+        supports_sdpa     = True,
         **kwargs,
     ):
         if model_types is None:
@@ -307,16 +303,11 @@ class FastBaseModel:
             bnb_compute_dtype = torch.float16
             do_forced_float32 = True
         pass
-
-        global FORCE_EAGER_ATTENTION
-        attn_implementation = "sdpa"
-        for disable_name in FORCE_EAGER_ATTENTION:
-            if (disable_name.lower() == model_type_arch.lower() or \
-                disable_name.lower() in model_name.lower()):
-
-                print(f"Unsloth: {model_type_arch} does not support SDPA - switching to eager!")
-                attn_implementation = "eager"
-                break
+        # Stop SDPA for some archs like Pixtral / Mistral3
+        kwargs["attn_implementation"] = "sdpa"
+        if not supports_sdpa:
+            print(f"Unsloth: {model_type_arch} does not support SDPA - switching to eager!")
+            del kwargs["attn_implementation"]
         pass
 
         bnb_config = None
@@ -355,8 +346,6 @@ class FastBaseModel:
             os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "0"
         pass
 
-        kwargs.pop("attn_implementation", None); # No need since we auto call it
-
         # Cannot be None, since HF now checks for the config
         if load_in_4bit: kwargs["quantization_config"] = bnb_config
 
@@ -370,7 +359,7 @@ class FastBaseModel:
             # quantization_config   = bnb_config,
             token                   = token,
             trust_remote_code       = trust_remote_code,
-            attn_implementation     = attn_implementation,
+            # attn_implementation   = attn_implementation,
             **kwargs,
         )
         # Return old flag