fix(gpt-oss): prefer flex attention over sdpa (#5701)

* fix(gpt-oss): prefer flex attention over sdpa

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

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

* fix(gpt-oss): use eager config for unsupported backends

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

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

---------

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

tests/test_attention_implementation.py

@@ -0,0 +1,56 @@
+from types import SimpleNamespace
+
+import unsloth  # noqa: F401
+from transformers.utils import import_utils
+
+from unsloth.models import _utils
+
+
+class SupportsFlexAndSdpa:
+    _supports_flash_attn_2 = True
+    _supports_flex_attn = True
+    _supports_sdpa = True
+
+
+def _config(model_type, **kwargs):
+    values = {"model_type": model_type, "attention_dropout": 0}
+    values.update(kwargs)
+    return SimpleNamespace(**values)
+
+
+def _set_flex_available(monkeypatch, available):
+    monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "1")
+    monkeypatch.setattr(
+        import_utils,
+        "is_torch_flex_attn_available",
+        lambda: available,
+        raising = False,
+    )
+
+
+def test_gpt_oss_uses_eager_instead_of_flash_flex_or_sdpa(monkeypatch):
+    _set_flex_available(monkeypatch, True)
+    config = _config("gpt_oss")
+
+    impl = _utils.resolve_attention_implementation(
+        SupportsFlexAndSdpa,
+        config,
+        supports_sdpa = True,
+    )
+
+    assert impl == "eager"
+    assert config._attn_implementation == "eager"
+
+
+def test_gpt_oss_falls_back_to_eager_when_flex_unavailable(monkeypatch):
+    _set_flex_available(monkeypatch, False)
+    config = _config("gpt_oss")
+
+    impl = _utils.resolve_attention_implementation(
+        SupportsFlexAndSdpa,
+        config,
+        supports_sdpa = True,
+    )
+
+    assert impl == "eager"
+    assert config._attn_implementation == "eager"

unsloth/models/_utils.py

@@ -226,8 +226,10 @@ def apply_unsloth_gradient_checkpointing(
     return use_gradient_checkpointing
 
 
-# Models that don't work with flex_attention:
-# GPT-OSS: left padding issues cause incorrect outputs.
+# Models that don't work with flex_attention as the global Transformers
+# attention implementation:
+# GPT-OSS: training uses the custom flex sink patch, but inference intentionally
+# falls back to eager because flex decoding gives incorrect outputs.
 # Mllama: BlockMask Q_LEN!=KV_LEN ValueError on decode.
 # NemotronH: hybrid Mamba-2 + Transformer, raises NotImplementedError.
 # Gemma3N: timm vision wrappers don't support flex_attention.
@@ -235,6 +237,8 @@ def apply_unsloth_gradient_checkpointing(
 # access on some GPU architectures (B200). Falls back to eager safely.
 _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert")
 _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2")
+_SDPA_EXCLUDED_MODELS = ("gpt_oss",)
+_FLASH_EXCLUDED_MODELS = ("gpt_oss",)
 _EAGER_ONLY_PREFIXES = ("gemma3n",)
 _FLASH_ATTENTION_MAX_HEAD_DIM = 256
 _FLASH_ATTENTION_DISABLED_WARNED = set()
@@ -244,6 +248,14 @@ def _is_flex_excluded(model_type):
     return model_type in _FLEX_EXCLUDED_MODELS
 
 
+def _is_sdpa_excluded(model_type):
+    return model_type in _SDPA_EXCLUDED_MODELS
+
+
+def _is_flash_excluded(model_type):
+    return model_type in _FLASH_EXCLUDED_MODELS
+
+
 def _config_prefers_flex_attention(config):
     return any(
         _config_get(attention_config, "model_type", "").lower()
@@ -360,6 +372,9 @@ def _get_max_attention_head_dim(config):
 
 
 def _get_flash_attention_disable_reason(config):
+    model_type = _config_get(config, "model_type", "").lower()
+    if _is_flash_excluded(model_type):
+        return f"{model_type} uses custom sink attention kernels"
     max_head_dim = _get_max_attention_head_dim(config)
     if max_head_dim is not None and max_head_dim > _FLASH_ATTENTION_MAX_HEAD_DIM:
         return (
@@ -489,9 +504,15 @@ def resolve_attention_implementation(
         supports_sdpa = model_class is not None and getattr(
             model_class, "_supports_sdpa", False
         )
-    supports_flash_attention = model_class is not None and (
-        getattr(model_class, "_supports_flash_attn_2", False)
-        or getattr(model_class, "_supports_flash_attn", False)
+    if _is_sdpa_excluded(model_type):
+        supports_sdpa = False
+    supports_flash_attention = (
+        model_class is not None
+        and (
+            getattr(model_class, "_supports_flash_attn_2", False)
+            or getattr(model_class, "_supports_flash_attn", False)
+        )
+        and not _is_flash_excluded(model_type)
     )
     supports_flex_attention = _supports_flex_attention(model_class, config, model_type)
     disable_reason = _get_flash_attention_disable_reason(config)

unsloth/models/loader.py

@@ -140,6 +140,7 @@ global DISABLE_SDPA_MODEL_NAMES
 DISABLE_SDPA_MODEL_NAMES = [
     "gemma3,",  # Add comma bc gemma3 will match gemma3n
     "gemma3_text",  # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore
+    "gpt_oss",
 ]