flash attn pytest marker (#41781)

* flash attn marker

* 111

---------

Co-authored-by: ydshieh <[email protected]>

Author: ydshieh

conftest.py

@@ -87,6 +87,8 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "not_device_test: mark the tests always running on cpu")
     config.addinivalue_line("markers", "torch_compile_test: mark test which tests torch compile functionality")
     config.addinivalue_line("markers", "torch_export_test: mark test which tests torch export functionality")
+    config.addinivalue_line("markers", "flash_attn_test: mark test which tests flash attention functionality")
+    config.addinivalue_line("markers", "flash_attn_3_test: mark test which tests flash attention 3 functionality")
 
     os.environ["DISABLE_SAFETENSORS_CONVERSION"] = "true"
 

tests/models/exaone4/test_modeling_exaone4.py

@@ -120,6 +120,7 @@ def test_model_generation_sdpa(self):
         text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
         self.assertEqual(EXPECTED_TEXT, text)
 
+    @pytest.mark.flash_attn_test
     @slow
     @require_torch_accelerator
     @require_flash_attn

tests/models/idefics2/test_modeling_idefics2.py

@@ -643,6 +643,7 @@ def test_integration_test_4bit_batch2(self):
         self.assertEqual(batched_generated_texts[0], generated_text_0[0])
         self.assertEqual(batched_generated_texts[1], generated_text_1[0])
 
+    @pytest.mark.flash_attn_test
     @require_flash_attn
     @require_torch_gpu
     @require_bitsandbytes

tests/models/ministral/test_modeling_ministral.py

@@ -208,6 +208,7 @@ def test_export_text_with_hybrid_cache(self):
 
         self.assertEqual(export_generated_text, eager_generated_text)
 
+    @pytest.mark.flash_attn_test
     @require_flash_attn
     @slow
     def test_past_sliding_window_generation(self):

tests/models/mistral/test_modeling_mistral.py

@@ -300,6 +300,7 @@ def test_compile_static_cache(self):
         static_compiled_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         self.assertEqual(EXPECTED_TEXT_COMPLETION, static_compiled_text)
 
+    @pytest.mark.flash_attn_test
     @parameterized.expand([("flash_attention_2",), ("sdpa",), ("flex_attention",), ("eager",)])
     @require_flash_attn
     @slow

tests/models/qwen2/test_modeling_qwen2.py

@@ -274,6 +274,7 @@ def test_export_static_cache(self):
         ep_generated_text = tokenizer.batch_decode(ep_generated_ids, skip_special_tokens=True)
         self.assertEqual(EXPECTED_TEXT_COMPLETION, ep_generated_text)
 
+    @pytest.mark.flash_attn_test
     @require_flash_attn
     @slow
     def test_3b_generation(self):

tests/models/qwen2_5_omni/test_modeling_qwen2_5_omni.py

@@ -816,6 +816,7 @@ def test_small_model_integration_test_w_audio(self):
     @slow
     @require_flash_attn
     @require_torch_gpu
+    @pytest.mark.flash_attn_test
     def test_small_model_integration_test_batch_flashatt2(self):
         model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2.5-Omni-7B",

tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py

@@ -17,6 +17,7 @@
 import tempfile
 import unittest
 
+import pytest
 import requests
 
 from transformers import (
@@ -630,6 +631,7 @@ def test_small_model_integration_test_batch_different_resolutions(self):
     @slow
     @require_flash_attn
     @require_torch_gpu
+    @pytest.mark.flash_attn_test
     def test_small_model_integration_test_batch_flashatt2(self):
         model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2.5-VL-7B-Instruct",
@@ -658,6 +660,7 @@ def test_small_model_integration_test_batch_flashatt2(self):
     @slow
     @require_flash_attn
     @require_torch_gpu
+    @pytest.mark.flash_attn_test
     def test_small_model_integration_test_batch_wo_image_flashatt2(self):
         model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2.5-VL-7B-Instruct",

tests/models/qwen2_vl/test_modeling_qwen2_vl.py

@@ -18,6 +18,7 @@
 import tempfile
 import unittest
 
+import pytest
 import requests
 
 from transformers import (
@@ -562,6 +563,7 @@ def test_small_model_integration_test_batch_different_resolutions(self):
     @slow
     @require_flash_attn
     @require_torch_gpu
+    @pytest.mark.flash_attn_test
     def test_small_model_integration_test_batch_flashatt2(self):
         model = Qwen2VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2-VL-7B-Instruct",
@@ -589,6 +591,7 @@ def test_small_model_integration_test_batch_flashatt2(self):
     @slow
     @require_flash_attn
     @require_torch_gpu
+    @pytest.mark.flash_attn_test
     def test_small_model_integration_test_batch_wo_image_flashatt2(self):
         model = Qwen2VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2-VL-7B-Instruct",

tests/models/qwen3/test_modeling_qwen3.py

@@ -266,6 +266,7 @@ def test_export_static_cache(self):
 
     @require_flash_attn
     @slow
+    @pytest.mark.flash_attn_test
     def test_600m_generation(self):
         model_id = "Qwen/Qwen3-0.6B-Base"
         tokenizer = AutoTokenizer.from_pretrained(model_id)