yindo/flash-attention-prebuild-wheels-rocm
1
0
Fork 0
mirror of https://github.com/BillyOutlast/flash-attention-prebuild-wheels-rocm.git synced 2026-06-30 23:57: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
matrix:
flash-attn-version: ["2.8.0"]
python-version: ["3.11", "3.11"]
python-version: ["3.11"]
torch-version: ["2.7.1"]
# https://developer.nvidia.com/cuda-toolkit-archive
cuda-version: ["12.8.1"]
@@ -51,6 +51,9 @@ jobs:
- name: Build wheels
id: build_wheels
shell: bash
env:
MAX_JOBS: 2
NVCC_THREADS: 1
run: |
chmod +x build_linux.sh
./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 " RAM: ${RAM_GB}GB"
# Calculate max product based on following constraints:
# - MAX_JOBS x NVCC_THREADS <= NUM_THREADS
# - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB
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))
# Determine MAX_JOBS and NVCC_THREADS based on system resources
if [[ -z "${MAX_JOBS:-}" && -z "${NVCC_THREADS:-}" ]]; then
# Calculate max product based on following constraints:
# - MAX_JOBS x NVCC_THREADS <= NUM_THREADS
# - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB
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))
MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}')
NVCC_THREADS=$MAX_JOBS
# Set MAX_JOBS = NVCC_THREADS = floor(sqrt(MAX_PRODUCT))
MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}')
NVCC_THREADS=$MAX_JOBS
# Ensure minimum values of 1
MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS))
NVCC_THREADS=$((NVCC_THREADS < 1 ? 1 : NVCC_THREADS))
# Ensure minimum values of 1
MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS))
NVCC_THREADS=$((NVCC_THREADS < 1 ? 1 : NVCC_THREADS))
fi
echo "Build parallelism settings:"
echo " MAX_JOBS: $MAX_JOBS"