yindo/flash-attention-prebuild-wheels-rocm
1
0
Fork 0
mirror of https://github.com/BillyOutlast/flash-attention-prebuild-wheels-rocm.git synced 2026-07-01 01:27:54 -04:00
Code Issues Packages Projects Releases 2 Wiki Activity

feat: Add new test for arm building and fix build prameters

Browse Source
This commit is contained in:
Junya Morioka
2025-12-04 19:07:32 +09:00
parent 3f80135940
commit 0bce0110f7
10 changed files with 87 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build_linux.yml → .github/workflows/_build_linux.yml
+5 -21
View File
@@ -1,5 +1,5 @@
# #########################################################
# Build wheels with GitHub hosted runner
# Build wheels with GitHub-hosted runner
# #########################################################
name: "[Linux x86_64] Build wheels and upload to GitHub Releases"
@@ -44,38 +44,21 @@ jobs:
sudo rm -rf /opt/hostedtoolcache/CodeQL
df -h
- name: Set Swap Space
uses: pierotofy/set-swap-space@master
with:
swap-size-gb: 48
- uses: actions/setup-python@v5
with:
python-version: ${{ inputs.python-version }}
- uses: Jimver/cuda-toolkit@v0.2.29
- uses: mjun0812/setup-cuda@v1
with:
cuda: ${{ inputs.cuda-version }}
sub-packages: '["nvcc", "toolkit"]'
method: "network"
use-github-cache: false
use-local-cache: false
version: ${{ inputs.cuda-version }}
- name: Install build dependencies
run: |
sudo apt install -y ninja-build clang
pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil
- name: Set environment variables
run: |
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
pip install -U pip setuptools==75.8.0 wheel packaging psutil
- name: Build wheels
id: build_wheels
env:
MAX_JOBS: 2
NVCC_THREADS: 2
run: |
chmod +x build_linux.sh
./build_linux.sh ${{ inputs.flash-attn-version }} ${{ inputs.python-version }} ${{ inputs.torch-version }} ${{ inputs.cuda-version }}
@@ -84,6 +67,7 @@ jobs:
- name: Install Test
run: |
pip uninstall -y flash-attn > /dev/null 2>&1
pip install --no-cache-dir flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
python -c "import flash_attn; print(flash_attn.__version__)"
.github/workflows/build_linux_self_host.yml → .github/workflows/_build_linux_self_host.yml
+3 -16
View File
@@ -80,33 +80,20 @@ jobs:
with:
python-version: ${{ inputs.python-version }}
- uses: Jimver/cuda-toolkit@v0.2.29
- uses: mjun0812/setup-cuda@v1
with:
cuda: ${{ inputs.cuda-version }}
sub-packages: '["nvcc", "toolkit"]'
method: "network"
use-github-cache: false
use-local-cache: false
version: ${{ inputs.cuda-version }}
- name: Install build dependencies
shell: bash
run: |
sudo apt install -y ninja-build clang
pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil
- name: Set environment variables
shell: bash
run: |
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
pip install -U pip setuptools==75.8.0 wheel packaging psutil
- name: Build wheels
timeout-minutes: 2160
id: build_wheels
shell: bash
env:
MAX_JOBS: 3
NVCC_THREADS: 3
run: |
chmod +x build_linux.sh
./build_linux.sh ${{ inputs.flash-attn-version }} ${{ inputs.python-version }} ${{ inputs.torch-version }} ${{ inputs.cuda-version }}
.github/workflows/build_windows.yml → .github/workflows/_build_windows.yml
+2 -5
View File
@@ -42,12 +42,9 @@ jobs:
with:
python-version: ${{ inputs.python-version }}
- uses: Jimver/cuda-toolkit@v0.2.29
- uses: mjun0812/setup-cuda@v1
with:
cuda: ${{ inputs.cuda-version }}
method: "network"
use-github-cache: false
use-local-cache: false
version: ${{ inputs.cuda-version }}
- name: Install Visual Studio BuildTools (C++/CMake + Windows SDK)
shell: pwsh
.github/workflows/build_windows_code_build.yml → .github/workflows/_build_windows_code_build.yml
+2 -5
View File
@@ -43,12 +43,9 @@ jobs:
with:
python-version: ${{ inputs.python-version }}
- uses: Jimver/cuda-toolkit@v0.2.29
- uses: mjun0812/setup-cuda@v1
with:
cuda: ${{ inputs.cuda-version }}
method: "network"
use-github-cache: false
use-local-cache: false
version: ${{ inputs.cuda-version }}
- name: Install Visual Studio BuildTools (C++/CMake + Windows SDK)
shell: pwsh
.github/workflows/build.yml
+4 -4
View File
@@ -47,7 +47,7 @@ jobs:
torch-version: ${{ fromjson(needs.create_matrix.outputs.matrix).linux.torch-version }}
cuda-version: ${{ fromjson(needs.create_matrix.outputs.matrix).linux.cuda-version }}
exclude: ${{ fromjson(needs.create_matrix.outputs.matrix).exclude }}
uses: ./.github/workflows/build_linux.yml
uses: ./.github/workflows/_build_linux.yml
with:
flash-attn-version: ${{ matrix.flash-attn-version }}
python-version: ${{ matrix.python-version }}
@@ -67,7 +67,7 @@ jobs:
torch-version: ${{ fromjson(needs.create_matrix.outputs.matrix).linux_self_hosted.torch-version }}
cuda-version: ${{ fromjson(needs.create_matrix.outputs.matrix).linux_self_hosted.cuda-version }}
exclude: ${{ fromjson(needs.create_matrix.outputs.matrix).exclude }}
uses: ./.github/workflows/build_linux_self_host.yml
uses: ./.github/workflows/_build_linux_self_host.yml
with:
flash-attn-version: ${{ matrix.flash-attn-version }}
python-version: ${{ matrix.python-version }}
@@ -90,7 +90,7 @@ jobs:
torch-version: ${{ fromjson(needs.create_matrix.outputs.matrix).windows.torch-version }}
cuda-version: ${{ fromjson(needs.create_matrix.outputs.matrix).windows.cuda-version }}
exclude: ${{ fromjson(needs.create_matrix.outputs.matrix).exclude }}
uses: ./.github/workflows/build_windows.yml
uses: ./.github/workflows/_build_windows.yml
with:
flash-attn-version: ${{ matrix.flash-attn-version }}
python-version: ${{ matrix.python-version }}
@@ -110,7 +110,7 @@ jobs:
torch-version: ${{ fromjson(needs.create_matrix.outputs.matrix).windows_code_build.torch-version }}
cuda-version: ${{ fromjson(needs.create_matrix.outputs.matrix).windows_code_build.cuda-version }}
exclude: ${{ fromjson(needs.create_matrix.outputs.matrix).exclude }}
uses: ./.github/workflows/build_windows_code_build.yml
uses: ./.github/workflows/_build_windows_code_build.yml
with:
flash-attn-version: ${{ matrix.flash-attn-version }}
python-version: ${{ matrix.python-version }}
.github/workflows/test-arm.yml
+65
View File
@@ -0,0 +1,65 @@
name: Test ARM Linux build
on:
workflow_dispatch:
jobs:
# #########################################################
# Build wheels with self-hosted runner
# #########################################################
build_wheels_self_hosted:
name: Build wheels and Upload (Linux arm64)
runs-on: ubuntu-22.04-arm
env:
DEBIAN_FRONTEND: noninteractive
TERM: xterm-256color
strategy:
fail-fast: false
matrix:
flash-attn-version: ["2.8.0"]
python-version: ["3.11", "3.11"]
torch-version: ["2.7.1"]
# https://developer.nvidia.com/cuda-toolkit-archive
cuda-version: ["12.8.1"]
steps:
- uses: actions/checkout@v4
- name: Maximize build space
run: |
df -h
echo "-----------------------------"
sudo rm -rf /usr/share/dotnet
sudo rm -rf /usr/local/lib/android
sudo rm -rf /opt/ghc
sudo rm -rf /opt/hostedtoolcache/CodeQL
df -h
- uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- uses: mjun0812/setup-cuda@v1
with:
version: "${{ matrix.cuda-version }}"
- name: Install build dependencies
shell: bash
run: |
sudo apt install -y ninja-build clang
pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil
- name: Build wheels
id: build_wheels
shell: bash
run: |
chmod +x build_linux.sh
./build_linux.sh ${{ matrix.flash-attn-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }}
wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1))
echo "wheel_name=$wheel_name" >> $GITHUB_ENV
- name: Install Test
shell: bash
run: |
pip uninstall -y flash-attn > /dev/null 2>&1
pip install --no-cache-dir flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
python -c "import flash_attn; print(flash_attn.__version__)"
.github/workflows/test-code-build.yml
+1 -1
View File
@@ -1,4 +1,4 @@
name: "Test Code build"
name: Test Code build
on:
workflow_dispatch:
.github/workflows/test-linux.yml
+1 -2
View File
@@ -1,4 +1,4 @@
name: Build wheels and upload to GitHub Releases
name: Test GHA runner environment
on:
workflow_dispatch:
@@ -16,4 +16,3 @@ jobs:
- shell: bash
run: |
python -c "import platform, os; print(platform.libc_ver())"
.github/workflows/test-self-hosted-linux.yml
+1 -1
View File
@@ -23,7 +23,7 @@ jobs:
fail-fast: false
matrix:
flash-attn-version: ["2.8.0"]
python-version: ["3.11", "3.12"]
python-version: ["3.11", "3.11"]
torch-version: ["2.7.1"]
# https://developer.nvidia.com/cuda-toolkit-archive
cuda-version: ["12.8.1"]
build_linux.sh
+3 -4
View File
@@ -51,15 +51,14 @@ echo "System resources:"
echo " CPU threads: $NUM_THREADS"
echo " RAM: ${RAM_GB}GB"
# Calculate max product based on constraints:
# Calculate max product based on following constraints:
# - MAX_JOBS x NVCC_THREADS <= NUM_THREADS
# - 4GB x MAX_JOBS x NVCC_THREADS <= RAM_GB
# - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB
MAX_PRODUCT_CPU=$NUM_THREADS
MAX_PRODUCT_RAM=$((RAM_GB / 4))
MAX_PRODUCT_RAM=$(awk -v ram="$RAM_GB" 'BEGIN {print int(ram / 2.5)}')
MAX_PRODUCT=$((MAX_PRODUCT_CPU < MAX_PRODUCT_RAM ? MAX_PRODUCT_CPU : MAX_PRODUCT_RAM))
# Set MAX_JOBS = NVCC_THREADS = floor(sqrt(MAX_PRODUCT))
# This balances parallelism across both dimensions
MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}')
NVCC_THREADS=$MAX_JOBS