Repository: huggingface/transformers
Pull request: XPU now supports backward for the FA2 fixed path

Per the title.
Enable the corresponding tests and re-enable the tests that were skipped before.

Changed files:
filename: tests/models/cohere2/test_modeling_cohere2.py
patch: @@ -26,6 +26,8 @@
     Expectations,
     cleanup,
     is_flash_attn_2_available,
+    is_kernels_available,
+    is_torch_xpu_available,
     require_flash_attn,
     require_torch,
     require_torch_large_accelerator,
@@ -249,12 +251,13 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):
                 reason="`flex_attention` gives `torch._inductor.exc.InductorError: RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_0 Required: 147456 Hardware limit:101376 Reducing block sizes or `num_stages` may help.`"
             )
 
-        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+        if (
+            attn_implementation == "flash_attention_2"
+            and not is_flash_attn_2_available()
+            and not (is_torch_xpu_available() and is_kernels_available())
+        ):
             self.skipTest("FlashAttention2 is required for this test.")
 
-        if torch_device == "xpu" and attn_implementation == "flash_attention_2":
-            self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
-
         model_id = "CohereForAI/c4ai-command-r7b-12-2024"
         EXPECTED_COMPLETIONS = [
             " the mountains, the lakes, the rivers, the forests, the trees, the birds, the animals",

filename: tests/models/gemma2/test_modeling_gemma2.py
patch: @@ -27,6 +27,8 @@
     Expectations,
     cleanup,
     is_flash_attn_2_available,
+    is_kernels_available,
+    is_torch_xpu_available,
     require_flash_attn,
     require_large_cpu_ram,
     require_torch,
@@ -343,12 +345,13 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):
                 reason="`flex_attention` gives `torch._inductor.exc.InductorError: RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_0 Required: 147456 Hardware limit:101376 Reducing block sizes or `num_stages` may help.`"
             )
 
-        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+        if (
+            attn_implementation == "flash_attention_2"
+            and not is_flash_attn_2_available()
+            and not (is_torch_xpu_available() and is_kernels_available())
+        ):
             self.skipTest("FlashAttention2 is required for this test.")
 
-        if torch_device == "xpu" and attn_implementation == "flash_attention_2":
-            self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
-
         model_id = "google/gemma-2-2b"
         EXPECTED_COMPLETIONS = [
             " the people, the food, the culture, the history, the music, the art, the architecture",
@@ -388,12 +391,13 @@ def test_generation_beyond_sliding_window_dynamic(self, attn_implementation: str
                 reason="`flex_attention` gives `torch._inductor.exc.InductorError: RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_0 Required: 147456 Hardware limit:101376 Reducing block sizes or `num_stages` may help.`"
             )
 
-        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+        if (
+            attn_implementation == "flash_attention_2"
+            and not is_flash_attn_2_available()
+            and not (is_torch_xpu_available() and is_kernels_available())
+        ):
             self.skipTest("FlashAttention2 is required for this test.")
 
-        if torch_device == "xpu" and attn_implementation == "flash_attention_2":
-            self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
-
         model_id = "google/gemma-2-2b"
         EXPECTED_COMPLETIONS = [
             " the people, the food, the culture, the history, the music, the art, the architecture",

filename: tests/models/vaultgemma/test_modeling_vaultgemma.py
patch: @@ -32,6 +32,8 @@
     Expectations,
     cleanup,
     is_flash_attn_2_available,
+    is_kernels_available,
+    is_torch_xpu_available,
     require_torch,
     require_torch_accelerator,
     slow,
@@ -183,12 +185,13 @@ def test_generation_beyond_sliding_window(self, attn_implementation: str):
                 reason="`flex_attention` gives `torch._inductor.exc.InductorError: RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_0 Required: 147456 Hardware limit:101376 Reducing block sizes or `num_stages` may help.`"
             )
 
-        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+        if (
+            attn_implementation == "flash_attention_2"
+            and not is_flash_attn_2_available()
+            and not (is_torch_xpu_available() and is_kernels_available())
+        ):
             self.skipTest("FlashAttention2 is required for this test.")
 
-        if torch_device == "xpu" and attn_implementation == "flash_attention_2":
-            self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
-
         model_id = "google/vaultgemma-1b"
         EXPECTED_COMPLETIONS = [
             " place pretty place pretty place. place pretty place pretty place. place pretty place pretty place. place pretty",
@@ -228,12 +231,13 @@ def test_generation_beyond_sliding_window_dynamic(self, attn_implementation: str
                 reason="`flex_attention` gives `torch._inductor.exc.InductorError: RuntimeError: No valid triton configs. OutOfMemoryError: out of resource: triton_tem_fused_0 Required: 147456 Hardware limit:101376 Reducing block sizes or `num_stages` may help.`"
             )
 
-        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+        if (
+            attn_implementation == "flash_attention_2"
+            and not is_flash_attn_2_available()
+            and not (is_torch_xpu_available() and is_kernels_available())
+        ):
             self.skipTest("FlashAttention2 is required for this test.")
 
-        if torch_device == "xpu" and attn_implementation == "flash_attention_2":
-            self.skipTest(reason="Intel XPU doesn't support flash_attention_2 as of now.")
-
         model_id = "google/vaultgemma-1b"
         EXPECTED_COMPLETIONS = [
             " place pretty place pretty place. place pretty place pretty place. place pretty place pretty place. place pretty",

filename: tests/test_modeling_common.py
patch: @@ -3873,9 +3873,6 @@ def test_flash_attn_2_can_compile_with_attention_mask_None_without_graph_break(s
         if not is_torch_fp16_available_on_device(torch_device):
             self.skipTest(f"float16 not supported on {torch_device} (on the specific device currently used)")
 
-        if torch_device == "xpu":
-            self.skipTest("XPU FA2 currently does not support backward.")
-
         torch.compiler.reset()
         dtype = torch.float16