Fix quantized kernels in JIT build (#4372)

Co-authored-by: Cheng <git@zcbenz.com>

Author: dwijenpatel

.github/actions/setup/action.yml

@@ -229,14 +229,7 @@ runs:
           else
             cmakeArgs+=("-DMLX_BUILD_METAL=ON")
             if ${{ inputs.toolkit == 'jit' }} ; then
-              cmakeArgs+=(
-                "-DBUILD_SHARED_LIBS=ON"
-                "-DCMAKE_BUILD_TYPE=MinSizeRel"
-                "-DMLX_BUILD_CPU=OFF"
-                "-DMLX_BUILD_SAFETENSORS=OFF"
-                "-DMLX_BUILD_GGUF=OFF"
-                "-DMLX_METAL_JIT=ON"
-              )
+              cmakeArgs+=("-DMLX_METAL_JIT=ON")
             fi
           fi
         fi

.github/workflows/build_and_test.yml

@@ -105,7 +105,7 @@ jobs:
         with:
           toolkit: 'cpu'
       - uses: actions/upload-artifact@v7
-        if: matrix.toolkit == 'metal'
+        if: matrix.toolkit != 'cpu'
         with:
           name: mlx-${{ matrix.toolkit }}-macos${{ matrix.macos-target }}
           path: |
@@ -115,24 +115,28 @@ jobs:
           if-no-files-found: error
 
   mac_test:
-    name: Test macOS
+    name: Test macOS (${{ matrix.toolkit }})
     if: github.repository == 'ml-explore/mlx'
     runs-on: [self-hosted, macos]
     needs: mac_build
+    strategy:
+      fail-fast: false
+      matrix:
+        toolkit: ['metal', 'jit']
     steps:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup
         with:
-          toolkit: 'metal'
+          toolkit: ${{ matrix.toolkit }}
           use-ccache: false
       - uses: actions/download-artifact@v8
         with:
           path: artifact
-          pattern: mlx-metal-*
+          pattern: mlx-${{ matrix.toolkit }}-*
       - run: ls -lhR artifact
       - uses: ./.github/actions/test-macos
         with:
-          toolkit: 'metal'
+          toolkit: ${{ matrix.toolkit }}
 
   build_documentation:
     name: Build Documentation

mlx/backend/metal/kernels/fp_quantized.h

@@ -1232,7 +1232,8 @@ template <
     int group_size,
     int bits,
     bool batched,
-    bool has_global_scale = false>
+    bool has_global_scale = false,
+    int results_per_simdgroup = 4>
 [[kernel]] void fp_qmv(
     const device uint32_t* w,
     const device uint8_t* scales,

mlx/backend/metal/kernels/quantized.h

@@ -1603,7 +1603,8 @@ template <
     int group_size,
     int bits,
     bool batched,
-    bool has_global_scale = false>
+    bool has_global_scale = false,
+    int results_per_simdgroup = 4>
 [[kernel]] void affine_qmv_fast(
     const device uint32_t* w [[buffer(0)]],
     const device T* scales [[buffer(1)]],
@@ -1660,7 +1661,8 @@ template <
     int group_size,
     const int bits,
     bool batched,
-    bool has_global_scale = false>
+    bool has_global_scale = false,
+    int results_per_simdgroup = 4>
 [[kernel]] void affine_qmv(
     const device uint32_t* w [[buffer(0)]],
     const device T* scales [[buffer(1)]],

mlx/backend/metal/quantized.cpp

@@ -972,7 +972,7 @@ void gather_qmm_nax(
     kernel = get_qmm_nax_kernel_wrapped(
         d,
         kname,
-        "gather_qmm_t_nax_",
+        "gather_qmm_t_nax",
         mode,
         type_string,
         group_size,
@@ -987,7 +987,7 @@ void gather_qmm_nax(
     kernel = get_qmm_nax_kernel_wrapped(
         d,
         kname,
-        "gather_qmm_n_nax_",
+        "gather_qmm_n_nax",
         mode,
         type_string,
         group_size,