ci: refactor ARM64 self-hosted runner workflow configuration

- Rename workflow file from _build_linux_self_host.yml to _build_linux_arm_self_host.yml for clarity
- Update workflow_call inputs to use flexible container image configuration
- Remove use-container flag and replace with container-image parameter
- Consolidate ubuntu image handling with --platform linux/arm64 option
- Update test-arm-self-hosted.yml to reference the renamed workflow

.github/workflows/_build_linux_arm_self_host.yml

@@ -0,0 +1,150 @@
+# #########################################################
+# Build wheels with self-hosted runner
+# #########################################################
+
+name: "[Linux Arm64, self-hosted] Build wheels and upload to GitHub Releases"
+
+on:
+  workflow_call:
+    inputs:
+      flash-attn-version:
+        description: "Flash-Attention version"
+        required: true
+        type: string
+      python-version:
+        description: "Python version"
+        required: true
+        type: string
+      torch-version:
+        description: "PyTorch version"
+        required: true
+        type: string
+      cuda-version:
+        description: "CUDA version"
+        required: true
+        type: string
+      runner:
+        description: "Runner type"
+        required: false
+        type: string
+        default: "self-hosted"
+      is-upload:
+        description: "Whether to upload the release asset"
+        required: false
+        type: boolean
+        default: true
+      container-image:
+        description: "Container image"
+        required: false
+        type: string
+        default: "ubuntu:22.04"
+
+jobs:
+  build_wheels_self_hosted:
+    name: Build wheels and Upload (Linux x86_64, self-hosted runner)
+    runs-on: ${{ inputs.runner }}
+    container:
+      image: ${{ inputs.container-image }}
+      options: --platform linux/arm64
+    defaults:
+      run:
+        shell: bash
+    env:
+      DEBIAN_FRONTEND: noninteractive
+      TERM: xterm-256color
+    timeout-minutes: 2160
+
+    steps:
+      - name: Install tools
+        shell: bash
+        run: |
+          apt-get update && apt-get install -y --no-install-recommends \
+            curl \
+            ca-certificates \
+            sudo \
+            software-properties-common \
+            wget \
+            unzip \
+            zip \
+            git \
+            build-essential \
+            gcc \
+            g++ \
+            clang \
+            ninja-build \
+            keyboard-configuration \
+            time
+
+      - name: Install gh
+        shell: bash
+        run: |
+          sudo mkdir -p -m 755 /etc/apt/keyrings
+          out=$(mktemp)
+          wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg
+          cat $out | sudo tee /etc/apt/keyrings/githubcli-archive-keyring.gpg > /dev/null
+          sudo chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg
+          echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null
+          sudo apt update
+          sudo apt install gh -y
+
+      - uses: actions/checkout@v4
+
+      - name: Configure Git safe directory
+        shell: bash
+        run: |
+          git config --global --add safe.directory $(pwd)
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ inputs.python-version }}
+
+      - uses: mjun0812/setup-cuda@v1
+        with:
+          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 packaging psutil
+
+      - name: Build wheels
+        timeout-minutes: 2160
+        id: build_wheels
+        shell: bash
+        run: |
+          chmod +x build_linux.sh
+          ./build_linux.sh ${{ inputs.flash-attn-version }} ${{ inputs.python-version }} ${{ inputs.torch-version }} ${{ inputs.cuda-version }}
+          wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1))
+          echo "WHEEL_NAME=$wheel_name" >> $GITHUB_OUTPUT
+
+      - 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__)"
+
+      - name: Upload Release Asset
+        if: ${{ inputs.is-upload }}
+        shell: bash
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          tag_name=${{ github.ref_name }}
+          wheel_path="flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}"
+
+          # Check if the file exists
+          if [ ! -f "$wheel_path" ]; then
+            echo "Error: Wheel file not found at $wheel_path"
+            exit 1
+          fi
+
+          # Upload the release asset using GitHub CLI
+          gh release upload "$tag_name" "$wheel_path" --clobber
+
+      - name: Clean up
+        shell: bash
+        if: always()
+        run: |
+          rm -rf /opt/hostedtoolcache/Python

.github/workflows/test-arm-self-hosted.yml

@@ -9,7 +9,7 @@ jobs:
   # #########################################################
   build_wheels_self_hosted:
     name: Build wheels and Upload (Linux ARM64, self-hosted runner)
-    uses: ./.github/workflows/_build_linux_self_host.yml
+    uses: ./.github/workflows/_build_linux_arm_self_host.yml
     strategy:
       fail-fast: false
       matrix:
@@ -26,5 +26,4 @@ jobs:
       cuda-version: ${{ matrix.cuda-version }}
       is-upload: false
       runner: ${{ matrix.runner}}
-      use-container: true
-      container-image: "arm64v8/ubuntu:22.04"
+      container-image: "ubuntu:22.04"