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__
@@ -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