Honor an explicit sdpa or flex_attention request when flash is disabled (#6847)

* Honor an explicit sdpa or flex_attention request when flash is disabled

When flash attention is disabled for a model, the fallback selection could
downgrade a caller who explicitly passed attn_implementation='sdpa' or
'flex_attention' to a different backend, because the disable reason is
flash-specific. Keep an explicit non-flash request as-is; flash requests
still fall back as before.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten comments

* Gate honor-explicit attention on provenance and flex support

Only honor an explicit non-flash attention request when it comes from the
caller argument, not from a config value the loaders synthesize (the language
path seeds attn_implementation=sdpa). Honor explicit flex_attention only when
supports_flex_attention is True so excluded/broken configs (e.g. gpt_oss) fall
back instead of selecting a known-broken backend. Explicit sdpa stays honored.

* Honor explicit sdpa through the resolver guard

* Keep SDPA exclusions when honoring an explicit sdpa request

An explicit attn_implementation="sdpa" was re-enabling sdpa for models in
_SDPA_EXCLUDED_MODELS (e.g. gpt_oss) where sdpa is known-broken: the helper
honored the request and the resolver's final not-supports_sdpa guard skipped
the eager downgrade for any explicit request. Honor an explicit sdpa only when
the model is not sdpa-excluded, mirroring the flex guard that already falls
back for _FLEX_EXCLUDED_MODELS via supports_flex_attention. Conservative
supports_sdpa=False (large head dim / attention-sink models) still honors an
explicit sdpa; a synthesized/default sdpa still downgrades to eager.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Honor DISABLE_SDPA_MODEL_NAMES when honoring explicit sdpa

The honor-explicit-sdpa guard only skipped the sdpa->eager downgrade for models
in _SDPA_EXCLUDED_MODELS (gpt_oss). Gemma3/Gemma3Text disable SDPA through the
loader's DISABLE_SDPA_MODEL_NAMES (their bundled SDPA modules are wrong), so an
explicit sdpa request bypassed the downgrade and re-enabled a known-wrong path.

Extend _is_sdpa_excluded to also treat DISABLE_SDPA_MODEL_NAMES membership as
excluded, replicating the loader's trailing-comma substring match so gemma3 and
gemma3_text match but gemma3n does not. Move the constant into _utils.py (single
source of truth, re-exported from loader.py) to avoid a loader -> _utils cycle.
Conservative supports_sdpa=False models not in either list still honor explicit
sdpa.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

tests/test_attn_impl_honor_explicit.py

@@ -0,0 +1,190 @@
+"""An explicit non-flash attention request must survive the flash disable path.
+
+When flash attention is disabled for a model, a caller who explicitly asked for
+"sdpa" or "flex_attention" should keep that choice instead of being downgraded
+to whatever the conservative supports_* fallback would pick.
+"""
+
+import pytest
+
+from unsloth.models._utils import (
+    _disable_flash_attention_if_needed,
+    resolve_attention_implementation,
+)
+
+
+def test_explicit_sdpa_is_honored_even_when_not_marked_supported():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "sdpa",
+        supports_sdpa = False,  # conservative flag would have skipped sdpa
+        supports_flex_attention = False,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_explicit_flex_is_honored_when_supported():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flex_attention",
+        supports_sdpa = True,
+        supports_flex_attention = True,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "flex_attention"
+    assert config.get("_attn_implementation") == "flex_attention"
+
+
+def test_explicit_flex_falls_back_when_not_supported():
+    # flex_attention is False for known-broken/excluded configs (e.g. gpt_oss),
+    # so an explicit flex request must not select that backend - it falls back.
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flex_attention",
+        supports_sdpa = True,
+        supports_flex_attention = False,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+
+
+def test_synthesized_config_sdpa_is_not_treated_as_explicit():
+    # The language loader seeds the config with attn_implementation="sdpa"; when the
+    # caller passes nothing, that synthesized value must not override the flex fallback
+    # for a model that supports flex but not sdpa.
+    config = {"attn_implementation": "sdpa"}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = None,
+        supports_sdpa = False,
+        supports_flex_attention = True,
+        would_use_flash_attention = False,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "flex_attention"
+
+
+def test_no_disable_reason_returns_request_untouched():
+    result = _disable_flash_attention_if_needed(
+        {},
+        attn_implementation = "flash_attention_2",
+        disable_reason = None,
+    )
+    assert result == "flash_attention_2"
+
+
+def test_flash_request_still_falls_back_when_disabled():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flash_attention_2",
+        supports_sdpa = True,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+
+
+def test_resolver_honors_explicit_sdpa_when_not_supported_and_flash_disabled():
+    # End-to-end through the public resolver: an explicit sdpa request with a
+    # flash-disabled config (oversized head dim) and supports_sdpa=False must not be
+    # rewritten to eager by the resolver's own not-supports_sdpa guard.
+    config = {"model_type": "test", "head_dim": 512}  # head_dim > 256 disables flash
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_resolver_downgrades_non_explicit_sdpa_when_not_supported():
+    # No explicit request: the model resolution seeds sdpa/eager and the guard must
+    # still downgrade a synthesized sdpa to eager for a model that cannot run it.
+    config = {"model_type": "test", "attn_implementation": "sdpa"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = None,
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+
+
+def test_resolver_downgrades_explicit_sdpa_for_sdpa_excluded_model():
+    # gpt_oss is in _SDPA_EXCLUDED_MODELS (sdpa is known-broken) and _FLASH_EXCLUDED_MODELS
+    # (flash disabled). Honoring an explicit sdpa request must not re-enable that broken
+    # backend: it downgrades to eager, mirroring how an explicit flex request falls back
+    # for _FLEX_EXCLUDED_MODELS. supports_sdpa=True proves the exclusion overrides even a
+    # model that otherwise advertises SDPA support.
+    config = {"model_type": "gpt_oss"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = True,
+    )
+    assert result == "eager"
+    assert config.get("_attn_implementation") == "eager"
+
+
+@pytest.mark.parametrize("model_type", ["gemma3", "gemma3_text"])
+def test_resolver_downgrades_explicit_sdpa_for_disable_sdpa_model(model_type):
+    # gemma3 / gemma3_text are in DISABLE_SDPA_MODEL_NAMES: the loader forces
+    # supports_sdpa=False because their bundled SDPA modules are wrong. An explicit
+    # sdpa request with flash disabled must NOT re-enable that known-wrong path - it
+    # downgrades to eager, exactly like _SDPA_EXCLUDED_MODELS (gpt_oss). head_dim>256
+    # disables flash to mirror the real flash-disabled scenario.
+    config = {"model_type": model_type, "head_dim": 512}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+    assert config.get("_attn_implementation") == "eager"
+
+
+def test_resolver_does_not_overmatch_gemma3n_for_explicit_sdpa():
+    # The "gemma3," trailing-comma guard must not match gemma3n: gemma3n is not in
+    # DISABLE_SDPA_MODEL_NAMES, so it stays a conservative (not known-wrong) model and an
+    # explicit sdpa request is still honored. Proves the substring match neither over- nor
+    # under-matches.
+    config = {"model_type": "gemma3n", "head_dim": 512}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_resolver_downgrades_synthesized_sdpa_for_disable_sdpa_model():
+    # A synthesized/default sdpa (requested is None; the value came from config) on a
+    # DISABLE_SDPA_MODEL_NAMES model must still downgrade to eager.
+    config = {"model_type": "gemma3", "attn_implementation": "sdpa"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = None,
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__, "-q"]))

