Use the correct repo for arm64 CUDA

.github/actions/build-linux/action.yml

@@ -9,9 +9,10 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Set cmake args
-      shell: bash
+    - name: Install Python package
+      shell: sh
       env:
+        DEBUG: 1
         CMAKE_ARGS_CPU: >-
           -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
           -DMLX_BUILD_CUDA=OFF
@@ -20,22 +21,16 @@ runs:
           -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
           -DMLX_BUILD_CUDA=ON
       run: |
-        CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}"
-        # FIXME: Unable to build tests without libnccl on arm64.
+        export CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}"
         if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
+          # Can not build tests when the built executables can not run.
           CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
         fi
+        pip install --no-build-isolation -e ".[dev]" -v
+        # Pass the CMAKE_ARGS to following steps.
         echo CMAKE_ARGS="$CMAKE_ARGS" > $GITHUB_ENV
 
-    - name: Install Python package
-      shell: sh
-      env:
-        DEBUG: 1
-      run: pip install --no-build-isolation -e ".[dev]" -v
-
     - name: Generate package stubs
-      # FIXME: Unable to generate stubs without libnccl on arm64.
-      if: ${{ !(startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64') }}
       shell: sh
       run: |
         pip install typing_extensions

.github/actions/setup-linux/action.yml

@@ -60,17 +60,13 @@ runs:
             "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
           }
       run: |
-        export ARCH=${{ runner.arch == 'arm64' && 'arm64' || 'x86_64' }}
+        export ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
         wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
         sudo dpkg -i cuda-keyring_1.1-1_all.deb
         sudo apt-get update
-        sudo apt-get install -y ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
-
-    - name: Install NCCL
-      # For some reason NVIDIA did not provide arm64 package for libnccl2.
-      if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch != 'arm64' }}
-      shell: bash
-      run: sudo apt-get install -y libnccl2 libnccl-dev
+        sudo apt-get install -y \
+            libnccl2 libnccl-dev \
+            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
 
     - name: CUDA packages and driver report
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}