From 0dbd2193cfb4c3262a59e5c3e662f0f53f8eb343 Mon Sep 17 00:00:00 2001
From: Junya Morioka <mjun@mjunya.com>
Date: Thu, 4 Dec 2025 19:23:55 +0900
Subject: [PATCH] fix: build for arm

---
 .github/workflows/test-arm.yml |  5 ++++-
 build_linux.sh                 | 27 +++++++++++++++------------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/test-arm.yml b/.github/workflows/test-arm.yml
index 9a3ae3d..9aa3b07 100644
--- a/.github/workflows/test-arm.yml
+++ b/.github/workflows/test-arm.yml
@@ -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 }}
diff --git a/build_linux.sh b/build_linux.sh
index b528bd4..31e8d70 100755
--- a/build_linux.sh
+++ b/build_linux.sh
@@ -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"