diff --git a/.github/actions/build-cuda-release/action.yml b/.github/actions/build-cuda-release/action.yml index b11b332f..f78b0c10 100644 --- a/.github/actions/build-cuda-release/action.yml +++ b/.github/actions/build-cuda-release/action.yml @@ -20,7 +20,7 @@ runs: run: | pip install auditwheel "build<=1.4.2" patchelf setuptools python setup.py clean --all - MLX_DISABLE_SM90A_KERNELS=1 MLX_BUILD_STAGE=2 python -m build -w + MLX_BUILD_STAGE=2 python -m build -w auditwheel repair dist/mlx_cuda*.whl \ --plat manylinux_2_35_${{ inputs.arch }} \ diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml index e7286d8a..bb397f52 100644 --- a/.github/actions/setup-linux/action.yml +++ b/.github/actions/setup-linux/action.yml @@ -14,6 +14,9 @@ inputs: description: 'Whether to enable ccache' required: false default: 'true' + ccache-key: + required: false + default: 'ccache' runs: using: "composite" @@ -33,7 +36,7 @@ runs: if: ${{ inputs.use-ccache == 'true' }} uses: hendrikmuhs/ccache-action@v1.2 with: - key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }} + key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }} max-size: 1GB # ccache-action bug: running "apt-get update" fails on large arm runner. update-package-index: false diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 09169e06..8266140d 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -85,20 +85,24 @@ jobs: build_cuda_release: if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22-large + strategy: + matrix: + arch: ['x86_64', 'aarch64'] + toolkit: ['cuda-12.9', 'cuda-13.0'] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }} steps: - uses: actions/checkout@v6 - uses: ./.github/actions/setup-linux with: - toolkit: 'cuda-12.9' + toolkit: ${{ matrix.toolkit }} + ccache-key: 'ccache-release' - name: Build Python package uses: ./.github/actions/build-cuda-release with: - toolkit: 'cuda-12.9' - arch: 'x86_64' + arch: ${{ matrix.arch }} - name: Upload artifacts uses: actions/upload-artifact@v7 with: - name: mlx-cuda + name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }} path: wheelhouse/mlx_cuda_*.whl retention-days: 7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 66f949e1..50530bf9 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -146,7 +146,7 @@ jobs: - uses: ./.github/actions/setup-linux with: toolkit: ${{ matrix.toolkit }} - use-ccache: false + ccache-key: 'ccache-release' - name: Build Python package uses: ./.github/actions/build-cuda-release with: diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 37af7584..bc772db5 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -168,9 +168,8 @@ set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") # Skip Hopper-only kernels when not building for sm90a. -if(NOT DEFINED ENV{MLX_DISABLE_SM90A_KERNELS} - AND (("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST - MLX_CUDA_ARCHITECTURES))) +if(("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST + MLX_CUDA_ARCHITECTURES)) target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED) endif()