Update PyTorch and use num_stages=1 for flex attn bwd autotuning to reduce CI time (#4785)

anmyachev · web-flow · commit 543ba723c4bf · 2025-07-29T18:04:48.000+02:00
Signed-off-by: Anatoly Myachev &lt;anatoly.myachev@intel.com&gt;
diff --git a/.github/pins/pytorch.txt b/.github/pins/pytorch.txt
@@ -1 +1 @@
-4869f7117009fb99a57482fce56b00c6163fbce6
+6d071bd65de9bdc354f32adf67e00d6e13475e76
diff --git a/.github/workflows/inductor-tests-reusable.yml b/.github/workflows/inductor-tests-reusable.yml
@@ -29,6 +29,7 @@ permissions: read-all
 env:
   PYTHON_VERSION: "${{ inputs.python_version }}"
   PYTORCH_REPO: "${{ inputs.pytorch_repo }}"
+  TRITON_LESS_FLEX_ATTN_BWD_CONFIGS: 1
 
 jobs:
   build:
diff --git a/.github/workflows/inductor-tests-windows.yml b/.github/workflows/inductor-tests-windows.yml
@@ -39,6 +39,7 @@ env:
   PYTHONIOENCODING: utf-8
   PYTHON_VERSION: "${{ inputs.python_version || '3.10' }}"
   NEW_WORKSPACE: C:\gh${{ github.run_id }}
+  TRITON_LESS_FLEX_ATTN_BWD_CONFIGS: 1
 
 jobs:
   build:
diff --git a/scripts/patch-pytorch.sh b/scripts/patch-pytorch.sh
@@ -40,3 +40,4 @@ apply_patch ./patch/pytorch_fp64.patch
 apply_patch ./patch/pytorch_global_scratch.patch
 apply_patch ./patch/test_compile_subprocess.patch
 apply_patch ./patch/flex_decoding.patch
+apply_patch ./patch/flex_attn_bwd_num_stage_1.patch
diff --git a/scripts/patch/flex_attn_143553.patch b/scripts/patch/flex_attn_143553.patch
diff --git a/scripts/patch/flex_attn_bwd_num_stage_1.patch b/scripts/patch/flex_attn_bwd_num_stage_1.patch
