[Flex Attention] Apply patch from pytorch#143553 instead of using fork (#3945)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

.github/actions/setup-pytorch/action.yml

@@ -45,14 +45,8 @@ runs:
       if: inputs.ref != ''
       shell: bash
       run: |
-        if [[ "${{ inputs.repository }}" = "liangan1/pytorch" ]]; then
-          PYTORCH_COMMIT_ID="$(<.github/pins/pytorchFlexAttention.txt)"
-          echo "PYTORCH_REPO=${{ inputs.repository }}" | tee -a "$GITHUB_ENV"
-          echo "PYTORCH_COMMIT_ID=$PYTORCH_COMMIT_ID" | tee -a "$GITHUB_ENV"
-        else
-          echo "PYTORCH_REPO=${{ inputs.repository }}" | tee -a "$GITHUB_ENV"
-          echo "PYTORCH_COMMIT_ID=${{ steps.commit-id.outputs.commit_id }}" | tee -a "$GITHUB_ENV"
-        fi
+        echo "PYTORCH_REPO=${{ inputs.repository }}" | tee -a "$GITHUB_ENV"
+        echo "PYTORCH_COMMIT_ID=${{ steps.commit-id.outputs.commit_id }}" | tee -a "$GITHUB_ENV"
 
     - name: Identify Python version
       shell: bash
@@ -105,7 +99,7 @@ runs:
         path: pytorch
 
     - name: Apply additional PR patches
-      if: ${{ steps.pytorch-cache.outputs.status == 'miss'  && inputs.mode == 'source' && (inputs.repository == 'pytorch/pytorch' || inputs.repository == 'liangan1/pytorch') }}
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.repository == 'pytorch/pytorch' && inputs.mode == 'source' }}
       shell: bash
       run: |
         cd pytorch

.github/pins/pytorchFlexAttention.txt

@@ -274,14 +274,6 @@ jobs:
           cd benchmarks/micro_benchmarks
           python run_benchmarks.py --reports $REPORTS
 
-      # Install Pytorch with FlexAttention XPU support enabled
-      - name: Setup PyTorch
-        if: ${{ steps.install.outcome == 'success' && !cancelled() }}
-        uses: ./.github/actions/setup-pytorch
-        with:
-          repository: liangan1/pytorch
-          ref: liangan1/flex_attention
-
       - name: Run Triton FlexAttention Causal Mask fwd kernel benchmark
         if: ${{ steps.install.outcome == 'success' && !cancelled() && (inputs.benchmarks == '' || contains(fromJson(inputs.benchmarks || '[]'), 'flex_attention_benchmark_causal_mask.py')) && !contains(fromJson(inputs.skip_benchmarks || '[]'), 'flex_attention_benchmark_causal_mask.py') }}
         run: |

.github/workflows/triton-benchmarks.yml

@@ -18,3 +18,4 @@ cd "$REPO_ROOT"
 
 # curl -sSL https://github.com/pytorch/pytorch/pull/126516.diff | git apply -
 git apply "${SCRIPT_DIR}/pytorch_fp64.patch"
+curl -sSL https://github.com/pytorch/pytorch/pull/143553.diff | git apply --exclude=test/inductor/test_flex_attention.py --exclude=test/inductor/test_flex_decoding.py -