Merge branch 'main' into agent/conv-transpose3d-metal

Author: ternaus

.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/cuda/cuda_utils.h

@@ -50,6 +50,12 @@ class CudaHandle {
     }
   }
 
+  Handle release() {
+    Handle handle = handle_;
+    handle_ = nullptr;
+    return handle;
+  }
+
   operator Handle() const {
     return handle_;
   }

mlx/backend/cuda/device.cpp

@@ -462,6 +462,42 @@ bool CommandEncoder::needs_commit() {
 
 void CommandEncoder::commit() {
   nvtx3::scoped_range r("CommandEncoder::commit");
+  try {
+    commit_impl();
+  } catch (...) {
+    // Clear pending CUDA error first.
+    cudaGetLastError();
+    // Clear states.
+    clear_graph_state();
+    node_count_ = 0;
+    bytes_in_graph_ = 0;
+    // Clear graph.
+    try {
+      graph_.reset();
+    } catch (...) {
+      // Destroying could fail.
+      graph_.release();
+    }
+    try {
+      graph_ = CudaGraph(device_);
+    } catch (...) {
+      // Keep the original error.
+    }
+    // Re-throw the error.
+    throw;
+  }
+}
+
+void CommandEncoder::synchronize() {
+  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
+  auto p = std::make_shared<std::promise<void>>();
+  std::future<void> f = p->get_future();
+  add_completed_handler([p = std::move(p)]() { p->set_value(); });
+  commit();
+  f.wait();
+}
+
+void CommandEncoder::commit_impl() {
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
   }
@@ -520,13 +556,8 @@ void CommandEncoder::commit() {
     }
 
     // Reset state
-    from_nodes_.clear();
-    to_nodes_.clear();
-    graph_deps_key_.clear();
-    graph_nodes_key_.clear();
-    node_map_.clear();
+    clear_graph_state();
     graph_ = CudaGraph(device_);
-    is_graph_updatable_ = true;
   }
 
   // Put completion handlers in a batch.
@@ -535,13 +566,16 @@ void CommandEncoder::commit() {
   bytes_in_graph_ = 0;
 }
 
-void CommandEncoder::synchronize() {
-  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
-  auto p = std::make_shared<std::promise<void>>();
-  std::future<void> f = p->get_future();
-  add_completed_handler([p = std::move(p)]() { p->set_value(); });
-  commit();
-  f.wait();
+void CommandEncoder::clear_graph_state() {
+  from_nodes_.clear();
+  to_nodes_.clear();
+  graph_deps_key_.clear();
+  graph_nodes_key_.clear();
+  node_map_.clear();
+  active_deps_.clear();
+  active_outputs_.clear();
+  concurrent_nodes_.clear();
+  is_graph_updatable_ = true;
 }
 
 Device& device(int cuda_device) {

mlx/backend/cuda/device.h

@@ -138,6 +138,8 @@ class CommandEncoder {
     std::string id;
   };
 
+  void commit_impl();
+  void clear_graph_state();
   void insert_graph_dependencies(GraphNode node);
   void insert_graph_dependencies(std::vector<GraphNode> nodes);
 

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,