Fix config stamping when SDPA is not supported

When supports_sdpa=False and the attention hierarchy selected sdpa,
the kwarg was deleted but config._attn_implementation remained as
"sdpa". Stamp it to "eager" for consistency before deleting the kwarg.

unsloth/models/_utils.py

@@ -64,7 +64,7 @@ __all__ = [
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
-    "prefer_flex_attn_if_supported",
+    "determine_attention_implementation",
     "patch_fast_lora",
     "validate_loftq_config",
     "RaiseUninitialized",
@@ -222,36 +222,70 @@ def apply_unsloth_gradient_checkpointing(
     return use_gradient_checkpointing
 
 
-def prefer_flex_attn_if_supported(model_class, config):
-    if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0":
-        return None
-    try:
-        from transformers.utils.import_utils import is_torch_flex_attn_available
+def _set_attn_impl(config, impl):
+    """Stamp the chosen attention implementation onto the config object."""
+    if config is not None:
+        setattr(config, "_attn_implementation", impl)
+        if hasattr(config, "attn_implementation"):
+            setattr(config, "attn_implementation", impl)
+    return impl
 
-        if not is_torch_flex_attn_available():
-            return None
-        if model_class is None or not getattr(
-            model_class, "_supports_flex_attn", False
-        ):
-            return None
-        # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during
-        # inference since flex attention returns incorrect results or errors out.
-        # GPT-OSS: left padding issues cause incorrect outputs.
-        # Mllama: _update_causal_mask uses make_flex_block_causal_mask which
-        # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during
-        # decode q_len=1, causing ValueError. Needs transformers update.
-        # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not
-        # support flex_attention.
-        model_type = getattr(config, "model_type", "") if config else ""
-        if model_type in ("gpt_oss", "mllama") or str(model_type).startswith("gemma3n"):
-            return None
-        if config is not None:
-            setattr(config, "_attn_implementation", "flex_attention")
-            if hasattr(config, "attn_implementation"):
-                setattr(config, "attn_implementation", "flex_attention")
-        return "flex_attention"
-    except Exception:
-        return None
+
+def determine_attention_implementation(model_class, config):
+    model_type = getattr(config, "model_type", "").lower() if config else ""
+
+    # 1. Flash Attention 2
+    if (
+        HAS_FLASH_ATTENTION
+        and model_type not in ("gpt_oss", "mllama", "nemotron_h")
+        and not model_type.startswith("gemma3n")
+    ):
+        supports_fa2 = False
+        if model_class is not None:
+            supports_fa2 = getattr(
+                model_class, "_supports_flash_attn_2", False
+            ) or getattr(model_class, "_supports_flash_attn", False)
+
+        if supports_fa2:
+            return _set_attn_impl(config, "flash_attention_2")
+
+    # 2. Flex Attention
+    if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0":
+        try:
+            from transformers.utils.import_utils import is_torch_flex_attn_available
+
+            if (
+                is_torch_flex_attn_available()
+                and (model_class is not None)
+                and getattr(model_class, "_supports_flex_attn", False)
+            ):
+                attention_dropout = getattr(config, "attention_dropout", 0) or 0
+                # GPT-OSS, Mllama, Gemma3N and NemotronH use eager/sdpa attention
+                # during inference since flex attention returns incorrect results
+                # or errors out.
+                # GPT-OSS: left padding issues cause incorrect outputs.
+                # Mllama: _update_causal_mask uses make_flex_block_causal_mask which
+                # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during
+                # decode q_len=1, causing ValueError. Needs transformers update.
+                # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not
+                # support flex_attention.
+                # NemotronH: hybrid Mamba-2 + Transformer model that does not
+                # support flex_attention (raises NotImplementedError from transformers).
+                is_excluded = (
+                    model_type in ("gpt_oss", "mllama", "nemotron_h")
+                    or model_type.startswith("gemma3n")
+                )
+                if attention_dropout == 0 and not is_excluded:
+                    return _set_attn_impl(config, "flex_attention")
+        except Exception:
+            pass
+
+    # 3. SDPA
+    if model_class is not None and getattr(model_class, "_supports_sdpa", False):
+        return _set_attn_impl(config, "sdpa")
+
+    # 4. Eager
+    return _set_attn_impl(config, "eager")
 
 
 def _run_temporary_patches(phase):

unsloth/models/llama.py

@@ -2341,8 +2341,8 @@ class FastLlamaModel:
         model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__]
         IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1")
 
-        preferred_attn_impl = (
-            prefer_flex_attn_if_supported(model_function, model_config) or "eager"
+        preferred_attn_impl = determine_attention_implementation(
+            model_function, model_config
         )
 
         has_rope_scaling = False

unsloth/models/vision.py

@@ -29,7 +29,7 @@ except:
 from ..kernels import (
     post_patch_loss_function,
 )
-from ._utils import __version__, importlib_version, _prepare_model_for_qat
+from ._utils import __version__, importlib_version, _prepare_model_for_qat, _set_attn_impl
 from ._utils import *
 from .loader_utils import _get_fp8_mode_and_check_settings
 from ..save import patch_saving_functions
@@ -597,8 +597,7 @@ class FastBaseModel:
                 custom_datatype = None
                 correct_dtype = None
 
-        # Stop SDPA for some archs like Pixtral / Mistral3
-        flex_attn_impl = None
+        # Unified hierarchical attention fallback: Flash > Flex > SDPA > Eager
         if auto_config is None:
             auto_config = AutoConfig.from_pretrained(
                 model_name,
@@ -609,7 +608,20 @@ class FastBaseModel:
             model_class = auto_model._model_mapping[auto_config.__class__]
         except Exception:
             model_class = None
-        flex_attn_impl = prefer_flex_attn_if_supported(model_class, auto_config)
+
+        model_type = str(getattr(auto_config, "model_type", "")).lower()
+        if model_type.startswith("gemma3n"):
+            # Gemma3N variants use timm-based vision towers which do not support
+            # flex_attention. The old code defaulted gemma3n to eager; preserve
+            # that behavior rather than letting the hierarchy pick sdpa.
+            attn_impl = _set_attn_impl(auto_config, "eager")
+        elif model_class is None and supports_sdpa:
+            # When model_class cannot be resolved (remote-code or unmapped
+            # configs), the old code defaulted to sdpa. Preserve that fallback
+            # instead of falling through to eager.
+            attn_impl = _set_attn_impl(auto_config, "sdpa")
+        else:
+            attn_impl = determine_attention_implementation(model_class, auto_config)
 
         # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with
         # FP8 weights. We just need to update it here for sanity.
@@ -620,20 +632,14 @@ class FastBaseModel:
         except Exception:
             model_class = None
 
-        model_type = str(getattr(auto_config, "model_type", "")).lower()
-        if model_type.startswith("gemma3n"):
-            # Gemma3N variants initialize timm-based vision towers which do
-            # not support flex_attention, so default to eager unless overridden.
-            default_attn_impl = "eager"
-        else:
-            default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa"
         if not ("attn_implementation" in kwargs):
-            kwargs["attn_implementation"] = default_attn_impl
+            kwargs["attn_implementation"] = attn_impl
         if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa":
             if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0":
                 print(
                     f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager."
                 )
+            _set_attn_impl(auto_config, "eager")
             del kwargs["attn_implementation"]
 
         bnb_config = None