From b626c528d68e8660f6bc1a9ee12b2afdb04549ff Mon Sep 17 00:00:00 2001
From: jllllll <3887729+jllllll@users.noreply.github.com>
Date: Mon, 20 Nov 2023 08:46:39 -0600
Subject: [PATCH] Add Windows workflows

---
 .github/workflows/batch-build-wheels.yml |  45 ++++++++++
 .github/workflows/build-wheels.yml       | 104 +++++++++++++++++++++++
 2 files changed, 149 insertions(+)
 create mode 100644 .github/workflows/batch-build-wheels.yml
 create mode 100644 .github/workflows/build-wheels.yml

diff --git a/.github/workflows/batch-build-wheels.yml b/.github/workflows/batch-build-wheels.yml
new file mode 100644
index 000000000..76748a02f
--- /dev/null
+++ b/.github/workflows/batch-build-wheels.yml
@@ -0,0 +1,45 @@
+name: Batch Build flash-attention Wheels for Windows
+
+on:
+  workflow_dispatch:
+    inputs:
+      versions:
+        description: 'Version tag of flash-attention to build: v2.3.4'
+        default: 'v2.3.2,v2.3.3,v2.3.4'
+        required: true
+        type: string
+
+permissions:
+  contents: write
+
+jobs:
+  define_matrix:
+    name: Define Workflow Matrix
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    defaults:
+      run:
+        shell: pwsh
+    env:
+      PCKGVERS: ${{ inputs.versions }}
+        
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Define Job Output
+        id: set-matrix
+        run: |
+          $x = ConvertTo-Json $env:PCKGVERS.Split(',').Trim() -Compress
+          Write-Output ('matrix=' + $x) >> $env:GITHUB_OUTPUT
+          
+  run_workflows:
+    name: Build ${{ matrix.version }} Wheels
+    needs: define_matrix
+    strategy:
+      max-parallel: 1
+      matrix:
+        version: ${{ fromJSON(needs.define_matrix.outputs.matrix) }}
+    uses: ./.github/workflows/build-wheels.yml
+    with:
+      version: ${{ matrix.version }}
diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
new file mode 100644
index 000000000..ae19cb11d
--- /dev/null
+++ b/.github/workflows/build-wheels.yml
@@ -0,0 +1,104 @@
+name: Build flash-attention Wheels for Windows
+
+on:
+  workflow_dispatch:
+    inputs:
+      version:
+        description: 'Version tag of flash-attention to build: v2.3.4'
+        default: 'v2.3.4'
+        required: true
+        type: string
+  workflow_call:
+    inputs:
+      version:
+        description: 'Version tag of flash-attention to build: v2.3.4'
+        default: 'v2.3.4'
+        required: true
+        type: string
+
+permissions:
+  contents: write
+
+jobs:
+  build_wheels:
+    name: Build wheels for Python ${{ matrix.pyver }} and CUDA ${{ matrix.cuda }}
+    runs-on: windows-latest
+    strategy:
+      matrix:
+        pyver: ["3.8", "3.9", "3.10", "3.11"]
+        cuda: ["12.1.1"]
+    defaults:
+      run:
+        shell: pwsh
+    env:
+        CUDAVER: ${{ matrix.cuda }}
+        PCKGVER: ${{ inputs.version }}
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          repository: 'Dao-AILab/flash-attention'
+          ref: ${{ inputs.version }}
+          submodules: 'recursive'
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.pyver }}
+
+      - name: Setup Mamba
+        uses: conda-incubator/setup-miniconda@v2.2.0
+        with:
+          activate-environment: "build"
+          python-version: ${{ matrix.pyver }}
+          miniforge-variant: Mambaforge
+          miniforge-version: latest
+          use-mamba: true
+          add-pip-as-python-dependency: true
+          auto-activate-base: false
+
+      - 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()
+          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 torch==2.1.0 --extra-index-url "https://download.pytorch.org/whl/cu$cudaVersionPytorch"
+
+      - name: Build Wheel
+        id: build-wheel
+        run: |
+          $cudaVersion = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','')
+          $packageVersion = $env:PCKGVER.TrimStart('v')
+          
+          $env:CUDA_PATH = $env:CONDA_PREFIX
+          $env:CUDA_HOME = $env:CONDA_PREFIX
+          
+          $env:MAX_JOBS = '1'
+          $env:FLASH_ATTENTION_FORCE_BUILD = 'TRUE'
+          $env:FLASH_ATTENTION_FORCE_SINGLE_THREAD = 'TRUE'
+          
+          python -m build -n --wheel
+          
+          $wheel = (gi '.\dist\*.whl')[0]
+          $wheelname = $wheel.name.replace("flash_attn-$packageVersion-","flash_attn-$packageVersion+cu$cudaVersion"+"torch2.1cxx11abiFALSE-")
+          Move-Item $wheel.fullname ".\dist\$wheelname"
+
+      - uses: actions/upload-artifact@v3
+        with:
+          name: 'windows-wheels'
+          path: ./dist/*.whl
+          
+      - name: Upload files to a GitHub release
+        uses: svenstaro/upload-release-action@2.6.1
+        continue-on-error: true
+        with:
+          file: ./dist/*.whl
+          tag: ${{ inputs.version }}
+          file_glob: true
+          overwrite: true
+          release_name: ${{ inputs.version }}