Use large runners on Windows to build with FA

.github/workflows/wheels.yml

@@ -19,8 +19,8 @@ jobs:
       fail-fast: false
       matrix:
         os:
-          - 4-core-ubuntu
-          - windows-2019
+          - 8-core-ubuntu
+          - windows-8-core
         python:
           - "3.8"
           - "3.9"
@@ -39,7 +39,7 @@ jobs:
           - torch_version: "2.1.0"
             publish: true
           # Publish source distribution only from this runner
-          - os: 4-core-ubuntu
+          - os: 8-core-ubuntu
             python: "3.10"
             torch_version: "2.1.0"
             sdist: true
@@ -59,4 +59,3 @@ jobs:
       twine_username: __token__
     secrets:
       twine_password: ${{ secrets.PYPI_TOKEN }}
-

.github/workflows/wheels_reusable.yml

@@ -39,7 +39,7 @@ on:
 env:
   # you need at least cuda 5.0 for some of the stuff compiled here.
   TORCH_CUDA_ARCH_LIST: "5.0+PTX 6.0 6.1 7.0 7.5 8.0+PTX"
-  MAX_JOBS: 1
+  MAX_JOBS: 4
   DISTUTILS_USE_SDK: 1 # otherwise distutils will complain on windows about multiple versions of msvc
   XFORMERS_BUILD_TYPE: "Release"
   TWINE_USERNAME: __token__
@@ -128,7 +128,7 @@ jobs:
         run: |
           set -Eeuo pipefail
           git config --global --add safe.directory "*"
-          $PY -m pip install packaging
+          $PY -m pip install packaging ninja
           version=`$PY packaging/compute_wheel_version.py`
           echo $version > version.txt
           echo "BUILD_VERSION=$version${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}" >> ${GITHUB_ENV}
@@ -152,14 +152,6 @@ jobs:
           ./aws/install
           aws --version
 
-      # TODO: This is needed because github-actions takes more than 6 hours (!)
-      # to build Flash-Attention on windows, and times out.
-      # We need to use bigger gh-action workers...
-      - if: runner.os == 'Windows'
-        name: (Windows) Disable Flash-Attention build
-        run: |
-          echo "XFORMERS_DISABLE_FLASH_ATTN=1" >> ${GITHUB_ENV}
-
       - if: runner.os == 'Windows'
         name: (Windows) Setup Runner
         uses: ./.github/actions/setup-windows-runner