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:37:53 -04:00
Code Issues Packages Projects Releases 2 Wiki Activity

fix: build for arm

Browse Source
This commit is contained in:
Junya Morioka
2025-12-04 19:23:55 +09:00
parent 4db0ed0589
commit 0dbd2193cf
2 changed files with 19 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test-arm.yml
+4 -1
View File
@@ -17,7 +17,7 @@ jobs:
fail-fast: false fail-fast: false
matrix: matrix:
flash-attn-version: ["2.8.0"] flash-attn-version: ["2.8.0"]
python-version: ["3.11", "3.11"] python-version: ["3.11"]
torch-version: ["2.7.1"] torch-version: ["2.7.1"]
# https://developer.nvidia.com/cuda-toolkit-archive # https://developer.nvidia.com/cuda-toolkit-archive
cuda-version: ["12.8.1"] cuda-version: ["12.8.1"]
@@ -51,6 +51,9 @@ jobs:
- name: Build wheels - name: Build wheels
id: build_wheels id: build_wheels
shell: bash shell: bash
env:
MAX_JOBS: 2
NVCC_THREADS: 1
run: | run: |
chmod +x build_linux.sh chmod +x build_linux.sh
./build_linux.sh ${{ matrix.flash-attn-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }} ./build_linux.sh ${{ matrix.flash-attn-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }}
build_linux.sh
+15 -12
View File
@@ -51,20 +51,23 @@ echo "System resources:"
echo " CPU threads: $NUM_THREADS" echo " CPU threads: $NUM_THREADS"
echo " RAM: ${RAM_GB}GB" echo " RAM: ${RAM_GB}GB"
# Calculate max product based on following constraints: # Determine MAX_JOBS and NVCC_THREADS based on system resources
# - MAX_JOBS x NVCC_THREADS <= NUM_THREADS if [[ -z "${MAX_JOBS:-}" && -z "${NVCC_THREADS:-}" ]]; then
# - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB # Calculate max product based on following constraints:
MAX_PRODUCT_CPU=$NUM_THREADS # - MAX_JOBS x NVCC_THREADS <= NUM_THREADS
MAX_PRODUCT_RAM=$(awk -v ram="$RAM_GB" 'BEGIN {print int(ram / 2.5)}') # - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB
MAX_PRODUCT=$((MAX_PRODUCT_CPU < MAX_PRODUCT_RAM ? MAX_PRODUCT_CPU : MAX_PRODUCT_RAM)) MAX_PRODUCT_CPU=$NUM_THREADS
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)) # Set MAX_JOBS = NVCC_THREADS = floor(sqrt(MAX_PRODUCT))
MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}') MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}')
NVCC_THREADS=$MAX_JOBS NVCC_THREADS=$MAX_JOBS
# Ensure minimum values of 1 # Ensure minimum values of 1
MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS)) MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS))
NVCC_THREADS=$((NVCC_THREADS < 1 ? 1 : NVCC_THREADS)) NVCC_THREADS=$((NVCC_THREADS < 1 ? 1 : NVCC_THREADS))
fi
echo "Build parallelism settings:" echo "Build parallelism settings:"
echo " MAX_JOBS: $MAX_JOBS" echo " MAX_JOBS: $MAX_JOBS"