Repository: huggingface/transformers
File: tests/models/gpt_oss/test_modeling_gpt_oss.py (Python)

Code before:
        GptOssModel,
    )

    NUM_GPUS = torch.cuda.device_count()


class GptOssModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = GptOssModel


@require_torch
class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
    _is_stateful = True
    model_split_percents = [0.5, 0.6]
    model_tester_class = GptOssModelTester

    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 = self.model_tester.get_config()
        config._attn_implementation = "flash_attention_2"
        with self.assertRaisesRegex(ValueError, "GPT-OSS models do not support"):
            GptOssModel(config)

        config._attn_implementation = "kernels-community/vllm-flash-attn3"
        try:
            GptOssModel(config)
        except ValueError:
            self.fail("GptOssModel raised ValueError unexpectedly with vllm-flash-attn3!")

    @unittest.skip("GptOss's forcefully disables sdpa due to Sink")
    def test_sdpa_can_dispatch_non_composite_models(self):
        pass

    @unittest.skip("GptOss's eager attn/sdpa attn outputs are expected to be different")
    def test_eager_matches_sdpa_generate(self):
        pass

    @unittest.skip("GptOss eager/FA2 attention outputs are expected to be different")
    def test_flash_attn_2_equivalence(self):
        pass

    @unittest.skip("Most probably because of the MOE, the moe and router does not ignore padding tokens")
    def test_eager_padding_matches_padding_free_with_position_ids(self):
        pass

    @unittest.skip("GptOss does not support flex officially")

Review comment (suggestion):
We have decorators for non-standard library, here you should use `@require_kernels` and `@require_flash_attn`

Code after:
if is_torch_available():
    import torch

    from transformers import (
        GptOssModel,
    )

    NUM_GPUS = torch.cuda.device_count()


class GptOssModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = GptOssModel


@require_torch
class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
    _is_stateful = True
    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"