Repository: huggingface/transformers
Pull request: gpt-oss is not working with flash-attention

When initializing `gpt-oss` model with `attn_implementation="flash_attention_2"` or `"flash_attention_3"` would result in silent failures and garbage generation output as reported in #42533.

`gpt-oss` models rely on attention sinks which are not yet implemented for the `flash_attention` as suggested the safest path is to strictly block unsupported attention backends rather than failing silently or assuming a fallback.

@vasqu can you see if this is what you had in mind?

Changed files:
filename: src/transformers/models/gpt_oss/configuration_gpt_oss.py
patch: @@ -117,5 +117,22 @@ def __init__(
             **kwargs,
         )
 
+    def __setattr__(self, key, value):
+        """
+        Overwritten to allow checking for the proper attention implementation to be used.
+
+        Due to `set_attn_implementation` which internally assigns `_attn_implementation_internal = "..."`, simply overwriting
+        the specific attention setter is not enough. Using a property/setter for `_attn_implementation_internal` would result in
+        a recursive dependency (as `_attn_implementation` acts as a wrapper around `_attn_implementation_internal`) - hence, this
+        workaround.
+        """
+        if key in ("_attn_implementation", "_attn_implementation_internal"):
+            if value and "flash" in value and value.removeprefix("paged|") != "kernels-community/vllm-flash-attn3":
+                raise ValueError(
+                    f"GPT-OSS model does not support the specified flash attention implementation: {value}. "
+                    "Only `kernels-community/vllm-flash-attn3` is supported."
+                )
+        super().__setattr__(key, value)
+
 
 __all__ = ["GptOssConfig"]

filename: src/transformers/models/gpt_oss/modeling_gpt_oss.py
patch: @@ -434,7 +434,7 @@ class GptOssPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True
     _supports_sdpa = False
-    _supports_flex_attn = True
+    _supports_flex_attn = False
 
     _can_compile_fullgraph = True
     _supports_attention_backend = True
@@ -444,8 +444,6 @@ class GptOssPreTrainedModel(PreTrainedModel):
         "attentions": GptOssAttention,
     }
     _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
-    _supports_flash_attention = False
-    _supports_flex_attention = False
 
     @torch.no_grad()
     def _init_weights(self, module):

filename: src/transformers/models/gpt_oss/modular_gpt_oss.py
patch: @@ -354,8 +354,7 @@ def forward(
 class GptOssPreTrainedModel(LlamaPreTrainedModel):
     _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
     _supports_sdpa = False
-    _supports_flash_attention = False
-    _supports_flex_attention = False
+    _supports_flex_attn = False
     _can_record_outputs = {
         "router_logits": OutputRecorder(GptOssTopKRouter, index=0),
         "hidden_states": GptOssDecoderLayer,

filename: tests/models/gpt_oss/test_modeling_gpt_oss.py
patch: @@ -22,6 +22,7 @@
 import unittest
 from pathlib import Path
 
+import pytest
 from parameterized import parameterized
 
 from transformers import (
@@ -31,9 +32,12 @@
 )
 from transformers.testing_utils import (
     cleanup,
+    require_flash_attn,
+    require_kernels,
     require_read_token,
     require_torch,
     require_torch_accelerator,
+    require_torch_gpu,
     slow,
     torch_device,
 )
@@ -62,6 +66,37 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
     model_split_percents = [0.5, 0.6]
     model_tester_class = GptOssModelTester
 
+    @require_kernels
+    @require_flash_attn
+    @pytest.mark.flash_attn_test
+    @require_torch_gpu
+    def test_initialization_raises_error_for_flash_attn(self):
+        """
+        Tests that initializing the model with unsupported Flash Attention implementations raises a ValueError,
+        but allows the specific vllm kernel.
+        """
+
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        kernel_attn = "kernels-community/vllm-flash-attn3"
+
+        # Checking each via `set_attn_implementation` and manually setting within the config
+        model = GptOssModel(config).to(device=torch_device, dtype=torch.bfloat16)
+        model.set_attn_implementation("kernels-community/vllm-flash-attn3")
+        self.assertTrue(model.config._attn_implementation == kernel_attn)
+
+        config._attn_implementation = kernel_attn
+        self.assertTrue(model.config._attn_implementation == kernel_attn)
+
+        with torch.no_grad():
+            output = model(**inputs_dict)
+        self.assertIsNotNone(output)
+
+        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
+            model.set_attn_implementation("flash_attention_2")
+
+        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
+            config._attn_implementation = "flash_attention_2"
+
     @unittest.skip("GptOss's forcefully disables sdpa due to Sink")
     def test_sdpa_can_dispatch_non_composite_models(self):
         pass