diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml new file mode 100644 index 000000000..a6baf8376 --- /dev/null +++ b/.github/workflows/build-wheels.yml @@ -0,0 +1,112 @@ +name: Build flash-attention Wheels for Windows + +on: + workflow_dispatch: + inputs: + version: + 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 (format: v2.3.4/v2.3.4.post1)' + default: 'v2.3.4' + required: true + type: string + +permissions: + contents: write + +jobs: + build_wheels: + name: Build wheels for Python ${{ matrix.pyver }}, CUDA ${{ matrix.cuda }}, and Torch ${{ matrix.torchver }} + runs-on: windows-2022 + strategy: + matrix: + pyver: ["3.8", "3.9", "3.10", "3.11", "3.12"] + cuda: ["12.2.2"] + torchver: ["2.2.2", "2.3.1"] + defaults: + run: + shell: pwsh + env: + CUDAVER: ${{ matrix.cuda }} + PCKGVER: ${{ inputs.version }} + + steps: + - uses: actions/checkout@v4 + with: + submodules: 'recursive' + + - uses: actions/setup-python@v4 + 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: + 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 = '${{ 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' $cudaVersion} + if (!(mamba list cuda)[-1].contains('cuda')) {throw 'CUDA Toolkit failed to install!'} + + 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') + + $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"+"torch${{ matrix.torchver }}cxx11abiFALSE-") + 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 }}