[ROCm][CI] Changed to flex attention for cross-attention

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Author: AndreasKaratzas

tests/entrypoints/openai/test_transcription_validation.py

@@ -28,20 +28,19 @@
 
 
 @pytest.fixture(scope="module", autouse=True)
-def rocm_aiter_fa_attention():
+def rocm_flex_attention():
     """
-    Automatically sets VLLM_ATTENTION_BACKEND=ROCM_AITER_FA for ROCm
-    for the duration of this test module.
+    Sets VLLM_ATTENTION_BACKEND=FLEX_ATTENTION for ROCm
+    for the duration of this test module. For now the only
+    attention backend that supports cross attention on ROCm
+    is FLEX_ATTENTION.
     """
     from vllm.platforms import current_platform
 
     if current_platform.is_rocm():
-        # Store previous value to restore later (cleanup)
         old_backend = os.environ.get("VLLM_ATTENTION_BACKEND")
-        # Set the specific backend required for audio models on ROCm
-        os.environ["VLLM_ATTENTION_BACKEND"] = "ROCM_AITER_FA"
+        os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"
         yield
-        # Cleanup: Restore the environment
         if old_backend is None:
             del os.environ["VLLM_ATTENTION_BACKEND"]
         else:

tests/entrypoints/openai/test_translation_validation.py

@@ -22,18 +22,15 @@
 @pytest.fixture(scope="module", autouse=True)
 def rocm_aiter_fa_attention():
     """
-    Automatically sets VLLM_ATTENTION_BACKEND=ROCM_AITER_FA for ROCm
+    Sets VLLM_ATTENTION_BACKEND=ROCM_AITER_FA for ROCm
     for the duration of this test module.
     """
     from vllm.platforms import current_platform
 
     if current_platform.is_rocm():
-        # Store previous value to restore later (cleanup)
         old_backend = os.environ.get("VLLM_ATTENTION_BACKEND")
-        # Set the specific backend required for audio models on ROCm
         os.environ["VLLM_ATTENTION_BACKEND"] = "ROCM_AITER_FA"
         yield
-        # Cleanup: Restore the environment
         if old_backend is None:
             del os.environ["VLLM_ATTENTION_BACKEND"]
         else: