Merge remote-tracking branch 'bdashore3/main'

Author: oobabooga

.github/workflows/build-wheels.yml

@@ -4,14 +4,14 @@ on:
   workflow_dispatch:
     inputs:
       version:
-        description: 'Version tag of flash-attention to build: v2.3.4'
+        description: 'Version tag of flash-attention to build: (format: v2.3.4/v2.3.4.post1)'
         default: 'v2.3.4'
         required: true
         type: string
   workflow_call:
     inputs:
       version:
-        description: 'Version tag of flash-attention to build: v2.3.4'
+        description: 'Version tag of flash-attention to build (format: v2.3.4/v2.3.4.post1)'
         default: 'v2.3.4'
         required: true
         type: string
@@ -22,12 +22,12 @@ permissions:
 jobs:
   build_wheels:
     name: Build wheels for Python ${{ matrix.pyver }}, CUDA ${{ matrix.cuda }}, and Torch ${{ matrix.torchver }}
-    runs-on: windows-latest
+    runs-on: windows-2022
     strategy:
       matrix:
         pyver: ["3.10", "3.11"]
         cuda: ["12.2.2"]
-        torchver: ["2.1.2", "2.2.0"]
+        torchver: ["2.2.2", "2.3.1"]
     defaults:
       run:
         shell: pwsh
@@ -44,6 +44,10 @@ jobs:
         with:
           python-version: ${{ matrix.pyver }}
 
+      - name: Install VS2022 BuildTools 17.9.7
+        run: choco install -y visualstudio2022buildtools --version=117.9.7.0 --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 --installChannelUri https://aka.ms/vs/17/release/180911598_-255012421/channel"
+        if: runner.os == 'Windows'
+
       - name: Setup Mamba
         uses: conda-incubator/setup-miniconda@v2.2.0
         with:
@@ -57,20 +61,25 @@ jobs:
 
       - name: Install Dependencies
         run: |
-          $cudaVersion = $env:CUDAVER
-          $cudaVersionPytorch = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','')
-          $cudaChannels = ''
-          $cudaNum = [int]$cudaVersion.substring($cudaVersion.LastIndexOf('.')+1)
-          while ($cudaNum -ge 0) { $cudaChannels += '-c nvidia/label/cuda-' + $cudaVersion.Remove($cudaVersion.LastIndexOf('.')+1) + $cudaNum + ' '; $cudaNum-- }
-          mamba install -y 'cuda' $cudaChannels.TrimEnd().Split()
+          $cudaVersion = '${{ matrix.cuda }}'
+          $cudaVersionPytorch = '${{ matrix.cuda }}'.Remove('${{ matrix.cuda }}'.LastIndexOf('.')).Replace('.','')
+
+          mamba install -y -c nvidia/label/cuda-$cudaVersion cuda-toolkit cuda-runtime
           if (!(mamba list cuda)[-1].contains('cuda')) {sleep -s 10; mamba install -y 'cuda' $cudaChannels.TrimEnd().Split()}
           if (!(mamba list cuda)[-1].contains('cuda')) {throw 'CUDA Toolkit failed to install!'}
           
-          python -m pip install --upgrade build setuptools wheel packaging ninja psutil torch==${{ matrix.torchver }} --extra-index-url "https://download.pytorch.org/whl/cu121"
+          python -m pip install --upgrade build setuptools wheel packaging ninja torch==${{ matrix.torchver }} psutil --extra-index-url "https://download.pytorch.org/whl/cu121"
 
       - name: Build Wheel
         id: build-wheel
         run: |
+          # --- Spawn the VS shell
+          if ($IsWindows) {
+            Import-Module 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\Microsoft.VisualStudio.DevShell.dll'
+            Enter-VsDevShell -VsInstallPath 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools' -DevCmdArguments '-arch=x64 -host_arch=x64'
+            $env:DISTUTILS_USE_SDK=1
+          }
+
           $cudaVersion = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','')
           $packageVersion = $env:PCKGVER.TrimStart('v')