From 07fdda3d9c72a2c14b443bdf13e29c6cbafea41e Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Mon, 17 Nov 2025 10:17:42 +0900
Subject: [PATCH] Merge build-cuda and build-linux actions

---
 .github/actions/build-cuda/action.yml  | 26 ------------------------
 .github/actions/build-linux/action.yml | 28 ++++++++++++++++++++------
 .github/workflows/pull_request.yml     |  2 +-
 3 files changed, 23 insertions(+), 33 deletions(-)
 delete mode 100644 .github/actions/build-cuda/action.yml

diff --git a/.github/actions/build-cuda/action.yml b/.github/actions/build-cuda/action.yml
deleted file mode 100644
index 2a3b39883..000000000
--- a/.github/actions/build-cuda/action.yml
+++ /dev/null
@@ -1,26 +0,0 @@
-name: 'Build and Test with CUDA'
-description: 'Build and test MLX with CUDA'
-
-inputs:
-  toolkit:
-    description: 'The CUDA toolkit'
-    required: true
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      shell: bash
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
-      run: pip install --no-build-isolation -e ".[dev]" -v
-
-    - name: Build CPP only
-      shell: bash
-      run: |
-        cmake . -B build \
-          -DMLX_BUILD_CUDA=ON \
-          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc \
-          -DCMAKE_BUILD_TYPE=DEBUG
-        cmake --build build -j $(nproc)
diff --git a/.github/actions/build-linux/action.yml b/.github/actions/build-linux/action.yml
index 6273ab8de..76dc86f6c 100644
--- a/.github/actions/build-linux/action.yml
+++ b/.github/actions/build-linux/action.yml
@@ -1,13 +1,30 @@
 name: 'Build and Test on Linux'
-description: 'Build and test MLX on Linux'
+
+inputs:
+  toolkit:
+    description: 'The toolkit to build with'
+    required: false
+    default: 'cpu'
 
 runs:
   using: "composite"
   steps:
+    - name: Set cmake args
+      shell: bash
+      env:
+        CMAKE_ARGS_CPU: >
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DMLX_BUILD_CUDA=OFF
+        CMAKE_ARGS_CUDA: >
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
+          -DMLX_BUILD_CUDA=ON
+      run: |
+        echo CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}" > $GITHUB_ENV
+
     - name: Install Python package
       shell: sh
       env:
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
         DEBUG: 1
       run: pip install --no-build-isolation -e ".[dev]" -v
 
@@ -16,10 +33,9 @@ runs:
       run: |
         pip install typing_extensions
         python setup.py generate_stubs
-    
+
     - name: Build CPP only
       shell: bash
       run: |
-        mkdir -p build && cd build
-        cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-        make -j $(nproc)
+        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ env.CMAKE_ARGS }}
+        cmake --build build -j $(nproc)
diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index db63a9ad1..fd8eeeb73 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -66,7 +66,7 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/build-cuda
+      - uses: ./.github/actions/build-linux
         with:
           toolkit: ${{ matrix.toolkit }}
       - uses: ./.github/actions/test-linux