unsloth/models/_utils.py

@@ -423,6 +423,18 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len
 _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert")
 _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2")
 _SDPA_EXCLUDED_MODELS = ("gpt_oss",)
+# The loader (loader.py) forces supports_sdpa=False for these because their bundled
+# SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor
+# them without a loader -> _utils import cycle (loader.py already imports from _utils
+# and re-exports this name for callers like sentence_transformer.py). Entries are matched
+# as substrings against a comma-joined model_types string ending in a comma, so "gemma3,"
+# matches a distinct "gemma3" entry but not "gemma3n", and "gemma3_text" matches the
+# EmbeddingGemma text model.
+DISABLE_SDPA_MODEL_NAMES = [
+    "gemma3,",  # Add comma bc gemma3 will match gemma3n
+    "gemma3_text",  # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore
+    "gpt_oss",
+]
 _FLASH_EXCLUDED_MODELS = ("gpt_oss",)
 _EAGER_ONLY_PREFIXES = ("gemma3n",)
 _FLASH_ATTENTION_MAX_HEAD_DIM = 256
@@ -433,8 +445,23 @@ def _is_flex_excluded(model_type):
     return model_type in _FLEX_EXCLUDED_MODELS
 
 
+def _is_sdpa_disabled_by_name(model_type):
+    # Mirror the loader's DISABLE_SDPA_MODEL_NAMES check: loader.py builds
+    # model_types_all = ",".join(model_types) + "," and tests `name in model_types_all`.
+    # Rebuild the same trailing-comma form for a single model_type so the match is
+    # identical (e.g. "gemma3," matches "gemma3" but not "gemma3n", and "gemma3_text"
+    # still matches "gemma3_text").
+    model_types_all = model_type.lower() + ","
+    return any(name.lower() in model_types_all for name in DISABLE_SDPA_MODEL_NAMES)
+
+
 def _is_sdpa_excluded(model_type):
