Skip to content

XPU now supports backward for the FA2 fixed path #43905

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vasqu merged 4 commits into from
Feb 26, 2026
Merged

XPU now supports backward for the FA2 fixed path #43905

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
befca71
XPU supports fixed FA2 bwd now
YangKai0616 Feb 10, 2026
115684a
Merge branch 'main' into FA2-XPU
YangKai0616 Feb 11, 2026
66cbfa9
Update judgment conditions
YangKai0616 Feb 12, 2026
61f6d84
Merge branch 'main' into FA2-XPU
YangKai0616 Feb 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions tests/models/cohere2/test_modeling_cohere2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
20 changes: 12 additions & 8 deletions tests/models/gemma2/test_modeling_gemma2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
20 changes: 12 additions & 8 deletions tests/models/vaultgemma/test_modeling_vaultgemma.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
Expectations,
cleanup,
is_flash_attn_2_available,
is_kernels_available,
is_torch_xpu_available,
require_torch,
require_torch_accelerator,
slow,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
3 changes: 0 additions & 3 deletions tests/test_modeling_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down