-    return model_type in _SDPA_EXCLUDED_MODELS
+    # SDPA is known-broken for these models, so an explicit sdpa request must not
+    # re-enable it. Two sources: _SDPA_EXCLUDED_MODELS (resolver-level, e.g. gpt_oss)
+    # and DISABLE_SDPA_MODEL_NAMES (loader-level, e.g. gemma3 / gemma3_text, which the
+    # loader also forces to supports_sdpa=False).
+    lowered = model_type.lower()
+    return lowered in _SDPA_EXCLUDED_MODELS or _is_sdpa_disabled_by_name(lowered)
 
 
 def _is_flash_excluded(model_type):
@@ -610,6 +637,12 @@ def _disable_flash_attention_if_needed(
     if disable_reason is None:
         return attn_implementation
 
+    # Only an implementation passed by the caller counts as an explicit request.
+    # Values read from the config are synthesized by the loaders (the language path
+    # seeds the config with attn_implementation="sdpa") or come from Transformers
+    # defaults, so they must not be treated as a deliberate user choice.
+    explicit_request = attn_implementation
+
     requested_attn_implementation = attn_implementation
     if requested_attn_implementation is None:
         requested_attn_implementation = _config_get(config, "_attn_implementation", None)
@@ -619,6 +652,20 @@ def _disable_flash_attention_if_needed(
     if requested_attn_implementation == "eager":
         return _set_attn_impl(config, "eager")
 
+    model_type = _config_get(config, "model_type", "")
+
+    # The disable reason is flash-specific: honor an explicit non-flash request from
+    # the caller instead of downgrading it. SDPA is honored unless the model's SDPA is
+    # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES
+    # (e.g. gemma3 / gemma3_text); flex_attention
+    # is honored only when it is actually usable, since supports_flex_attention already
+    # rejects the excluded/broken/unavailable configs. This keeps an explicit request
+    # from selecting a backend the repo marks as wrong.
+    if explicit_request == "sdpa" and not _is_sdpa_excluded(model_type.lower()):
+        return _set_attn_impl(config, "sdpa")
+    if explicit_request == "flex_attention" and supports_flex_attention:
+        return _set_attn_impl(config, "flex_attention")
+
     if supports_sdpa:
         fallback_attn_implementation = "sdpa"
     elif supports_flex_attention:
@@ -631,7 +678,6 @@ def _disable_flash_attention_if_needed(
             if _is_flash_attention_requested(requested_attn_implementation)
             else "flash_attention_2"
         )
-        model_type = _config_get(config, "model_type", "")
         warning_key = (
             model_type,
             logged_attn_implementation,
@@ -845,7 +891,19 @@ def resolve_attention_implementation(
         final_attn_impl = requested_attn_implementation
         _set_attn_impl(config, final_attn_impl)
 
-    if not supports_sdpa and final_attn_impl == "sdpa":
+    # A caller who explicitly passes requested_attn_implementation="sdpa" keeps it even
+    # on a conservatively unsupported model, mirroring _disable_flash_attention_if_needed
+    # which honors an explicit sdpa request. The exception is a model whose SDPA is
+    # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES
+    # (e.g. gemma3 / gemma3_text, which the loader also forces to supports_sdpa=False):
+    # an explicit request must not re-enable it, so it still downgrades to eager, just
+    # like flex falls back for _FLEX_EXCLUDED_MODELS. A synthesized/default sdpa
+    # (requested is None, so the value came from the model resolution above or the
+    # config) also downgrades.
+    honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded(
+        model_type
+    )
+    if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa:
         print(
             f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager."
         )

unsloth/models/loader.py

@@ -21,6 +21,10 @@ from ._utils import (
     USE_MODELSCOPE,
     get_transformers_model_type,
     hf_login,
+    # Single source of truth is _utils.py; re-exported here so callers doing
+    # `from unsloth.models.loader import DISABLE_SDPA_MODEL_NAMES` keep working and so
+    # _is_sdpa_excluded (in _utils) can honor it without a loader -> _utils cycle.
+    DISABLE_SDPA_MODEL_NAMES,
 )
 from .granite import FastGraniteModel
 from .llama import FastLlamaModel, logger
@@ -196,14 +200,6 @@ DISABLE_COMPILE_MODEL_NAMES = [
     "granite,llava_next",  # Granite-vision 3
 ]
 
-global DISABLE_SDPA_MODEL_NAMES
-# Disables some SDPA modules since it's wrong
-DISABLE_SDPA_MODEL_NAMES = [
-    "gemma3,",  # Add comma bc gemma3 will match gemma3n
-    "gemma3_text",  # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore
-    "gpt_oss",
-]
-
 
 def _fix_rope_inv_freq(model):
     """Fix inv_freq corruption caused by transformers v5 meta-device loading.