Fix the sig_handler check

* redesign for faster cpu/gpu synch

* load + more async CPU

* use command encoder API and move more ops to use it

* make fence back-end generic + CPU only fence

* faster build

* fix async eval

* fixes + handle temporaries

* fix / improve cpu conv

* remove unused status, fix siblings

* fix extensions

* fix

* fix no cpu build

* format

* comments

* fix perf regression, remove unecessary abort

* fix events, task limit cpu

* fix waiting

* fix donation / temporaries in normalization

CMakeLists.txt

@@ -1,6 +1,23 @@
 cmake_minimum_required(VERSION 3.25)
 
-project(mlx LANGUAGES C CXX)
+if(NOT MLX_VERSION)
+  file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
+  string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_major ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_minor ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}")
+  set(_patch ${CMAKE_MATCH_1})
+  set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}")
+else()
+  string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION
+                       ${MLX_VERSION})
+endif()
+
+project(
+  mlx
+  LANGUAGES C CXX
+  VERSION ${MLX_PROJECT_VERSION})
 
 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
@@ -24,9 +41,6 @@ option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
-if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.23.1)
-endif()
 add_compile_definitions("MLX_VERSION=${MLX_VERSION}")
 
 # --------------------- Processor tests -------------------------
@@ -218,6 +232,14 @@ if(MPI_FOUND)
   endif()
 endif()
 
+message(STATUS "Downloading json")
+FetchContent_Declare(
+  json
+  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
+FetchContent_MakeAvailable(json)
+target_include_directories(
+  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
+
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(

benchmarks/python/layer_norm_bench.py

@@ -10,7 +10,12 @@ def layer_norm(x, w, b, eps):
     x = x.astype(mx.float32)
     mu = mx.mean(x, -1, keepdims=True)
     v = mx.var(x, -1, keepdims=True)
-    return (x - mu) * mx.rsqrt(v + eps) * w + b
+    y = (x - mu) * mx.rsqrt(v + eps)
+    if w is not None:
+        y = y * w
+    if b is not None:
+        y = y + b
+    return y
 
 
 def time_layer_norm():
@@ -36,6 +41,28 @@ def time_layer_norm():
     time_fn(layer_norm_loop, mx.compile(g1), x, w, b)
     time_fn(layer_norm_loop, mx.compile(g2), x, w, b)
 
+    f1 = lambda x, y: (layer_norm(x, None, None, 1e-5) * y).sum()
+    f2 = lambda x, y: (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum()
+    g1 = mx.grad(f1, argnums=(0,))
+    g2 = mx.grad(f2, argnums=(0,))
+
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    b = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    mx.eval(x, w, b, y)
+
+    def layer_norm_loop(g, x):
+        gx = x
+        for _ in range(32):
+            gx = g(gx, y)
+        return gx
+
+    time_fn(layer_norm_loop, g1, x)
+    time_fn(layer_norm_loop, g2, x)
+    time_fn(layer_norm_loop, mx.compile(g1), x)
+    time_fn(layer_norm_loop, mx.compile(g2), x)
+
 
 if __name__ == "__main__":
     time_layer_norm()

benchmarks/python/rms_norm_bench.py

@@ -9,7 +9,10 @@ def rms_norm(x, w, eps):
     ot = x.dtype
     x = x.astype(mx.float32)
     n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
-    return (x * n).astype(ot) * w
+    y = (x * n).astype(ot)
+    if w is not None:
+        y = y * w
+    return y
 
 
 def time_rms_norm():
@@ -34,6 +37,27 @@ def time_rms_norm():
     time_fn(rms_norm_loop, mx.compile(g1), x, w)
     time_fn(rms_norm_loop, mx.compile(g2), x, w)
 
+    f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum()
+    f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum()
+    g1 = mx.grad(f1, argnums=(0,))
+    g2 = mx.grad(f2, argnums=(0,))
+
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    mx.eval(x, w, y)
+
+    def rms_norm_loop(g, x):
+        gx = x
+        for _ in range(32):
+            gx = g(gx, y)
+        return gx
+
+    time_fn(rms_norm_loop, g1, x)
+    time_fn(rms_norm_loop, g2, x)
+    time_fn(rms_norm_loop, mx.compile(g1), x)
+    time_fn(rms_norm_loop, mx.compile(g2), x)
+
 
 if __name__ == "__main__":
     time_rms_norm()

cmake/extension.cmake

@@ -1,5 +1,7 @@
 include(CMakeParseArguments)
 
+# clang format off
+#
 # ##############################################################################
 # Build metal library
 #
@@ -11,6 +13,8 @@ include(CMakeParseArguments)
 # of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
 # files (like headers)
 #
+# clang format on
+
 macro(mlx_build_metallib)
   # Parse args
   set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
@@ -21,7 +25,7 @@ macro(mlx_build_metallib)
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
   # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math)
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
 
   # Prepare metallib build command
   add_custom_command(

docs/src/dev/extensions.rst

@@ -22,12 +22,12 @@ You can do that in MLX directly:
 This function performs that operation while leaving the implementation and
 function transformations to MLX.
 
-However you may need to customize the underlying implementation, perhaps to
-make it faster or for custom differentiation. In this tutorial we will go
-through adding custom extensions. It will cover:
+However, you may want to customize the underlying implementation, perhaps to
+make it faster. In this tutorial we will go through adding custom extensions.
+It will cover:
 
 * The structure of the MLX library.
-* Implementing a CPU operation that redirects to Accelerate_ when appropriate.
+* Implementing a CPU operation.
 * Implementing a GPU operation using metal.
 * Adding the ``vjp`` and ``jvp`` function transformation.
 * Building a custom extension and binding it to python.
@@ -45,7 +45,7 @@ Operations
 Operations are the front-end functions that operate on arrays. They are defined
 in the C++ API (:ref:`cpp_ops`), and the Python API (:ref:`ops`) binds them.
 
-We would like an operation, :meth:`axpby` that takes in two arrays ``x`` and
+We would like an operation :meth:`axpby` that takes in two arrays, ``x`` and
 ``y``, and two scalars, ``alpha`` and ``beta``. This is how to define it in
 C++:
 
@@ -55,7 +55,7 @@ C++:
     *  Scale and sum two vectors element-wise
     *  z = alpha * x + beta * y
     *
-    *  Follow numpy style broadcasting between x and y
+    *  Use NumPy-style broadcasting between x and y
     *  Inputs are upcasted to floats if needed
     **/
     array axpby(
@@ -66,7 +66,7 @@ C++:
         StreamOrDevice s = {} // Stream on which to schedule the operation
     );
 
-The simplest way to this operation is in terms of existing operations:
+The simplest way to implement this is with existing operations:
 
 .. code-block:: C++
 
@@ -153,9 +153,6 @@ more concrete:
       private:
         float alpha_;
         float beta_;
-
-        /** Fall back implementation for evaluation on CPU */
-        void eval(const std::vector<array>& inputs, array& out);
     };
 
 The :class:`Axpby` class derives from the base :class:`Primitive` class. The
@@ -188,7 +185,7 @@ Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
         auto promoted_dtype = promote_types(x.dtype(), y.dtype());
 
         // Upcast to float32 for non-floating point inputs x and y
-        auto out_dtype = is_floating_point(promoted_dtype)
+        auto out_dtype = issubdtype(promoted_dtype, float32)
             ? promoted_dtype
             : promote_types(promoted_dtype, float32);
 
@@ -234,49 +231,59 @@ the execution of the computation graph, and calls :meth:`Axpby::eval_cpu` or
 Implementing the CPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Let's start by implementing a naive and generic version of
-:meth:`Axpby::eval_cpu`. We declared this as a private member function of
-:class:`Axpby` earlier called :meth:`Axpby::eval`.
+Let's start by implementing :meth:`Axpby::eval_cpu`.
 
-Our naive method will go over each element of the output array, find the
+The method will go over each element of the output array, find the
 corresponding input elements of ``x`` and ``y`` and perform the operation
 point-wise. This is captured in the templated function :meth:`axpby_impl`.
 
 .. code-block:: C++
 
-    template <typename T>
-    void axpby_impl(
-            const array& x,
-            const array& y,
-            array& out,
-            float alpha_,
-            float beta_) {
-        // We only allocate memory when we are ready to fill the output
-        // malloc_or_wait synchronously allocates available memory
-        // There may be a wait executed here if the allocation is requested
-        // under memory-pressured conditions
-        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  template <typename T>
+  void axpby_impl(
+      const mx::array& x,
+      const mx::array& y,
+      mx::array& out,
+      float alpha_,
+      float beta_,
+      mx::Stream stream) {
+    // Allocate the output with `malloc_or_wait` which synchronously allocates
+    // memory, potentially waiting if the system is under memory pressure
+    out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
 
-        // Collect input and output data pointers
-        const T* x_ptr = x.data<T>();
-        const T* y_ptr = y.data<T>();
-        T* out_ptr = out.data<T>();
+    // Get the CPU command encoder and register input and output arrays
+    auto& encoder = mx::cpu::get_command_encoder(stream);
+    encoder.set_input_array(x);
+    encoder.set_input_array(y);
+    encoder.set_output_array(out);
 
-        // Cast alpha and beta to the relevant types
-        T alpha = static_cast<T>(alpha_);
-        T beta = static_cast<T>(beta_);
+    // Launch the CPU kernel
+    encoder.dispatch([x_ptr = x.data<T>(),
+                      y_ptr = y.data<T>(),
+                      out_ptr = out.data<T>(),
+                      size = out.size(),
+                      shape = out.shape(),
+                      x_strides = x.strides(),
+                      y_strides = y.strides(),
+                      alpha_,
+                      beta_]() {
 
-        // Do the element-wise operation for each output
-        for (size_t out_idx = 0; out_idx < out.size(); out_idx++) {
-            // Map linear indices to offsets in x and y
-            auto x_offset = elem_to_loc(out_idx, x.shape(), x.strides());
-            auto y_offset = elem_to_loc(out_idx, y.shape(), y.strides());
+      // Cast alpha and beta to the relevant types
+      T alpha = static_cast<T>(alpha_);
+      T beta = static_cast<T>(beta_);
 
-            // We allocate the output to be contiguous and regularly strided
-            // (defaults to row major) and hence it doesn't need additional mapping
-            out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
-        }
-    }
+      // Do the element-wise operation for each output
+      for (size_t out_idx = 0; out_idx < size; out_idx++) {
+        // Map linear indices to offsets in x and y
+        auto x_offset = mx::elem_to_loc(out_idx, shape, x_strides);
+        auto y_offset = mx::elem_to_loc(out_idx, shape, y_strides);
+
+        // We allocate the output to be contiguous and regularly strided
+        // (defaults to row major) and hence it doesn't need additional mapping
+        out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
+      }
+    });
+  }
 
 Our implementation should work for all incoming floating point arrays.
 Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
@@ -284,112 +291,32 @@ Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
 
 .. code-block:: C++
 
-    /** Fall back implementation for evaluation on CPU */
-    void Axpby::eval(
-      const std::vector<array>& inputs,
-      const std::vector<array>& outputs) {
-        auto& x = inputs[0];
-        auto& y = inputs[1];
-        auto& out = outputs[0];
-
-        // Dispatch to the correct dtype
-        if (out.dtype() == float32) {
-            return axpby_impl<float>(x, y, out, alpha_, beta_);
-        } else if (out.dtype() == float16) {
-            return axpby_impl<float16_t>(x, y, out, alpha_, beta_);
-        } else if (out.dtype() == bfloat16) {
-            return axpby_impl<bfloat16_t>(x, y, out, alpha_, beta_);
-        } else if (out.dtype() == complex64) {
-            return axpby_impl<complex64_t>(x, y, out, alpha_, beta_);
-        } else {
-            throw std::runtime_error(
-                "[Axpby] Only supports floating point types.");
-        }
-    }
-
-This is good as a fallback implementation. We can use the ``axpby`` routine
-provided by the Accelerate_ framework for a faster implementation in certain
-cases:
-
-#.  Accelerate does not provide implementations of ``axpby`` for half precision
-    floats. We can only use it for ``float32`` types.
-#.  Accelerate assumes the inputs ``x`` and ``y`` are contiguous and all
-    elements have fixed strides between them. We only direct to Accelerate
-    if both ``x`` and ``y`` are row contiguous or column contiguous.
-#.  Accelerate performs the routine ``Y = (alpha * X) + (beta * Y)`` in-place.
-    MLX expects to write the output to a new array. We must copy the elements
-    of ``y`` into the output and use that as an input to ``axpby``.
-
-Let's write an implementation that uses Accelerate in the right conditions.
-It allocates data for the output, copies ``y`` into it, and then calls the
-:func:`catlas_saxpby` from accelerate.
-
-.. code-block:: C++
-
-    template <typename T>
-    void axpby_impl_accelerate(
-            const array& x,
-            const array& y,
-            array& out,
-            float alpha_,
-            float beta_) {
-        // Accelerate library provides catlas_saxpby which does
-        // Y = (alpha * X) + (beta * Y) in place
-        // To use it, we first copy the data in y over to the output array
-        out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-        // We then copy over the elements using the contiguous vector specialization
-        copy_inplace(y, out, CopyType::Vector);
-
-        // Get x and y pointers for catlas_saxpby
-        const T* x_ptr = x.data<T>();
-        T* y_ptr = out.data<T>();
-
-        T alpha = static_cast<T>(alpha_);
-        T beta = static_cast<T>(beta_);
-
-        // Call the inplace accelerate operator
-        catlas_saxpby(
-            /* N = */ out.size(),
-            /* ALPHA = */ alpha,
-            /* X = */ x_ptr,
-            /* INCX = */ 1,
-            /* BETA = */ beta,
-            /* Y = */ y_ptr,
-            /* INCY = */ 1);
-    }
-
-For inputs that do not fit the criteria for accelerate, we fall back to
-:meth:`Axpby::eval`. With this in mind, let's finish our
-:meth:`Axpby::eval_cpu`.
-
-.. code-block:: C++
-
-    /** Evaluate primitive on CPU using accelerate specializations */
     void Axpby::eval_cpu(
-      const std::vector<array>& inputs,
-      const std::vector<array>& outputs) {
-        assert(inputs.size() == 2);
-        auto& x = inputs[0];
-        auto& y = inputs[1];
-        auto& out = outputs[0];
+        const std::vector<mx::array>& inputs,
+        std::vector<mx::array>& outputs) {
+      auto& x = inputs[0];
+      auto& y = inputs[1];
+      auto& out = outputs[0];
 
-        // Accelerate specialization for contiguous single precision float arrays
-        if (out.dtype() == float32 &&
-            ((x.flags().row_contiguous && y.flags().row_contiguous) ||
-            (x.flags().col_contiguous && y.flags().col_contiguous))) {
-            axpby_impl_accelerate<float>(x, y, out, alpha_, beta_);
-            return;
-        }
-
-        // Fall back to common back-end if specializations are not available
-        eval(inputs, outputs);
+      // Dispatch to the correct dtype
+      if (out.dtype() == mx::float32) {
+        return axpby_impl<float>(x, y, out, alpha_, beta_, stream());
+      } else if (out.dtype() == mx::float16) {
+        return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_, stream());
+      } else if (out.dtype() == mx::bfloat16) {
+        return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_, stream());
+      } else if (out.dtype() == mx::complex64) {
+        return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_, stream());
+      } else {
+        throw std::runtime_error(
+            "Axpby is only supported for floating point types.");
+      }
     }
 
 Just this much is enough to run the operation :meth:`axpby` on a CPU stream! If
 you do not plan on running the operation on the GPU or using transforms on
 computation graphs that contain :class:`Axpby`, you can stop implementing the
-primitive here and enjoy the speed-ups you get from the Accelerate library.
+primitive here.
 
 Implementing the GPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -824,7 +751,7 @@ Results
 ^^^^^^^
 
 Let's run a quick benchmark and see how our new ``axpby`` operation compares
-with the naive :meth:`simple_axpby` we first defined on the CPU.
+with the naive :meth:`simple_axpby` we first defined.
 
 .. code-block:: python
 
@@ -832,13 +759,11 @@ with the naive :meth:`simple_axpby` we first defined on the CPU.
     from mlx_sample_extensions import axpby
     import time
 
-    mx.set_default_device(mx.cpu)
-
     def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
         return alpha * x + beta * y
 
-    M = 256
-    N = 512
+    M = 4096
+    N = 4096
 
     x = mx.random.normal((M, N))
     y = mx.random.normal((M, N))
@@ -849,24 +774,24 @@ with the naive :meth:`simple_axpby` we first defined on the CPU.
 
     def bench(f):
         # Warm up
-        for i in range(100):
+        for i in range(5):
             z = f(x, y, alpha, beta)
             mx.eval(z)
 
         # Timed run
         s = time.time()
-        for i in range(5000):
+        for i in range(100):
             z = f(x, y, alpha, beta)
             mx.eval(z)
         e = time.time()
-        return e - s
+        return 1000 * (e - s) / 100
 
     simple_time = bench(simple_axpby)
     custom_time = bench(axpby)
 
-    print(f"Simple axpby: {simple_time:.3f} s | Custom axpby: {custom_time:.3f} s")
+    print(f"Simple axpby: {simple_time:.3f} ms | Custom axpby: {custom_time:.3f} ms")
 
-The results are ``Simple axpby: 0.114 s | Custom axpby: 0.109 s``. We see
+The results are ``Simple axpby: 1.559 ms | Custom axpby: 0.774 ms``. We see
 modest improvements right away!
 
 This operation is now good to be used to build other operations, in

docs/src/python/nn.rst

@@ -174,6 +174,7 @@ In detail:
 
    value_and_grad
    quantize
+   average_gradients
 
 .. toctree::
 

docs/src/usage/distributed.rst

@@ -5,21 +5,27 @@ Distributed Communication
 
 .. currentmodule:: mlx.core.distributed
 
-MLX utilizes `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ to
-provide distributed communication operations that allow the computational cost
-of training or inference to be shared across many physical machines. You can
-see a list of the supported operations in the :ref:`API docs<distributed>`.
+MLX supports distributed communication operations that allow the computational cost
+of training or inference to be shared across many physical machines. At the
+moment we support two different communication backends:
+
+* `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
+  full-featured and mature distributed communications library
+* A **ring** backend of our own that uses native TCP sockets and should be
+  faster for thunderbolt connections.
+
+The list of all currently supported operations and their documentation can be
+seen in the :ref:`API docs<distributed>`.
 
 .. note::
-   A lot of operations may not be supported or not as fast as they should be.
+   Some operations may not be supported or not as fast as they should be.
    We are adding more and tuning the ones we have as we are figuring out the
    best way to do distributed computing on Macs using MLX.
 
 Getting Started
 ---------------
 
-MLX already comes with the ability to "talk" to MPI if it is installed on the
-machine. The minimal distributed program in MLX is as simple as:
+A distributed program in MLX is as simple as:
 
 .. code:: python
 
@@ -30,74 +36,79 @@ machine. The minimal distributed program in MLX is as simple as:
     print(world.rank(), x)
 
 The program above sums the array ``mx.ones(10)`` across all
-distributed processes. If simply run with ``python``, however, only one
-process is launched and no distributed communication takes place.
+distributed processes. However, when this script is run with ``python`` only
+one process is launched and no distributed communication takes place. Namely,
+all operations in ``mx.distributed`` are noops when the distributed group has a
+size of one. This property allows us to avoid code that checks if we are in a
+distributed setting similar to the one below:
 
-To launch the program in distributed mode we need to use ``mpirun`` or
-``mpiexec`` depending on the MPI installation. The simplest possible way is the
-following:
+.. code:: python
+
+    import mlx.core as mx
+
+    x = ...
+    world = mx.distributed.init()
+    # No need for the check we can simply do x = mx.distributed.all_sum(x)
+    if world.size() > 1:
+        x = mx.distributed.all_sum(x)
+
+Running Distributed Programs
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+MLX provides ``mlx.launch`` a helper script to launch distributed programs.
+Continuing with our initial example we can run it on localhost with 4 processes using
 
 .. code:: shell
 
-    $ mpirun -np 2 python test.py
-    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    $ mlx.launch -n 4 my_script.py
+    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
 
-The above launches two processes on the same (local) machine and we can see
-both standard output streams. The processes send the array of 1s to each other
-and compute the sum which is printed. Launching with ``mpirun -np 4 ...`` would
-print 4 etc.
-
-Installing MPI
----------------
-
-MPI can be installed with Homebrew, using the Anaconda package manager or
-compiled from source. Most of our testing is done using ``openmpi`` installed
-with the Anaconda package manager as follows:
+We can also run it on some remote hosts by providing their IPs (provided that
+the script exists on all hosts and they are reachable by ssh)
 
 .. code:: shell
 
-    $ conda install conda-forge::openmpi
+    $ mlx.launch --hosts ip1,ip2,ip3,ip4 my_script.py
+    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
 
-Installing with Homebrew may require specifying the location of ``libmpi.dyld``
-so that MLX can find it and load it at runtime. This can simply be achieved by
-passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun``.
+Consult the dedicated :doc:`usage guide<launching_distributed>` for more
+information on using ``mlx.launch``.
 
-.. code:: shell
+Selecting Backend
+^^^^^^^^^^^^^^^^^
 
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
-
-Setting up Remote Hosts
------------------------
-
-MPI can automatically connect to remote hosts and set up the communication over
-the network if the remote hosts can be accessed via ssh. A good checklist to
-debug connectivity issues is the following:
-
-* ``ssh hostname`` works from all machines to all machines without asking for
-  password or host confirmation
-* ``mpirun`` is accessible on all machines. You can call ``mpirun`` using its
-  full path to force all machines to use a specific path.
-* Ensure that the ``hostname`` used by MPI is the one that you have configured
-  in the ``.ssh/config`` files on all machines.
+You can select the backend you want to use when calling :func:`init` by passing
+one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to
+initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they
+both fail then a singleton group is created.
 
 .. note::
-  For an example hostname ``foo.bar.com`` MPI can use only ``foo`` as
-  the hostname passed to ssh if the current hostname matches ``*.bar.com``.
+   After a distributed backend is successfully initialized :func:`init` will
+   return **the same backend** if called without arguments or with backend set to
+   ``any``.
 
-An easy way to pass the host names to MPI is using a host file. A host file
-looks like the following, where ``host1`` and ``host2`` should be the fully
-qualified domain names or IPs for these hosts.
+The following examples aim to clarify the backend initialization logic in MLX:
 
-.. code::
+.. code:: python
 
-    host1 slots=1
-    host2 slots=1
+    # Case 1: Initialize MPI regardless if it was possible to initialize the ring backend
+    world = mx.distributed.init(backend="mpi")
+    world2 = mx.distributed.init()  # subsequent calls return the MPI backend!
 
-When using MLX, it is very likely that you want to use 1 slot per host, ie one
-process per host.  The hostfile also needs to contain the current
-host if you want to run on the local host. Passing the host file to
-``mpirun`` is simply done using the ``--hostfile`` command line argument.
+    # Case 2: Initialize any backend
+    world = mx.distributed.init(backend="any")  # equivalent to no arguments
+    world2 = mx.distributed.init()  # same as above
+
+    # Case 3: Initialize both backends at the same time
+    world_mpi = mx.distributed.init(backend="mpi")
+    world_ring = mx.distributed.init(backend="ring")
+    world_any = mx.distributed.init()  # same as MPI because it was initialized first!
 
 Training Example
 ----------------
@@ -155,13 +166,179 @@ everything else remaining the same.
         optimizer.update(model, grads)
         return loss
 
-Tuning All Reduce
------------------
+Utilizing ``nn.average_gradients``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-We are working on improving the performance of all reduce on MLX but for now
-the two main things one can do to extract the most out of distributed training with MLX are:
+Although the code example above works correctly; it performs one communication
+per gradient. It is significantly more efficient to aggregate several gradients
+together and perform fewer communication steps.
 
-1. Perform a few large reductions instead of many small ones to improve
-   bandwidth and latency
-2. Pass ``--mca btl_tcp_links 4`` to ``mpirun`` to configure it to use 4 tcp
-   connections between each host to improve bandwidth
+This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
+almost identical to the example above:
+
+.. code:: python
+
+    model = ...
+    optimizer = ...
+    dataset = ...
+
+    def step(model, x, y):
+        loss, grads = loss_grad_fn(model, x, y)
+        grads = mlx.nn.average_gradients(grads) # <---- This line was added
+        optimizer.update(model, grads)
+        return loss
+
+    for x, y in dataset:
+        loss = step(model, x, y)
+        mx.eval(loss, model.parameters())
+
+
+Getting Started with MPI
+------------------------
+
+MLX already comes with the ability to "talk" to MPI if it is installed on the
+machine. Launching distributed MLX programs that use MPI can be done with
+``mpirun`` as expected. However, in the following examples we will be using
+``mlx.launch --backend mpi`` which takes care of some nuisances such as setting
+absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` shared
+library.
+
+The simplest possible usage is the following which, assuming the minimal
+example in the beginning of this page, should result in:
+
+.. code:: shell
+
+    $ mlx.launch --backend mpi -n 2 test.py
+    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+
+The above launches two processes on the same (local) machine and we can see
+both standard output streams. The processes send the array of 1s to each other
+and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would
+print 4 etc.
+
+Installing MPI
+^^^^^^^^^^^^^^
+
+MPI can be installed with Homebrew, using the Anaconda package manager or
+compiled from source. Most of our testing is done using ``openmpi`` installed
+with the Anaconda package manager as follows:
+
+.. code:: shell
+
+    $ conda install conda-forge::openmpi
+
+Installing with Homebrew may require specifying the location of ``libmpi.dyld``
+so that MLX can find it and load it at runtime. This can simply be achieved by
+passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
+done automatically by ``mlx.launch``.
+
+.. code:: shell
+
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
+    $ # or simply
+    $ mlx.launch -n 2 test.py
+
+Setting up Remote Hosts
+^^^^^^^^^^^^^^^^^^^^^^^
+
+MPI can automatically connect to remote hosts and set up the communication over
+the network if the remote hosts can be accessed via ssh. A good checklist to
+debug connectivity issues is the following:
+
+* ``ssh hostname`` works from all machines to all machines without asking for
+  password or host confirmation
+* ``mpirun`` is accessible on all machines.
+* Ensure that the ``hostname`` used by MPI is the one that you have configured
+  in the ``.ssh/config`` files on all machines.
+
+Tuning MPI All Reduce
+^^^^^^^^^^^^^^^^^^^^^
+
+.. note::
+
+    For faster all reduce consider using the ring backend either with Thunderbolt
+    connections or over Ethernet.
+
+Configure MPI to use N tcp connections between each host to improve bandwidth
+by passing ``--mca btl_tcp_links N``.
+
+Force MPI to use the most performant network interface by setting ``--mca
+btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want
+to use.
+
+Getting Started with Ring
+-------------------------
+
+The ring backend does not depend on any third party library so it is always
+available. It uses TCP sockets so the nodes need to be reachable via a network.
+As the name suggests the nodes are connected in a ring which means that rank 1
+can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3
+and so on and so forth. As a result :func:`send` and :func:`recv` with
+arbitrary sender and receiver is not supported in the ring backend.
+
+Defining a Ring
+^^^^^^^^^^^^^^^
+
+The easiest way to define and use a ring is via a JSON hostfile and the
+``mlx.launch`` :doc:`helper script <launching_distributed>`. For each node one
+defines a hostname to ssh into to run commands on this node and one or more IPs
+that this node will listen to for connections.
+
+For example the hostfile below defines a 4 node ring. ``hostname1`` will be
+rank 0, ``hostname2`` rank 1 etc.
+
+.. code:: json
+
+    [
+        {"ssh": "hostname1", "ips": ["123.123.123.1"]},
+        {"ssh": "hostname2", "ips": ["123.123.123.2"]},
+        {"ssh": "hostname3", "ips": ["123.123.123.3"]},
+        {"ssh": "hostname4", "ips": ["123.123.123.4"]}
+    ]
+
+Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each
+node, run the script which will listen for connections in each of the provided
+IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a
+connection from ``123.123.123.4`` and so on and so forth.
+
+Thunderbolt Ring
+^^^^^^^^^^^^^^^^
+
+Although the ring backend can have benefits over MPI even for Ethernet, its
+main purpose is to use Thunderbolt rings for higher bandwidth communication.
+Setting up such thunderbolt rings can be done manually, but is a relatively
+tedious process. To simplify this, we provide the utility ``mlx.distributed_config``.
+
+To use ``mlx.distributed_config`` your computers need to be accessible by ssh via
+Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the
+utility as follows:
+
+.. code:: shell
+
+   mlx.distributed_config --verbose --hosts host1,host2,host3,host4
+
+By default the script will attempt to discover the thunderbolt ring and provide
+you with the commands to configure each node as well as the ``hostfile.json``
+to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes
+then ``--auto-setup`` can be used to configure them automatically.
+
+To validate your connection without configuring anything
+``mlx.distributed_config`` can also plot the ring using DOT format.
+
+.. code:: shell
+
+   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot
+   dot -Tpng ring.dot >ring.png
+   open ring.png
+
+If you want to go through the process manually, the steps are as follows:
+
+* Disable the thunderbolt bridge interface
+* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces
+  corresponding to that cable in nodes ``i`` and ``i + 1``.
+* Set up a unique subnetwork connecting the two nodes for the corresponding
+  interfaces. For instance if the cable corresponds to ``en2`` on node ``i``
+  and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and
+  ``192.168.0.2`` respectively to the two nodes. For more details you can see
+  the commands prepared by the utility script.

docs/src/usage/launching_distributed.rst

@@ -0,0 +1,105 @@
+:orphan:
+
+.. _usage_launch_distributed:
+
+Launching Distributed Programs
+==============================
+
+.. currentmodule:: mlx.core.distributed
+
+Installing the MLX python package provides a helper script ``mlx.launch`` that
+can be used to run python scripts distributed on several nodes. It allows
+launching using either the MPI backend or the ring backend. See the
+:doc:`distributed docs <distributed>` for the different backends.
+
+Usage
+-----
+
+The minimal usage example of ``mlx.launch`` is simply
+
+.. code:: shell
+
+    mlx.launch --hosts ip1,ip2 my_script.py
+
+or for testing on localhost
+
+.. code:: shell
+
+    mlx.launch -n 2 my_script.py
+
+The ``mlx.launch`` command connects to the provided host and launches the input
+script on each host. It monitors each of the launched processes and terminates
+the rest if one of them fails unexpectedly or if ``mlx.launch`` is terminated.
+It also takes care of forwarding the output of each remote process to stdout
+and stderr respectively.
+
+Providing Hosts
+^^^^^^^^^^^^^^^^
+
+Hosts can be provided as command line arguments, like above, but the way that
+allows to fully define a list of hosts is via a JSON hostfile. The hostfile has
+a very simple schema. It is simply a list of objects that define each host via
+a hostname to ssh to and a list of IPs to utilize for the communication.
+
+.. code:: json
+
+    [
+        {"ssh": "hostname1", "ips": ["123.123.1.1", "123.123.2.1"]},
+        {"ssh": "hostname2", "ips": ["123.123.1.2", "123.123.2.2"]}
+    ]
+
+You can use ``mlx.distributed_config --over ethernet`` to create a hostfile
+with IPs corresponding to the ``en0`` interface.
+
+Setting up Remote Hosts
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+In order to be able to launch the script on each host we need to be able to
+connect via ssh. Moreover the input script and python binary need to be on each
+host and on the same path. A good checklist to debug errors is the following:
+
+* ``ssh hostname`` works without asking for password or host confirmation
+* the python binary is available on all hosts at the same path. You can use
+  ``mlx.launch --print-python`` to see what that path is.
+* the script you want to run is available on all hosts at the same path
+
+.. _mpi_specifics:
+
+MPI Specifics
+-------------
+
+One can use MPI by passing ``--backend mpi`` to ``mlx.launch``. In that case,
+``mlx.launch`` is a thin wrapper over ``mpirun``. Moreover,
+
+* The IPs in the hostfile are ignored
+* The ssh connectivity requirement is stronger as every node needs to be able
+  to connect to every other node
+* ``mpirun`` needs to be available on every node at the same path
+
+Finally, one can pass arguments to ``mpirun`` using ``--mpi-arg``. For instance
+to choose a specific interface for the byte-transfer-layer of MPI we can call
+``mlx.launch`` as follows:
+
+.. code:: shell
+
+    mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py
+
+
+.. _ring_specifics:
+
+Ring Specifics
+--------------
+
+The ring backend, which is also the default backend, can be explicitly selected
+with the argument ``--backend ring``. The ring backend has some specific
+requirements and arguments that are different to MPI:
+
+* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to
+  ssh to a hostname that does not correspond to the IP we want to bind to we
+  have to provide a hostfile.
+* ``--starting-port`` defines the port to bind to on the remote hosts.
+  Specifically rank 0 for the first IP will use this port and each subsequent
+  IP or rank will add 1 to this port.
+* ``--connections-per-ip`` allows us to increase the number of connections
+  between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for
+  ``mpirun``.

examples/extensions/CMakeLists.txt

@@ -10,7 +10,6 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
 
 # ----------------------------- Dependencies -----------------------------
-find_package(MLX CONFIG REQUIRED)
 find_package(
   Python 3.8
   COMPONENTS Interpreter Development.Module
@@ -21,6 +20,12 @@ execute_process(
   OUTPUT_VARIABLE nanobind_ROOT)
 find_package(nanobind CONFIG REQUIRED)
 
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE MLX_ROOT)
+find_package(MLX CONFIG REQUIRED)
+
 # ----------------------------- Extensions -----------------------------
 
 # Add library

examples/extensions/axpby/axpby.cpp

@@ -1,20 +1,14 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
-#include <cassert>
 #include <iostream>
 #include <sstream>
 
-#include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/utils.h"
 
 #include "axpby/axpby.h"
 
-#ifdef ACCELERATE_NEW_LAPACK
-#include <vecLib/cblas_new.h>
-#endif
-
 #ifdef _METAL_
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/utils.h"
@@ -76,136 +70,67 @@ void axpby_impl(
     const mx::array& y,
     mx::array& out,
     float alpha_,
-    float beta_) {
-  // We only allocate memory when we are ready to fill the output
-  // malloc_or_wait synchronously allocates available memory
-  // There may be a wait executed here if the allocation is requested
-  // under memory-pressured conditions
+    float beta_,
+    mx::Stream stream) {
+  // Allocate the output with `malloc_or_wait` which synchronously allocates
+  // memory, potentially waiting if the system is under memory pressure
   out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
 
-  // Collect input and output data pointers
-  const T* x_ptr = x.data<T>();
-  const T* y_ptr = y.data<T>();
-  T* out_ptr = out.data<T>();
+  // Get the CPU command encoder and register input and output arrays
+  auto& encoder = mx::cpu::get_command_encoder(stream);
+  encoder.set_input_array(x);
+  encoder.set_input_array(y);
+  encoder.set_output_array(out);
 
-  // Cast alpha and beta to the relevant types
-  T alpha = static_cast<T>(alpha_);
-  T beta = static_cast<T>(beta_);
+  // Launch the CPU kernel
+  encoder.dispatch([x_ptr = x.data<T>(),
+                    y_ptr = y.data<T>(),
+                    out_ptr = out.data<T>(),
+                    size = out.size(),
+                    shape = out.shape(),
+                    x_strides = x.strides(),
+                    y_strides = y.strides(),
+                    alpha_,
+                    beta_]() {
+    // Cast alpha and beta to the relevant types
+    T alpha = static_cast<T>(alpha_);
+    T beta = static_cast<T>(beta_);
 
-  // Do the element-wise operation for each output
-  for (size_t out_idx = 0; out_idx < out.size(); out_idx++) {
-    // Map linear indices to offsets in x and y
-    auto x_offset = mx::elem_to_loc(out_idx, x.shape(), x.strides());
-    auto y_offset = mx::elem_to_loc(out_idx, y.shape(), y.strides());
+    // Do the element-wise operation for each output
+    for (size_t out_idx = 0; out_idx < size; out_idx++) {
+      // Map linear indices to offsets in x and y
+      auto x_offset = mx::elem_to_loc(out_idx, shape, x_strides);
+      auto y_offset = mx::elem_to_loc(out_idx, shape, y_strides);
 
-    // We allocate the output to be contiguous and regularly strided
-    // (defaults to row major) and hence it doesn't need additional mapping
-    out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
-  }
+      // We allocate the output to be contiguous and regularly strided
+      // (defaults to row major) and hence it doesn't need additional mapping
+      out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
+    }
+  });
 }
 
-/** Fall back implementation for evaluation on CPU */
-void Axpby::eval(
+void Axpby::eval_cpu(
     const std::vector<mx::array>& inputs,
     std::vector<mx::array>& outputs) {
-  // Check the inputs (registered in the op while constructing the out array)
-  assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
   auto& out = outputs[0];
 
   // Dispatch to the correct dtype
   if (out.dtype() == mx::float32) {
-    return axpby_impl<float>(x, y, out, alpha_, beta_);
+    return axpby_impl<float>(x, y, out, alpha_, beta_, stream());
   } else if (out.dtype() == mx::float16) {
-    return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_);
+    return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_, stream());
   } else if (out.dtype() == mx::bfloat16) {
-    return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_);
+    return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_, stream());
   } else if (out.dtype() == mx::complex64) {
-    return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_);
+    return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_, stream());
   } else {
     throw std::runtime_error(
         "Axpby is only supported for floating point types.");
   }
 }
 
-///////////////////////////////////////////////////////////////////////////////
-// Primitive Accelerate Backend Implementation
-///////////////////////////////////////////////////////////////////////////////
-
-#ifdef ACCELERATE_NEW_LAPACK
-
-template <typename T>
-void axpby_impl_accelerate(
-    const mx::array& x,
-    const mx::array& y,
-    mx::array& out,
-    float alpha_,
-    float beta_) {
-  // Accelerate library provides catlas_saxpby which does
-  // Y = (alpha * X) + (beta * Y) in place
-  // To use it, we first copy the data in y over to the output array
-
-  // This specialization requires both x and y be contiguous in the same mode
-  // i.e: corresponding linear indices in both point to corresponding elements
-  // The data in the output array is allocated to match the strides in y
-  // such that x, y, and out are contiguous in the same mode and
-  // no transposition is needed
-  out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
-
-  // We then copy over the elements using the contiguous vector specialization
-  copy_inplace(y, out, mx::CopyType::Vector);
-
-  // Get x and y pointers for catlas_saxpby
-  const T* x_ptr = x.data<T>();
-  T* y_ptr = out.data<T>();
-
-  T alpha = static_cast<T>(alpha_);
-  T beta = static_cast<T>(beta_);
-
-  // Call the inplace accelerate operator
-  catlas_saxpby(
-      /* N = */ out.size(),
-      /* ALPHA = */ alpha,
-      /* X = */ x_ptr,
-      /* INCX = */ 1,
-      /* BETA = */ beta,
-      /* Y = */ y_ptr,
-      /* INCY = */ 1);
-}
-
-/** Evaluate primitive on CPU using accelerate specializations */
-void Axpby::eval_cpu(
-    const std::vector<mx::array>& inputs,
-    std::vector<mx::array>& outputs) {
-  assert(inputs.size() == 2);
-  auto& x = inputs[0];
-  auto& y = inputs[1];
-  auto& out = outputs[0];
-
-  // Accelerate specialization for contiguous single precision float arrays
-  if (out.dtype() == mx::float32 &&
-      ((x.flags().row_contiguous && y.flags().row_contiguous) ||
-       (x.flags().col_contiguous && y.flags().col_contiguous))) {
-    axpby_impl_accelerate<float>(x, y, out, alpha_, beta_);
-    return;
-  }
-
-  // Fall back to common backend if specializations are not available
-  eval(inputs, outputs);
-}
-
-#else // Accelerate not available
-
-/** Evaluate primitive on CPU falling back to common backend */
-void Axpby::eval_cpu(
-    const std::vector<mx::array>& inputs,
-    std::vector<mx::array>& outputs) {
-  eval(inputs, outputs);
-}
-
-#endif
-
 ///////////////////////////////////////////////////////////////////////////////
 // Primitive Metal Backend Implementation
 ///////////////////////////////////////////////////////////////////////////////
@@ -217,7 +142,6 @@ void Axpby::eval_gpu(
     const std::vector<mx::array>& inputs,
     std::vector<mx::array>& outputs) {
   // Prepare inputs
-  assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
   auto& out = outputs[0];

examples/extensions/axpby/axpby.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #pragma once
 
@@ -85,11 +85,6 @@ class Axpby : public mx::Primitive {
  private:
   float alpha_;
   float beta_;
-
-  /** Fall back implementation for evaluation on CPU */
-  void eval(
-      const std::vector<mx::array>& inputs,
-      std::vector<mx::array>& outputs);
 };
 
 } // namespace my_ext

examples/extensions/axpby/axpby.metal

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #include <metal_stdlib>
 

mlx/CMakeLists.txt

@@ -17,6 +17,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
 
 if(MSVC)

mlx/array.cpp

@@ -76,35 +76,27 @@ array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter)
   set_data(data, deleter);
 }
 
-array::array(
-    allocator::Buffer data,
-    Shape shape,
-    Dtype dtype,
-    Strides strides,
-    size_t data_size,
-    Flags flags,
-    Deleter deleter)
-    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
-  set_data(data, data_size, std::move(strides), flags, deleter);
-}
-
 void array::detach() {
+  array_desc_->primitive = nullptr;
+  for (auto& s : array_desc_->siblings) {
+    s.array_desc_->primitive = nullptr;
+  }
   for (auto& s : array_desc_->siblings) {
     s.array_desc_->inputs.clear();
     s.array_desc_->siblings.clear();
     s.array_desc_->position = 0;
-    s.array_desc_->primitive = nullptr;
   }
   array_desc_->inputs.clear();
   array_desc_->siblings.clear();
   array_desc_->position = 0;
-  array_desc_->primitive = nullptr;
 }
 
 bool array::is_available() const {
   if (status() == Status::available) {
     return true;
-  } else if (status() == Status::evaluated && event().is_signaled()) {
+  } else if (
+      status() == Status::evaluated &&
+      (!event().valid() || event().is_signaled())) {
     set_status(Status::available);
     return true;
   }
@@ -113,7 +105,10 @@ bool array::is_available() const {
 
 void array::wait() {
   if (!is_available()) {
-    event().wait();
+    if (event().valid()) {
+      event().wait();
+      detach_event();
+    }
     set_status(Status::available);
   }
 }
@@ -174,34 +169,13 @@ void array::copy_shared_buffer(const array& other) {
   copy_shared_buffer(other, other.strides(), other.flags(), other.data_size());
 }
 
-void array::move_shared_buffer(
-    array other,
-    const Strides& strides,
-    Flags flags,
-    size_t data_size,
-    size_t offset /* = 0 */) {
-  array_desc_->data = std::move(other.array_desc_->data);
-  array_desc_->strides = strides;
-  array_desc_->flags = flags;
-  array_desc_->data_size = data_size;
-  auto char_offset = sizeof(char) * itemsize() * offset;
-  auto data_ptr = other.array_desc_->data_ptr;
-  other.array_desc_->data_ptr = nullptr;
-  array_desc_->data_ptr =
-      static_cast<void*>(static_cast<char*>(data_ptr) + char_offset);
-}
-
-void array::move_shared_buffer(array other) {
-  move_shared_buffer(other, other.strides(), other.flags(), other.data_size());
-}
-
 array::~array() {
   if (array_desc_ == nullptr) {
     return;
   }
 
-  // Ignore arrays that might be detached during eval
-  if (status() == array::Status::scheduled) {
+  // Detached/detaching
+  if (array_desc_->primitive == nullptr) {
     return;
   }
 

mlx/array.h

@@ -243,18 +243,6 @@ class array {
     bool col_contiguous : 1;
   };
 
-  /** Build an array from all the info held by the array description. Including
-   * the buffer, strides, flags.
-   */
-  explicit array(
-      allocator::Buffer data,
-      Shape shape,
-      Dtype dtype,
-      Strides strides,
-      size_t data_size,
-      Flags flags,
-      Deleter deleter = allocator::free);
-
   /** The array's primitive. */
   Primitive& primitive() const {
     return *(array_desc_->primitive);
@@ -365,11 +353,6 @@ class array {
     // For example, the status of `x` in `auto x = a + b`.
     unscheduled,
 
-    // The ouptut of a computation which has been scheduled but `eval_*` has
-    // not yet been called on the array's primitive. A possible
-    // status of `x` in `auto x = a + b; eval(x);`
-    scheduled,
-
     // The array's `eval_*` function has been run, but the computation is not
     // necessarily complete. The array will have memory allocated and if it is
     // not a tracer then it will be detached from the graph.
@@ -406,6 +389,10 @@ class array {
     array_desc_->event = std::move(e);
   }
 
+  void detach_event() const {
+    array_desc_->event = Event{};
+  }
+
   // Mark the array as a tracer array (true) or not.
   void set_tracer(bool is_tracer) {
     array_desc_->is_tracer = is_tracer;
@@ -431,15 +418,6 @@ class array {
 
   void copy_shared_buffer(const array& other);
 
-  void move_shared_buffer(
-      array other,
-      const Strides& strides,
-      Flags flags,
-      size_t data_size,
-      size_t offset = 0);
-
-  void move_shared_buffer(array other);
-
   void overwrite_descriptor(const array& other) {
     array_desc_ = other.array_desc_;
   }

mlx/backend/common/binary.h

@@ -38,8 +38,7 @@ inline void set_binary_op_output_data(
     const array& a,
     const array& b,
     array& out,
-    BinaryOpType bopt,
-    bool donate_with_move = false) {
+    BinaryOpType bopt) {
   bool b_donatable = is_donatable(b, out);
   bool a_donatable = is_donatable(a, out);
   switch (bopt) {
@@ -49,11 +48,7 @@ inline void set_binary_op_output_data(
       break;
     case BinaryOpType::ScalarVector:
       if (b_donatable) {
-        if (donate_with_move) {
-          out.move_shared_buffer(b);
-        } else {
-          out.copy_shared_buffer(b);
-        }
+        out.copy_shared_buffer(b);
       } else {
         out.set_data(
             allocator::malloc_or_wait(b.data_size() * out.itemsize()),
@@ -64,11 +59,7 @@ inline void set_binary_op_output_data(
       break;
     case BinaryOpType::VectorScalar:
       if (a_donatable) {
-        if (donate_with_move) {
-          out.move_shared_buffer(a);
-        } else {
-          out.copy_shared_buffer(a);
-        }
+        out.copy_shared_buffer(a);
       } else {
         out.set_data(
             allocator::malloc_or_wait(a.data_size() * out.itemsize()),
@@ -79,17 +70,9 @@ inline void set_binary_op_output_data(
       break;
     case BinaryOpType::VectorVector:
       if (a_donatable) {
-        if (donate_with_move) {
-          out.move_shared_buffer(a);
-        } else {
-          out.copy_shared_buffer(a);
-        }
+        out.copy_shared_buffer(a);
       } else if (b_donatable) {
-        if (donate_with_move) {
-          out.move_shared_buffer(b);
-        } else {
-          out.copy_shared_buffer(b);
-        }
+        out.copy_shared_buffer(b);
       } else {
         out.set_data(
             allocator::malloc_or_wait(a.data_size() * out.itemsize()),
@@ -100,18 +83,10 @@ inline void set_binary_op_output_data(
       break;
     case BinaryOpType::General:
       if (a_donatable && a.flags().row_contiguous && a.size() == out.size()) {
-        if (donate_with_move) {
-          out.move_shared_buffer(a);
-        } else {
-          out.copy_shared_buffer(a);
-        }
+        out.copy_shared_buffer(a);
       } else if (
           b_donatable && b.flags().row_contiguous && b.size() == out.size()) {
-        if (donate_with_move) {
-          out.move_shared_buffer(b);
-        } else {
-          out.copy_shared_buffer(b);
-        }
+        out.copy_shared_buffer(b);
       } else {
         out.set_data(allocator::malloc_or_wait(out.nbytes()));
       }

mlx/backend/common/common.cpp

@@ -39,7 +39,7 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
   // rely on data_size anyway.
   size_t data_size = out.size();
 
-  return move_or_copy(in, out, strides_, flags, data_size, offset_);
+  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 
 void broadcast(const array& in, array& out) {
@@ -56,7 +56,7 @@ void broadcast(const array& in, array& out) {
   if (out.size() > in.size()) {
     flags.row_contiguous = flags.col_contiguous = false;
   }
-  move_or_copy(in, out, strides, flags, in.data_size());
+  out.copy_shared_buffer(in, strides, flags, in.data_size());
 }
 
 void Broadcast::eval(const std::vector<array>& inputs, array& out) {
@@ -69,7 +69,7 @@ void BroadcastAxes::eval(const std::vector<array>& inputs, array& out) {
 
 void Copy::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  move_or_copy(inputs[0], out);
+  out.copy_shared_buffer(inputs[0]);
 }
 
 void CustomTransforms::eval(
@@ -78,7 +78,7 @@ void CustomTransforms::eval(
   assert(inputs.size() > outputs.size());
   for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
        i++, j++) {
-    move_or_copy(inputs[j], outputs[i]);
+    outputs[i].copy_shared_buffer(inputs[j]);
   }
 }
 
@@ -87,7 +87,7 @@ void Depends::eval(
     std::vector<array>& outputs) {
   assert(inputs.size() > outputs.size());
   for (int i = 0; i < outputs.size(); i++) {
-    move_or_copy(inputs[i], outputs[i]);
+    outputs[i].copy_shared_buffer(inputs[i]);
   }
 }
 
@@ -98,7 +98,7 @@ void ExpandDims::eval(const std::vector<array>& inputs, array& out) {
   for (auto ax : axes_) {
     strides.insert(strides.begin() + ax, 1);
   }
-  move_or_copy(in, out, strides, in.flags(), in.data_size());
+  out.copy_shared_buffer(in, strides, in.flags(), in.data_size());
 }
 
 void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
@@ -210,7 +210,7 @@ void shared_buffer_reshape(
     auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
     flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
   }
-  move_or_copy(in, out, out_strides, flags, in.data_size());
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
 }
 
 void Split::eval(
@@ -276,12 +276,12 @@ void Squeeze::eval(const std::vector<array>& inputs, array& out) {
       strides.push_back(in.strides(i));
     }
   }
-  move_or_copy(in, out, strides, in.flags(), in.data_size());
+  out.copy_shared_buffer(in, strides, in.flags(), in.data_size());
 }
 
 void StopGradient::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  move_or_copy(inputs[0], out);
+  out.copy_shared_buffer(inputs[0]);
 }
 
 void Transpose::eval(const std::vector<array>& inputs, array& out) {
@@ -315,7 +315,7 @@ void Transpose::eval(const std::vector<array>& inputs, array& out) {
       b_stride *= out.shape(ri);
     }
   }
-  move_or_copy(in, out, out_strides, flags, in.data_size());
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
 }
 
 } // namespace mlx::core

mlx/backend/common/compiled.cpp

@@ -161,8 +161,7 @@ void compiled_allocate_outputs(
     std::vector<array>& outputs,
     const std::vector<array>& inputs_,
     const std::unordered_set<uintptr_t>& constant_ids_,
-    bool contiguous,
-    bool move_buffers /* = false */) {
+    bool contiguous) {
   if (contiguous) {
     int o = 0;
     Strides strides;
@@ -178,11 +177,7 @@ void compiled_allocate_outputs(
       if (in.itemsize() == outputs[o].itemsize() && !is_scalar(in) &&
           in.is_donatable() &&
           constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
-        if (move_buffers) {
-          outputs[o++].move_shared_buffer(in);
-        } else {
-          outputs[o++].copy_shared_buffer(in);
-        }
+        outputs[o++].copy_shared_buffer(in);
       }
       // Get representative input flags to properly set non-donated outputs
       if (strides.empty() && in.size() == outputs[0].size()) {
@@ -210,13 +205,8 @@ void compiled_allocate_outputs(
       if (in.flags().row_contiguous && in.size() == outputs[o].size() &&
           in.itemsize() == outputs[o].itemsize() && in.is_donatable() &&
           constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
-        if (move_buffers) {
-          outputs[o].move_shared_buffer(
-              in, outputs[o].strides(), in.flags(), in.data_size());
-        } else {
-          outputs[o].copy_shared_buffer(
-              in, outputs[o].strides(), in.flags(), in.data_size());
-        }
+        outputs[o].copy_shared_buffer(
+            in, outputs[o].strides(), in.flags(), in.data_size());
         o++;
       }
     }

mlx/backend/common/compiled.h

@@ -62,7 +62,6 @@ void compiled_allocate_outputs(
     std::vector<array>& outputs,
     const std::vector<array>& inputs_,
     const std::unordered_set<uintptr_t>& constant_ids_,
-    bool contiguous,
-    bool move_buffers = false);
+    bool contiguous);
 
 } // namespace mlx::core

mlx/backend/common/copy.h

@@ -22,4 +22,25 @@ enum class CopyType {
   GeneralGeneral
 };
 
+inline bool set_copy_output_data(const array& in, array& out, CopyType ctype) {
+  if (ctype == CopyType::Vector) {
+    // If the input is donateable, we are doing a vector copy and the types
+    // have the same size, then the input buffer can hold the output.
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(in);
+      return true;
+    } else {
+      out.set_data(
+          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+      return false;
+    }
+  } else {
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    return false;
+  }
+}
+
 } // namespace mlx::core

mlx/backend/common/load.cpp

@@ -3,7 +3,8 @@
 #include <algorithm>
 #include <utility>
 
-#include "mlx/backend/common/load.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
 
 namespace {
 
@@ -26,26 +27,31 @@ void swap_endianness(uint8_t* data_bytes, size_t N) {
 
 namespace mlx::core {
 
-void load(
-    array& out,
-    size_t offset,
-    const std::shared_ptr<io::Reader>& reader,
-    bool swap_endianness_) {
-  reader->read(out.data<char>(), out.nbytes(), offset);
-
-  if (swap_endianness_) {
-    switch (out.itemsize()) {
-      case 2:
-        swap_endianness<2>(out.data<uint8_t>(), out.data_size());
-        break;
-      case 4:
-        swap_endianness<4>(out.data<uint8_t>(), out.data_size());
-        break;
-      case 8:
-        swap_endianness<8>(out.data<uint8_t>(), out.data_size());
-        break;
+void Load::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  auto read_task = [out_ptr = out.data<char>(),
+                    size = out.size(),
+                    itemsize = out.itemsize(),
+                    offset = offset_,
+                    reader = reader_,
+                    swap_endianness_ = swap_endianness_]() mutable {
+    reader->read(out_ptr, size * itemsize, offset);
+    if (swap_endianness_) {
+      switch (itemsize) {
+        case 2:
+          swap_endianness<2>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+        case 4:
+          swap_endianness<4>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+        case 8:
+          swap_endianness<8>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+      }
     }
-  }
+  };
+  auto fut = io::thread_pool().enqueue(std::move(read_task)).share();
+  scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.wait(); });
 }
 
 } // namespace mlx::core

mlx/backend/common/load.h

@@ -1,14 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include "mlx/array.h"
-#include "mlx/io/load.h"
-
-namespace mlx::core {
-
-void load(
-    array& out,
-    size_t offset,
-    const std::shared_ptr<io::Reader>& reader,
-    bool swap_endianess);
-
-} // namespace mlx::core

mlx/backend/common/slicing.cpp

@@ -14,6 +14,10 @@ std::tuple<int64_t, Strides> prepare_slice(
     data_offset += start_indices[i] * in.strides()[i];
     inp_strides[i] = in.strides()[i] * strides[i];
   }
+  // Normalize the offset
+  if (data_offset < 0) {
+    data_offset += in.data_size();
+  }
   return std::make_tuple(data_offset, inp_strides);
 }
 
@@ -32,7 +36,7 @@ void shared_buffer_slice(
   flags.col_contiguous = is_col_contiguous;
   flags.contiguous = (no_bsx_size == data_size);
 
-  move_or_copy(in, out, out_strides, flags, data_size, data_offset);
+  out.copy_shared_buffer(in, out_strides, flags, data_size, data_offset);
 }
 
 void slice(
@@ -54,9 +58,10 @@ void slice(
       data_end += end_idx * in.strides()[i];
     }
   }
-  // data_end can be -1
-  size_t data_size =
-      data_end < 0 ? (data_offset - data_end) : (data_end - data_offset);
+  if (data_end < 0) {
+    data_end += in.data_size();
+  }
+  size_t data_size = (data_end - data_offset);
   shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
 }
 

mlx/backend/common/ternary.h

@@ -36,15 +36,10 @@ inline void set_ternary_op_output_data(
     const array& b,
     const array& c,
     array& out,
-    TernaryOpType topt,
-    bool donate_with_move = false) {
-  auto maybe_donate = [&out, donate_with_move](const array& x) {
+    TernaryOpType topt) {
+  auto maybe_donate = [&out](const array& x) {
     if (is_donatable(x, out)) {
-      if (donate_with_move) {
-        out.move_shared_buffer(x);
-      } else {
-        out.copy_shared_buffer(x);
-      }
+      out.copy_shared_buffer(x);
       return true;
     }
     return false;

mlx/backend/common/utils.cpp

@@ -4,28 +4,6 @@
 
 namespace mlx::core {
 
-void move_or_copy(const array& in, array& out) {
-  if (in.is_donatable()) {
-    out.move_shared_buffer(in);
-  } else {
-    out.copy_shared_buffer(in);
-  }
-}
-
-void move_or_copy(
-    const array& in,
-    array& out,
-    const Strides& strides,
-    array::Flags flags,
-    size_t data_size,
-    size_t offset /* = 0 */) {
-  if (in.is_donatable()) {
-    out.move_shared_buffer(in, strides, flags, data_size, offset);
-  } else {
-    out.copy_shared_buffer(in, strides, flags, data_size, offset);
-  }
-}
-
 std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
     const Shape& shape,
     const std::vector<Strides>& strides,

mlx/backend/common/utils.h

@@ -159,15 +159,6 @@ inline bool is_donatable(const array& in, const array& out) {
       in.buffer_size() <= out.nbytes() + donation_extra;
 }
 
-void move_or_copy(const array& in, array& out);
-void move_or_copy(
-    const array& in,
-    array& out,
-    const Strides& strides,
-    array::Flags flags,
-    size_t data_size,
-    size_t offset = 0);
-
 std::pair<bool, Strides> prepare_reshape(const array& in, const array& out);
 
 void shared_buffer_reshape(

mlx/backend/cpu/CMakeLists.txt

@@ -44,7 +44,9 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/eigh.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/encoder.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
@@ -65,6 +67,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/inverse.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/unary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
           ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp)
 
 if(MLX_BUILD_ACCELERATE)

mlx/backend/cpu/arange.h

@@ -2,76 +2,27 @@
 
 #pragma once
 
-#include "mlx/allocator.h"
 #include "mlx/array.h"
+#include "mlx/backend/cpu/encoder.h"
 
 namespace mlx::core {
 
 namespace {
 
 template <typename T>
-void arange(T start, T next, array& out, size_t size) {
+void arange(T start, T next, array& out, size_t size, Stream stream) {
   auto ptr = out.data<T>();
   auto step_size = next - start;
-  for (int i = 0; i < size; ++i) {
-    ptr[i] = start;
-    start += step_size;
-  }
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([ptr, start, step_size, size]() mutable {
+    for (int i = 0; i < size; ++i) {
+      ptr[i] = start;
+      start += step_size;
+    }
+  });
 }
 
 } // namespace
 
-void arange(
-    const std::vector<array>& inputs,
-    array& out,
-    double start,
-    double step) {
-  assert(inputs.size() == 0);
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  switch (out.dtype()) {
-    case bool_:
-      throw std::runtime_error("Bool type unsupported for arange.");
-      break;
-    case uint8:
-      arange<uint8_t>(start, start + step, out, out.size());
-      break;
-    case uint16:
-      arange<uint16_t>(start, start + step, out, out.size());
-      break;
-    case uint32:
-      arange<uint32_t>(start, start + step, out, out.size());
-      break;
-    case uint64:
-      arange<uint64_t>(start, start + step, out, out.size());
-      break;
-    case int8:
-      arange<int8_t>(start, start + step, out, out.size());
-      break;
-    case int16:
-      arange<int16_t>(start, start + step, out, out.size());
-      break;
-    case int32:
-      arange<int32_t>(start, start + step, out, out.size());
-      break;
-    case int64:
-      arange<int64_t>(start, start + step, out, out.size());
-      break;
-    case float16:
-      arange<float16_t>(start, start + step, out, out.size());
-      break;
-    case float32:
-      arange<float>(start, start + step, out, out.size());
-      break;
-    case float64:
-      arange<double>(start, start + step, out, out.size());
-      break;
-    case bfloat16:
-      arange<bfloat16_t>(start, start + step, out, out.size());
-      break;
-    case complex64:
-      arange<complex64_t>(start, start + step, out, out.size());
-      break;
-  }
-}
-
 } // namespace mlx::core

mlx/backend/cpu/arg_reduce.cpp

@@ -3,6 +3,7 @@
 #include <cassert>
 
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
@@ -10,23 +11,43 @@ namespace mlx::core {
 namespace {
 
 template <typename InT, typename OpT>
-void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
+void arg_reduce(
+    const array& in,
+    array& out,
+    const OpT& op,
+    int axis,
+    Stream stream) {
   auto axis_size = in.shape()[axis];
   auto axis_stride = in.strides()[axis];
   Strides strides = in.strides();
   Shape shape = in.shape();
   strides.erase(strides.begin() + axis);
   shape.erase(shape.begin() + axis);
-  for (uint32_t i = 0; i < out.size(); ++i) {
-    auto loc = elem_to_loc(i, shape, strides);
-    auto in_ptr = in.data<InT>() + loc;
-    uint32_t ind_v = 0;
-    InT v = (*in_ptr);
-    for (uint32_t j = 0; j < axis_size; ++j, in_ptr += axis_stride) {
-      op(j, (*in_ptr), &ind_v, &v);
+  auto in_ptr = in.data<InT>();
+  auto out_ptr = out.data<uint32_t>();
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in_ptr,
+                    out_ptr,
+                    axis_size,
+                    axis_stride,
+                    op = std::move(op),
+                    shape = std::move(shape),
+                    strides = std::move(strides),
+                    size = out.size()]() {
+    for (uint32_t i = 0; i < size; ++i) {
+      auto loc = elem_to_loc(i, shape, strides);
+      auto local_in_ptr = in_ptr + loc;
+      uint32_t ind_v = 0;
+      InT v = (*local_in_ptr);
+      for (uint32_t j = 0; j < axis_size; ++j, local_in_ptr += axis_stride) {
+        op(j, (*local_in_ptr), &ind_v, &v);
+      }
+      out_ptr[i] = ind_v;
     }
-    out.data<uint32_t>()[i] = ind_v;
-  }
+  });
 }
 
 template <typename InT>
@@ -34,7 +55,8 @@ void arg_reduce_dispatch(
     const array& in,
     array& out,
     ArgReduce::ReduceType rtype,
-    int axis) {
+    int axis,
+    Stream stream) {
   switch (rtype) {
     case ArgReduce::ArgMin: {
       auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
@@ -43,7 +65,7 @@ void arg_reduce_dispatch(
           (*ind_y) = ind_x;
         }
       };
-      arg_reduce<InT>(in, out, op, axis);
+      arg_reduce<InT>(in, out, op, axis, stream);
       break;
     }
     case ArgReduce::ArgMax: {
@@ -53,7 +75,7 @@ void arg_reduce_dispatch(
           (*ind_y) = ind_x;
         }
       };
-      arg_reduce<InT>(in, out, op, axis);
+      arg_reduce<InT>(in, out, op, axis, stream);
       break;
     }
   }
@@ -68,46 +90,46 @@ void ArgReduce::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   switch (in.dtype()) {
     case bool_:
-      arg_reduce_dispatch<bool>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<bool>(in, out, reduce_type_, axis_, stream());
       break;
     case uint8:
-      arg_reduce_dispatch<uint8_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<uint8_t>(in, out, reduce_type_, axis_, stream());
       break;
     case uint16:
-      arg_reduce_dispatch<uint16_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<uint16_t>(in, out, reduce_type_, axis_, stream());
       break;
     case uint32:
-      arg_reduce_dispatch<uint32_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<uint32_t>(in, out, reduce_type_, axis_, stream());
       break;
     case uint64:
-      arg_reduce_dispatch<uint64_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<uint64_t>(in, out, reduce_type_, axis_, stream());
       break;
     case int8:
-      arg_reduce_dispatch<int8_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<int8_t>(in, out, reduce_type_, axis_, stream());
       break;
     case int16:
-      arg_reduce_dispatch<int16_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<int16_t>(in, out, reduce_type_, axis_, stream());
       break;
     case int32:
-      arg_reduce_dispatch<int32_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<int32_t>(in, out, reduce_type_, axis_, stream());
       break;
     case int64:
-      arg_reduce_dispatch<int64_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<int64_t>(in, out, reduce_type_, axis_, stream());
       break;
     case float16:
-      arg_reduce_dispatch<float16_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<float16_t>(in, out, reduce_type_, axis_, stream());
       break;
     case float32:
-      arg_reduce_dispatch<float>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<float>(in, out, reduce_type_, axis_, stream());
       break;
     case bfloat16:
-      arg_reduce_dispatch<bfloat16_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<bfloat16_t>(in, out, reduce_type_, axis_, stream());
       break;
     case float64:
-      arg_reduce_dispatch<double>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<double>(in, out, reduce_type_, axis_, stream());
       break;
     case complex64:
-      arg_reduce_dispatch<complex64_t>(in, out, reduce_type_, axis_);
+      arg_reduce_dispatch<complex64_t>(in, out, reduce_type_, axis_, stream());
       break;
   }
 }

mlx/backend/cpu/binary.cpp

@@ -16,49 +16,49 @@ namespace mlx::core {
 namespace {
 
 template <typename Op>
-void comparison_op(const array& a, const array& b, array& out, Op op) {
+void comparison_op(const array& a, const array& b, array& out) {
   switch (a.dtype()) {
     case bool_:
-      binary_op<bool, bool>(a, b, out, op);
+      binary_op<bool, bool, Op>(a, b, out);
       break;
     case uint8:
-      binary_op<uint8_t, bool>(a, b, out, op);
+      binary_op<uint8_t, bool, Op>(a, b, out);
       break;
     case uint16:
-      binary_op<uint16_t, bool>(a, b, out, op);
+      binary_op<uint16_t, bool, Op>(a, b, out);
       break;
     case uint32:
-      binary_op<uint32_t, bool>(a, b, out, op);
+      binary_op<uint32_t, bool, Op>(a, b, out);
       break;
     case uint64:
-      binary_op<uint64_t, bool>(a, b, out, op);
+      binary_op<uint64_t, bool, Op>(a, b, out);
       break;
     case int8:
-      binary_op<int8_t, bool>(a, b, out, op);
+      binary_op<int8_t, bool, Op>(a, b, out);
       break;
     case int16:
-      binary_op<int16_t, bool>(a, b, out, op);
+      binary_op<int16_t, bool, Op>(a, b, out);
       break;
     case int32:
-      binary_op<int32_t, bool>(a, b, out, op);
+      binary_op<int32_t, bool, Op>(a, b, out);
       break;
     case int64:
-      binary_op<int64_t, bool>(a, b, out, op);
+      binary_op<int64_t, bool, Op>(a, b, out);
       break;
     case float16:
-      binary_op<float16_t, bool>(a, b, out, op);
+      binary_op<float16_t, bool, Op>(a, b, out);
       break;
     case float32:
-      binary_op<float, bool>(a, b, out, op);
+      binary_op<float, bool, Op>(a, b, out);
       break;
     case float64:
-      binary_op<double, bool>(a, b, out, op);
+      binary_op<double, bool, Op>(a, b, out);
       break;
     case bfloat16:
-      binary_op<bfloat16_t, bool>(a, b, out, op);
+      binary_op<bfloat16_t, bool, Op>(a, b, out);
       break;
     case complex64:
-      binary_op<complex64_t, bool>(a, b, out, op);
+      binary_op<complex64_t, bool, Op>(a, b, out);
       break;
   }
 }
@@ -151,47 +151,47 @@ void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
   if (equal_nan_) {
     switch (a.dtype()) {
       case float16:
-        binary_op<float16_t, bool>(a, b, out, detail::NaNEqual());
+        binary_op<float16_t, bool, detail::NaNEqual>(a, b, out);
         break;
       case float32:
-        binary_op<float, bool>(a, b, out, detail::NaNEqual());
+        binary_op<float, bool, detail::NaNEqual>(a, b, out);
         break;
       case float64:
-        binary_op<double, bool>(a, b, out, detail::NaNEqual());
+        binary_op<double, bool, detail::NaNEqual>(a, b, out);
         break;
       case bfloat16:
-        binary_op<bfloat16_t, bool>(a, b, out, detail::NaNEqual());
+        binary_op<bfloat16_t, bool, detail::NaNEqual>(a, b, out);
         break;
       case complex64:
-        binary_op<complex64_t, bool>(a, b, out, detail::NaNEqual());
+        binary_op<complex64_t, bool, detail::NaNEqual>(a, b, out);
         break;
       default:
         throw std::runtime_error(
             "[NanEqual::eval_cpu] Only for floating point types.");
     }
   } else {
-    comparison_op(a, b, out, detail::Equal());
+    comparison_op<detail::Equal>(a, b, out);
   }
 }
 
 void Greater::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::Greater());
+  comparison_op<detail::Greater>(inputs[0], inputs[1], out);
 }
 
 void GreaterEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::GreaterEqual());
+  comparison_op<detail::GreaterEqual>(inputs[0], inputs[1], out);
 }
 
 void Less::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::Less());
+  comparison_op<detail::Less>(inputs[0], inputs[1], out);
 }
 
 void LessEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::LessEqual());
+  comparison_op<detail::LessEqual>(inputs[0], inputs[1], out);
 }
 
 void LogAddExp::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -200,16 +200,16 @@ void LogAddExp::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
   switch (out.dtype()) {
     case float16:
-      binary_op<float16_t>(a, b, out, detail::LogAddExp());
+      binary_op<float16_t, detail::LogAddExp>(a, b, out);
       break;
     case float32:
-      binary_op<float>(a, b, out, detail::LogAddExp());
+      binary_op<float, detail::LogAddExp>(a, b, out);
       break;
     case float64:
-      binary_op<double>(a, b, out, detail::LogAddExp());
+      binary_op<double, detail::LogAddExp>(a, b, out);
       break;
     case bfloat16:
-      binary_op<bfloat16_t>(a, b, out, detail::LogAddExp());
+      binary_op<bfloat16_t, detail::LogAddExp>(a, b, out);
       break;
     default:
       throw std::runtime_error(
@@ -254,7 +254,7 @@ void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
 
 void NotEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::NotEqual());
+  comparison_op<detail::NotEqual>(inputs[0], inputs[1], out);
 }
 
 void Power::eval_cpu(const std::vector<array>& inputs, array& out) {

mlx/backend/cpu/binary.h

@@ -7,6 +7,8 @@
 #include "mlx/array.h"
 #include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
 
 #include "mlx/backend/cpu/simd/simd.h"
 
@@ -14,22 +16,18 @@ namespace mlx::core {
 
 template <typename Op>
 struct VectorScalar {
-  Op op;
-
-  VectorScalar(Op op_) : op(op_) {}
-
   template <typename T, typename U>
   void operator()(const T* a, const T* b, U* dst, int size) {
     T scalar = *b;
     constexpr int N = simd::max_size<T>;
     while (size >= N) {
-      simd::store(dst, op(simd::load<T, N>(a), simd::Simd<T, N>(scalar)));
+      simd::store(dst, Op{}(simd::load<T, N>(a), simd::Simd<T, N>(scalar)));
       dst += N;
       a += N;
       size -= N;
     }
     while (size-- > 0) {
-      *dst = op(*a, scalar);
+      *dst = Op{}(*a, scalar);
       dst++;
       a++;
     }
@@ -38,22 +36,18 @@ struct VectorScalar {
 
 template <typename Op>
 struct ScalarVector {
-  Op op;
-
-  ScalarVector(Op op_) : op(op_) {}
-
   template <typename T, typename U>
   void operator()(const T* a, const T* b, U* dst, int size) {
     T scalar = *a;
     constexpr int N = simd::max_size<T>;
     while (size >= N) {
-      simd::store(dst, op(simd::Simd<T, N>(scalar), simd::load<T, N>(b)));
+      simd::store(dst, Op{}(simd::Simd<T, N>(scalar), simd::load<T, N>(b)));
       dst += N;
       b += N;
       size -= N;
     }
     while (size-- > 0) {
-      *dst = op(scalar, *b);
+      *dst = Op{}(scalar, *b);
       dst++;
       b++;
     }
@@ -62,22 +56,18 @@ struct ScalarVector {
 
 template <typename Op>
 struct VectorVector {
-  Op op;
-
-  VectorVector(Op op_) : op(op_) {}
-
   template <typename T, typename U>
   void operator()(const T* a, const T* b, U* dst, int size) {
     constexpr int N = simd::max_size<T>;
     while (size >= N) {
-      simd::store(dst, op(simd::load<T, N>(a), simd::load<T, N>(b)));
+      simd::store(dst, Op{}(simd::load<T, N>(a), simd::load<T, N>(b)));
       dst += N;
       a += N;
       b += N;
       size -= N;
     }
     while (size-- > 0) {
-      *dst = op(*a, *b);
+      *dst = Op{}(*a, *b);
       dst++;
       a++;
       b++;
@@ -90,7 +80,6 @@ void binary_op_dims(
     const T* a,
     const T* b,
     U* out,
-    Op op,
     const Shape& shape,
     const Strides& a_strides,
     const Strides& b_strides,
@@ -104,12 +93,12 @@ void binary_op_dims(
   for (int i = 0; i < N; i++) {
     if constexpr (D > 1) {
       binary_op_dims<T, U, Op, D - 1, Strided>(
-          a, b, out, op, shape, a_strides, b_strides, out_strides, axis + 1);
+          a, b, out, shape, a_strides, b_strides, out_strides, axis + 1);
     } else {
       if constexpr (Strided) {
-        op(a, b, out, stride_out);
+        Op{}(a, b, out, stride_out);
       } else {
-        *out = op(*a, *b);
+        *out = Op{}(*a, *b);
       }
     }
     out += stride_out;
@@ -120,66 +109,38 @@ void binary_op_dims(
 
 template <typename T, typename U, bool Strided, typename Op>
 void binary_op_dispatch_dims(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
+    const T* a,
+    const T* b,
+    U* out,
     int dim,
+    int size,
     const Shape& shape,
     const Strides& a_strides,
     const Strides& b_strides,
     const Strides& out_strides) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* out_ptr = out.data<U>();
   switch (dim) {
     case 1:
       binary_op_dims<T, U, Op, 1, Strided>(
-          a_ptr,
-          b_ptr,
-          out_ptr,
-          op,
-          shape,
-          a_strides,
-          b_strides,
-          out_strides,
-          0);
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
       return;
     case 2:
       binary_op_dims<T, U, Op, 2, Strided>(
-          a_ptr,
-          b_ptr,
-          out_ptr,
-          op,
-          shape,
-          a_strides,
-          b_strides,
-          out_strides,
-          0);
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
       return;
     case 3:
       binary_op_dims<T, U, Op, 3, Strided>(
-          a_ptr,
-          b_ptr,
-          out_ptr,
-          op,
-          shape,
-          a_strides,
-          b_strides,
-          out_strides,
-          0);
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
       return;
   }
 
   ContiguousIterator a_it(shape, a_strides, dim - 3);
   ContiguousIterator b_it(shape, b_strides, dim - 3);
   auto stride = out_strides[dim - 4];
-  for (int64_t elem = 0; elem < a.size(); elem += stride) {
+  for (int64_t elem = 0; elem < size; elem += stride) {
     binary_op_dims<T, U, Op, 3, Strided>(
-        a_ptr + a_it.loc,
-        b_ptr + b_it.loc,
-        out_ptr + elem,
-        op,
+        a + a_it.loc,
+        b + b_it.loc,
+        out + elem,
         shape,
         a_strides,
         b_strides,
@@ -191,181 +152,216 @@ void binary_op_dispatch_dims(
 }
 
 template <typename T, typename U, typename Op>
-void binary_op(const array& a, const array& b, array& out, Op op) {
+void binary_op(const array& a, const array& b, array& out) {
   auto bopt = get_binary_op_type(a, b);
   set_binary_op_output_data(a, b, out, bopt);
 
   // The full computation is scalar scalar so call the base op once
-  if (bopt == BinaryOpType::ScalarScalar) {
-    *(out.data<U>()) = op(*a.data<T>(), *b.data<T>());
-    return;
-  }
+  auto a_ptr = a.data<T>();
+  auto b_ptr = b.data<T>();
 
-  // The full computation is scalar vector so delegate to the op
-  if (bopt == BinaryOpType::ScalarVector) {
-    ScalarVector{op}(a.data<T>(), b.data<T>(), out.data<U>(), b.data_size());
-    return;
-  }
-
-  // The full computation is vector scalar so delegate to the op
-  if (bopt == BinaryOpType::VectorScalar) {
-    VectorScalar{op}(a.data<T>(), b.data<T>(), out.data<U>(), a.data_size());
-    return;
-  }
-
-  // The full computation is vector vector so delegate to the op
-  if (bopt == BinaryOpType::VectorVector) {
-    VectorVector{op}(a.data<T>(), b.data<T>(), out.data<U>(), out.size());
-    return;
-  }
-
-  // General computation so let's try to optimize
-  auto [new_shape, new_strides] = collapse_contiguous_dims(
-      a.shape(), {a.strides(), b.strides(), out.strides()});
-  const auto& a_strides = new_strides[0];
-  const auto& b_strides = new_strides[1];
-  const auto& strides = new_strides[2];
-
-  // Get the left-most dim such that the array is row contiguous after
-  auto leftmost_rc_dim = [&strides](const auto& arr_strides) {
-    int d = arr_strides.size() - 1;
-    for (; d >= 0 && arr_strides[d] == strides[d]; d--) {
+  auto out_ptr = out.data<U>();
+  auto& encoder = cpu::get_command_encoder(out.primitive().stream());
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([bopt,
+                    a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_data_size = a.data_size(),
+                    b_data_size = b.data_size(),
+                    size = a.size(),
+                    shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_strides = b.strides(),
+                    strides = out.strides()]() mutable {
+    if (bopt == BinaryOpType::ScalarScalar) {
+      *out_ptr = Op{}(*a_ptr, *b_ptr);
+      return;
     }
-    return d + 1;
-  };
-  auto a_rc_dim = leftmost_rc_dim(a_strides);
-  auto b_rc_dim = leftmost_rc_dim(b_strides);
 
-  // Get the left-most dim such that the array is a broadcasted "scalar" after
-  auto leftmost_s_dim = [](const auto& arr_strides) {
-    int d = arr_strides.size() - 1;
-    for (; d >= 0 && arr_strides[d] == 0; d--) {
+    // The full computation is scalar vector so delegate to the op
+    if (bopt == BinaryOpType::ScalarVector) {
+      ScalarVector<Op>{}(a_ptr, b_ptr, out_ptr, b_data_size);
+      return;
     }
-    return d + 1;
-  };
-  auto a_s_dim = leftmost_s_dim(a_strides);
-  auto b_s_dim = leftmost_s_dim(b_strides);
 
-  auto ndim = new_shape.size();
+    // The full computation is vector scalar so delegate to the op
+    if (bopt == BinaryOpType::VectorScalar) {
+      VectorScalar<Op>{}(a_ptr, b_ptr, out_ptr, a_data_size);
+      return;
+    }
 
-  // Case 1: LxM and FxM where L and F are broadcastable and M is row contiguous
-  int dim = ndim;
-  if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
-    bopt = BinaryOpType::VectorVector;
-    dim = d;
-    // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
+    // The full computation is vector vector so delegate to the op
+    if (bopt == BinaryOpType::VectorVector) {
+      VectorVector<Op>{}(a_ptr, b_ptr, out_ptr, size);
+      return;
+    }
+
+    // General computation so let's try to optimize
+    auto [new_shape, new_strides] = collapse_contiguous_dims(
+        shape,
+        {std::move(a_strides), std::move(b_strides), std::move(strides)});
+    a_strides = new_strides[0];
+    b_strides = new_strides[1];
+    strides = new_strides[2];
+
+    // Get the left-most dim such that the array is row contiguous after
+    auto leftmost_rc_dim = [&strides](const auto& arr_strides) {
+      int d = arr_strides.size() - 1;
+      for (; d >= 0 && arr_strides[d] == strides[d]; d--) {
+      }
+      return d + 1;
+    };
+    auto a_rc_dim = leftmost_rc_dim(a_strides);
+    auto b_rc_dim = leftmost_rc_dim(b_strides);
+
+    // Get the left-most dim such that the array is a broadcasted "scalar" after
+    auto leftmost_s_dim = [](const auto& arr_strides) {
+      int d = arr_strides.size() - 1;
+      for (; d >= 0 && arr_strides[d] == 0; d--) {
+      }
+      return d + 1;
+    };
+    auto a_s_dim = leftmost_s_dim(a_strides);
+    auto b_s_dim = leftmost_s_dim(b_strides);
+
+    auto ndim = new_shape.size();
+
+    // Case 1: LxM and FxM where L and F are broadcastable and M is row
     // contiguous
-  } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
-    bopt = BinaryOpType::VectorScalar;
-    dim = d;
-    // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
-    // contiguous
-  } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
-    bopt = BinaryOpType::ScalarVector;
-    dim = d;
-  }
+    int dim = ndim;
+    if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
+      bopt = BinaryOpType::VectorVector;
+      dim = d;
+      // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
+      // contiguous
+    } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
+      bopt = BinaryOpType::VectorScalar;
+      dim = d;
+      // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
+      // contiguous
+    } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
+      bopt = BinaryOpType::ScalarVector;
+      dim = d;
+    }
 
-  // Can be sure dim > 0 since otherwise we would have used one of the fully
-  // contiguous methods above. Except for the case that the flags do not
-  // correspond to the underlying contiguity.
-  if (dim == 0 || strides[dim - 1] < 16) {
-    bopt = BinaryOpType::General;
-    dim = ndim;
-  }
+    // Can be sure dim > 0 since otherwise we would have used one of the fully
+    // contiguous methods above. Except for the case that the flags do not
+    // correspond to the underlying contiguity.
+    if (dim == 0 || strides[dim - 1] < 16) {
+      bopt = BinaryOpType::General;
+      dim = ndim;
+    }
 
-  switch (bopt) {
-    case BinaryOpType::VectorVector:
-      binary_op_dispatch_dims<T, U, true>(
-          a,
-          b,
-          out,
-          VectorVector{op},
-          dim,
-          new_shape,
-          a_strides,
-          b_strides,
-          strides);
-      break;
-    case BinaryOpType::VectorScalar:
-      binary_op_dispatch_dims<T, U, true>(
-          a,
-          b,
-          out,
-          VectorScalar{op},
-          dim,
-          new_shape,
-          a_strides,
-          b_strides,
-          strides);
-      break;
-    case BinaryOpType::ScalarVector:
-      binary_op_dispatch_dims<T, U, true>(
-          a,
-          b,
-          out,
-          ScalarVector{op},
-          dim,
-          new_shape,
-          a_strides,
-          b_strides,
-          strides);
-      break;
-    default:
-      binary_op_dispatch_dims<T, U, false>(
-          a, b, out, op, dim, new_shape, a_strides, b_strides, strides);
-      break;
-  }
+    switch (bopt) {
+      case BinaryOpType::VectorVector:
+        binary_op_dispatch_dims<T, U, true, VectorVector<Op>>(
+            a_ptr,
+            b_ptr,
+            out_ptr,
+            dim,
+            size,
+            new_shape,
+            a_strides,
+            b_strides,
+            strides);
+        break;
+      case BinaryOpType::VectorScalar:
+        binary_op_dispatch_dims<T, U, true, VectorScalar<Op>>(
+            a_ptr,
+            b_ptr,
+            out_ptr,
+            dim,
+            size,
+            new_shape,
+            a_strides,
+            b_strides,
+            strides);
+        break;
+      case BinaryOpType::ScalarVector:
+        binary_op_dispatch_dims<T, U, true, ScalarVector<Op>>(
+            a_ptr,
+            b_ptr,
+            out_ptr,
+            dim,
+            size,
+            new_shape,
+            a_strides,
+            b_strides,
+            strides);
+        break;
+      default:
+        binary_op_dispatch_dims<T, U, false, Op>(
+            a_ptr,
+            b_ptr,
+            out_ptr,
+            dim,
+            size,
+            new_shape,
+            a_strides,
+            b_strides,
+            strides);
+        break;
+    }
+  });
+}
+
+template <typename T, typename Op>
+void binary_op(const array& a, const array& b, array& out) {
+  binary_op<T, T, Op>(a, b, out);
 }
 
 template <typename T, typename Op>
 void binary_op(const array& a, const array& b, array& out, Op op) {
-  binary_op<T, T>(a, b, out, op);
+  binary_op<T, T, Op>(a, b, out);
 }
 
 template <typename Op>
 void binary(const array& a, const array& b, array& out, Op op) {
   switch (out.dtype()) {
     case bool_:
-      binary_op<bool>(a, b, out, op);
+      binary_op<bool, Op>(a, b, out);
       break;
     case uint8:
-      binary_op<uint8_t>(a, b, out, op);
+      binary_op<uint8_t, Op>(a, b, out);
       break;
     case uint16:
-      binary_op<uint16_t>(a, b, out, op);
+      binary_op<uint16_t, Op>(a, b, out);
       break;
     case uint32:
-      binary_op<uint32_t>(a, b, out, op);
+      binary_op<uint32_t, Op>(a, b, out);
       break;
     case uint64:
-      binary_op<uint64_t>(a, b, out, op);
+      binary_op<uint64_t, Op>(a, b, out);
       break;
     case int8:
-      binary_op<int8_t>(a, b, out, op);
+      binary_op<int8_t, Op>(a, b, out);
       break;
     case int16:
-      binary_op<int16_t>(a, b, out, op);
+      binary_op<int16_t, Op>(a, b, out);
       break;
     case int32:
-      binary_op<int32_t>(a, b, out, op);
+      binary_op<int32_t, Op>(a, b, out);
       break;
     case int64:
-      binary_op<int64_t>(a, b, out, op);
+      binary_op<int64_t, Op>(a, b, out);
       break;
     case float16:
-      binary_op<float16_t>(a, b, out, op);
+      binary_op<float16_t, Op>(a, b, out);
       break;
     case float32:
-      binary_op<float>(a, b, out, op);
+      binary_op<float, Op>(a, b, out);
       break;
     case float64:
-      binary_op<double>(a, b, out, op);
+      binary_op<double, Op>(a, b, out);
       break;
     case bfloat16:
-      binary_op<bfloat16_t>(a, b, out, op);
+      binary_op<bfloat16_t, Op>(a, b, out);
       break;
     case complex64:
-      binary_op<complex64_t>(a, b, out, op);
+      binary_op<complex64_t, Op>(a, b, out);
       break;
   }
 }

mlx/backend/cpu/binary_two.h

@@ -4,6 +4,8 @@
 
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/binary.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
 
 namespace mlx::core {
 
@@ -55,65 +57,81 @@ void binary_op_dispatch_dims(
     const array& b,
     array& out_a,
     array& out_b,
+    Stream stream,
     Op op) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out_a);
+  encoder.set_output_array(out_b);
+
   auto [shape, strides] = collapse_contiguous_dims(
       a.shape(), {a.strides(), b.strides(), out_a.strides()});
-  const auto& a_strides = strides[0];
-  const auto& b_strides = strides[1];
-  const auto& out_strides = strides[2];
   const T* a_ptr = a.data<T>();
   const T* b_ptr = b.data<T>();
   U* out_a_ptr = out_a.data<U>();
   U* out_b_ptr = out_b.data<U>();
 
-  int ndim = shape.size();
-  switch (ndim) {
-    case 1:
-      binary_op_dims<T, U, Op, 1>(
-          a_ptr,
-          b_ptr,
-          out_a_ptr,
-          out_b_ptr,
-          op,
-          shape,
-          a_strides,
-          b_strides,
-          out_strides,
-          0);
-      return;
-    case 2:
-      binary_op_dims<T, U, Op, 2>(
-          a_ptr,
-          b_ptr,
-          out_a_ptr,
-          out_b_ptr,
-          op,
-          shape,
-          a_strides,
-          b_strides,
-          out_strides,
-          0);
-      return;
-  }
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_a_ptr,
+                    out_b_ptr,
+                    size = a.size(),
+                    shape = std::move(shape),
+                    strides = std::move(strides),
+                    op = std::move(op)]() {
+    const auto& a_strides = strides[0];
+    const auto& b_strides = strides[1];
+    const auto& out_strides = strides[2];
+    int ndim = shape.size();
+    switch (ndim) {
+      case 1:
+        binary_op_dims<T, U, Op, 1>(
+            a_ptr,
+            b_ptr,
+            out_a_ptr,
+            out_b_ptr,
+            op,
+            shape,
+            a_strides,
+            b_strides,
+            out_strides,
+            0);
+        return;
+      case 2:
+        binary_op_dims<T, U, Op, 2>(
+            a_ptr,
+            b_ptr,
+            out_a_ptr,
+            out_b_ptr,
+            op,
+            shape,
+            a_strides,
+            b_strides,
+            out_strides,
+            0);
+        return;
+    }
 
-  ContiguousIterator a_it(shape, a_strides, ndim - 2);
-  ContiguousIterator b_it(shape, b_strides, ndim - 2);
-  auto stride = out_strides[ndim - 3];
-  for (size_t elem = 0; elem < a.size(); elem += stride) {
-    binary_op_dims<T, U, Op, 2>(
-        a_ptr + a_it.loc,
-        b_ptr + b_it.loc,
-        out_a_ptr + elem,
-        out_b_ptr + elem,
-        op,
-        shape,
-        a_strides,
-        b_strides,
-        out_strides,
-        ndim - 2);
-    a_it.step();
-    b_it.step();
-  }
+    ContiguousIterator a_it(shape, a_strides, ndim - 2);
+    ContiguousIterator b_it(shape, b_strides, ndim - 2);
+    auto stride = out_strides[ndim - 3];
+    for (size_t elem = 0; elem < size; elem += stride) {
+      binary_op_dims<T, U, Op, 2>(
+          a_ptr + a_it.loc,
+          b_ptr + b_it.loc,
+          out_a_ptr + elem,
+          out_b_ptr + elem,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          ndim - 2);
+      a_it.step();
+      b_it.step();
+    }
+  });
 }
 
 template <typename T, typename U = T, typename Op>
@@ -128,40 +146,71 @@ void binary_op(
   set_binary_op_output_data(a, b, out_a, bopt);
   set_binary_op_output_data(a, b, out_b, bopt);
 
+  auto stream = out_a.primitive().stream();
   // The full computation is scalar scalar so call the base op once
   if (bopt == BinaryOpType::General) {
-    binary_op_dispatch_dims<T, U, Op>(a, b, out_a, out_b, op);
+    binary_op_dispatch_dims<T, U, Op>(a, b, out_a, out_b, stream, op);
     return;
   }
 
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out_a);
+  encoder.set_output_array(out_b);
+
   auto a_ptr = a.data<T>();
   auto b_ptr = b.data<T>();
   auto out_a_ptr = out_a.data<U>();
   auto out_b_ptr = out_b.data<U>();
   if (bopt == BinaryOpType::ScalarScalar) {
-    std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+    encoder.dispatch(
+        [a_ptr, b_ptr, out_a_ptr, out_b_ptr, op = std::move(op)]() mutable {
+          std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+        });
   } else if (bopt == BinaryOpType::ScalarVector) {
-    for (size_t i = 0; i < b.size(); ++i) {
-      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
-      out_a_ptr++;
-      out_b_ptr++;
-      b_ptr++;
-    }
+    encoder.dispatch([a_ptr,
+                      b_ptr,
+                      out_a_ptr,
+                      out_b_ptr,
+                      size = b.size(),
+                      op = std::move(op)]() mutable {
+      for (size_t i = 0; i < size; ++i) {
+        std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+        out_a_ptr++;
+        out_b_ptr++;
+        b_ptr++;
+      }
+    });
   } else if (bopt == BinaryOpType::VectorScalar) {
-    for (size_t i = 0; i < a.size(); ++i) {
-      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
-      out_a_ptr++;
-      out_b_ptr++;
-      a_ptr++;
-    }
+    encoder.dispatch([a_ptr,
+                      b_ptr,
+                      out_a_ptr,
+                      out_b_ptr,
+                      size = a.size(),
+                      op = std::move(op)]() mutable {
+      for (size_t i = 0; i < size; ++i) {
+        std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+        out_a_ptr++;
+        out_b_ptr++;
+        a_ptr++;
+      }
+    });
   } else { // VectorVector
-    for (size_t i = 0; i < a.size(); ++i) {
-      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
-      out_a_ptr++;
-      out_b_ptr++;
-      a_ptr++;
-      b_ptr++;
-    }
+    encoder.dispatch([a_ptr,
+                      b_ptr,
+                      out_a_ptr,
+                      out_b_ptr,
+                      size = a.size(),
+                      op = std::move(op)]() mutable {
+      for (size_t i = 0; i < size; ++i) {
+        std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+        out_a_ptr++;
+        out_b_ptr++;
+        a_ptr++;
+        b_ptr++;
+      }
+    });
   }
 }
 

mlx/backend/cpu/cholesky.cpp

@@ -2,13 +2,15 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/linalg.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
-void cholesky_impl(const array& a, array& factor, bool upper) {
+template <typename T>
+void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) {
   // Lapack uses the column-major convention. We take advantage of the fact that
   // the matrix should be symmetric:
   //   (A)ᵀ = A
@@ -16,59 +18,68 @@ void cholesky_impl(const array& a, array& factor, bool upper) {
   // triangular matrix, so uplo is the opposite of what we would expect from
   // upper
 
-  char uplo = (upper) ? 'L' : 'U';
-
   // The decomposition is computed in place, so just copy the input to the
   // output.
   copy(
       a,
       factor,
-      a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
 
-  const int N = a.shape(-1);
-  const size_t num_matrices = a.size() / (N * N);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(factor);
+  encoder.dispatch([matrix = factor.data<T>(),
+                    upper,
+                    N = a.shape(-1),
+                    size = a.size()]() mutable {
+    char uplo = (upper) ? 'L' : 'U';
+    size_t num_matrices = size / (N * N);
+    for (int i = 0; i < num_matrices; i++) {
+      // Compute Cholesky factorization.
+      int info;
+      potrf<T>(
+          /* uplo = */ &uplo,
+          /* n = */ &N,
+          /* a = */ matrix,
+          /* lda = */ &N,
+          /* info = */ &info);
 
-  float* matrix = factor.data<float>();
-
-  for (int i = 0; i < num_matrices; i++) {
-    // Compute Cholesky factorization.
-    int info;
-    MLX_LAPACK_FUNC(spotrf)
-    (
-        /* uplo = */ &uplo,
-        /* n = */ &N,
-        /* a = */ matrix,
-        /* lda = */ &N,
-        /* info = */ &info);
-
-    // TODO: We do nothing when the matrix is not positive semi-definite
-    // because throwing an error would result in a crash. If we figure out how
-    // to catch errors from the implementation we should throw.
-    if (info < 0) {
-      std::stringstream msg;
-      msg << "[cholesky] Cholesky decomposition failed with error code "
-          << info;
-      throw std::runtime_error(msg.str());
-    }
-
-    // Zero out the upper/lower triangle while advancing the pointer to the
-    // next matrix at the same time.
-    for (int row = 0; row < N; row++) {
-      if (upper) {
-        std::fill(matrix, matrix + row, 0);
-      } else {
-        std::fill(matrix + row + 1, matrix + N, 0);
+      // TODO: We do nothing when the matrix is not positive semi-definite
+      // because throwing an error would result in a crash. If we figure out how
+      // to catch errors from the implementation we should throw.
+      if (info < 0) {
+        std::stringstream msg;
+        msg << "[Cholesky::eval_cpu] Cholesky decomposition failed with error code "
+            << info;
+        throw std::runtime_error(msg.str());
+      }
+
+      // Zero out the upper/lower triangle while advancing the pointer to the
+      // next matrix at the same time.
+      for (int row = 0; row < N; row++) {
+        if (upper) {
+          std::fill(matrix, matrix + row, 0);
+        } else {
+          std::fill(matrix + row + 1, matrix + N, 0);
+        }
+        matrix += N;
       }
-      matrix += N;
     }
-  }
+  });
 }
 
 void Cholesky::eval_cpu(const std::vector<array>& inputs, array& output) {
-  if (inputs[0].dtype() != float32) {
-    throw std::runtime_error("[Cholesky::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      cholesky_impl<float>(inputs[0], output, upper_, stream());
+      break;
+    case float64:
+      cholesky_impl<double>(inputs[0], output, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Cholesky::eval_cpu] only supports float32 or float64.");
   }
-  cholesky_impl(inputs[0], output, upper_);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/compiled.cpp

@@ -11,6 +11,7 @@
 
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/cpu/compiled_preamble.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/jit_compiler.h"
 #include "mlx/device.h"
 #include "mlx/graph_utils.h"
@@ -288,6 +289,7 @@ void Compiled::eval_cpu(
   // Figure out which kernel we are using
   auto& shape = outputs[0].shape();
   auto contiguous = compiled_check_contiguity(inputs, shape);
+  auto& encoder = cpu::get_command_encoder(stream());
 
   // Handle all broadcasting and collect function input arguments
   std::vector<void*> args;
@@ -298,6 +300,7 @@ void Compiled::eval_cpu(
       continue;
     }
     auto& x = inputs[i];
+    encoder.set_input_array(x);
     args.push_back((void*)x.data<void>());
 
     if (contiguous || is_scalar(x)) {
@@ -356,18 +359,25 @@ void Compiled::eval_cpu(
   });
 
   compiled_allocate_outputs(
-      inputs, outputs, inputs_, constant_ids_, contiguous, false);
+      inputs, outputs, inputs_, constant_ids_, contiguous);
 
   for (auto& x : outputs) {
     args.push_back(x.data<void>());
+    encoder.set_output_array(x);
   }
+  Shape out_shape;
   if (!contiguous) {
-    args.push_back((void*)outputs[0].shape().data());
+    out_shape = outputs[0].shape();
+    args.push_back((void*)out_shape.data());
   } else {
     args.push_back((void*)outputs[0].data_size());
   }
   auto fun = (void (*)(void**))fn_ptr;
-  fun(args.data());
+  encoder.dispatch(
+      [fun,
+       args = std::move(args),
+       strides = std::move(strides),
+       out_shape = std::move(out_shape)]() mutable { fun(args.data()); });
 }
 
 } // namespace mlx::core

mlx/backend/cpu/copy.cpp

@@ -5,6 +5,7 @@
 #include "mlx/allocator.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 
 namespace mlx::core {
@@ -12,20 +13,29 @@ namespace mlx::core {
 namespace {
 
 template <typename SrcT, typename DstT>
-void copy_single(const array& src, array& dst) {
-  auto val = static_cast<DstT>(src.data<SrcT>()[0]);
+void copy_single(const array& src, array& dst, Stream stream) {
+  auto src_ptr = src.data<SrcT>();
   auto dst_ptr = dst.data<DstT>();
-  for (int i = 0; i < dst.size(); ++i) {
-    dst_ptr[i] = val;
-  }
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_output_array(dst);
+  encoder.dispatch([src_ptr, dst_ptr, size = dst.size()]() {
+    auto val = static_cast<DstT>(src_ptr[0]);
+    std::fill_n(dst_ptr, size, val);
+  });
 }
 
 template <typename SrcT, typename DstT>
-void copy_vector(const array& src, array& dst) {
+void copy_vector(const array& src, array& dst, Stream stream) {
   auto src_ptr = src.data<SrcT>();
   auto dst_ptr = dst.data<DstT>();
   size_t size = src.data_size();
-  std::copy(src_ptr, src_ptr + src.data_size(), dst_ptr);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_output_array(dst);
+  encoder.dispatch([src_ptr, dst_ptr, size = src.data_size()]() {
+    std::copy(src_ptr, src_ptr + size, dst_ptr);
+  });
 }
 
 template <typename SrcT, typename DstT, int D>
@@ -56,151 +66,220 @@ template <typename SrcT, typename DstT>
 void copy_general_general(
     const array& src,
     array& dst,
+    Stream stream,
     const Shape& data_shape,
     const Strides& i_strides,
     const Strides& o_strides,
     int64_t i_offset,
-    int64_t o_offset) {
-  if (data_shape.empty()) {
-    auto val = static_cast<DstT>(*(src.data<SrcT>() + i_offset));
-    auto dst_ptr = dst.data<DstT>() + o_offset;
-    *dst_ptr = val;
-    return;
-  }
-  auto [shape, strides] =
-      collapse_contiguous_dims(data_shape, {i_strides, o_strides});
+    int64_t o_offset,
+    const std::optional<array>& dynamic_i_offset,
+    const std::optional<array>& dynamic_o_offset) {
   auto src_ptr = src.data<SrcT>() + i_offset;
   auto dst_ptr = dst.data<DstT>() + o_offset;
-  int ndim = shape.size();
-  if (ndim == 1) {
-    copy_dims<SrcT, DstT, 1>(
-        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
-    return;
-  } else if (ndim == 2) {
-    copy_dims<SrcT, DstT, 2>(
-        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
-    return;
-  } else if (ndim == 3) {
-    copy_dims<SrcT, DstT, 3>(
-        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
-    return;
-  }
-  ContiguousIterator in(shape, strides[0], ndim - 3);
-  ContiguousIterator out(shape, strides[1], ndim - 3);
-  auto stride = std::accumulate(
-      shape.end() - 3, shape.end(), 1, std::multiplies<int64_t>());
-  for (int64_t elem = 0; elem < src.size(); elem += stride) {
-    copy_dims<SrcT, DstT, 3>(
-        src_ptr + in.loc,
-        dst_ptr + out.loc,
-        shape,
-        strides[0],
-        strides[1],
-        ndim - 3);
-    in.step();
-    out.step();
-  }
+  auto i_offset_ptr =
+      dynamic_i_offset ? dynamic_i_offset->data<int64_t>() : nullptr;
+  auto o_offset_ptr =
+      dynamic_o_offset ? dynamic_o_offset->data<int64_t>() : nullptr;
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_output_array(dst);
+  encoder.dispatch([src_ptr,
+                    dst_ptr,
+                    size = src.size(),
+                    data_shape = data_shape,
+                    i_strides = i_strides,
+                    o_strides = o_strides,
+                    i_offset_ptr,
+                    o_offset_ptr]() mutable {
+    if (data_shape.empty()) {
+      auto val = static_cast<DstT>(*src_ptr);
+      *dst_ptr = val;
+      return;
+    }
+    auto [shape, strides] =
+        collapse_contiguous_dims(data_shape, {i_strides, o_strides});
+
+    int ndim = shape.size();
+    if (ndim < 3) {
+      if (i_offset_ptr) {
+        src_ptr += i_offset_ptr[0];
+      }
+      if (o_offset_ptr) {
+        dst_ptr += o_offset_ptr[0];
+      }
+
+      if (ndim == 1) {
+        copy_dims<SrcT, DstT, 1>(
+            src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+      } else if (ndim == 2) {
+        copy_dims<SrcT, DstT, 2>(
+            src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+      } else if (ndim == 3) {
+        copy_dims<SrcT, DstT, 3>(
+            src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+      }
+      return;
+    }
+    if (i_offset_ptr) {
+      src_ptr += i_offset_ptr[0];
+    }
+    if (o_offset_ptr) {
+      dst_ptr += o_offset_ptr[0];
+    }
+
+    ContiguousIterator in(shape, strides[0], ndim - 3);
+    ContiguousIterator out(shape, strides[1], ndim - 3);
+    auto stride = std::accumulate(
+        shape.end() - 3, shape.end(), 1, std::multiplies<int64_t>());
+    for (int64_t elem = 0; elem < size; elem += stride) {
+      copy_dims<SrcT, DstT, 3>(
+          src_ptr + in.loc,
+          dst_ptr + out.loc,
+          shape,
+          strides[0],
+          strides[1],
+          ndim - 3);
+      in.step();
+      out.step();
+    }
+  });
 }
 
 template <typename SrcT, typename DstT>
-inline void copy_general_general(const array& src, array& dst) {
+inline void copy_general_general(const array& src, array& dst, Stream stream) {
   copy_general_general<SrcT, DstT>(
-      src, dst, src.shape(), src.strides(), dst.strides(), 0, 0);
+      src,
+      dst,
+      stream,
+      src.shape(),
+      src.strides(),
+      dst.strides(),
+      0,
+      0,
+      std::nullopt,
+      std::nullopt);
 }
 
 template <typename SrcT, typename DstT>
 void copy_general(
     const array& src,
     array& dst,
+    Stream stream,
     const Shape& data_shape,
     const Strides& i_strides,
     const Strides&,
     int64_t i_offset,
-    int64_t o_offset) {
+    int64_t o_offset,
+    const std::optional<array>& dynamic_i_offset,
+    const std::optional<array>& dynamic_o_offset) {
   copy_general_general<SrcT, DstT>(
       src,
       dst,
+      stream,
       data_shape,
       i_strides,
       make_contiguous_strides(data_shape),
       i_offset,
-      o_offset);
+      o_offset,
+      dynamic_i_offset,
+      dynamic_o_offset);
 }
 
 template <typename SrcT, typename DstT>
-inline void copy_general(const array& src, array& dst) {
+inline void copy_general(const array& src, array& dst, Stream stream) {
   copy_general_general<SrcT, DstT>(
       src,
       dst,
+      stream,
       src.shape(),
       src.strides(),
       make_contiguous_strides(src.shape()),
       0,
-      0);
+      0,
+      std::nullopt,
+      std::nullopt);
 }
 
 template <typename SrcT, typename DstT, typename... Args>
-void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
+void copy(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Stream stream,
+    Args&&... args) {
   switch (ctype) {
     case CopyType::Scalar:
-      copy_single<SrcT, DstT>(src, dst);
+      copy_single<SrcT, DstT>(src, dst, stream);
       return;
     case CopyType::Vector:
-      copy_vector<SrcT, DstT>(src, dst);
+      copy_vector<SrcT, DstT>(src, dst, stream);
       return;
     case CopyType::General:
-      copy_general<SrcT, DstT>(src, dst, std::forward<Args>(args)...);
+      copy_general<SrcT, DstT>(src, dst, stream, std::forward<Args>(args)...);
       return;
     case CopyType::GeneralGeneral:
-      copy_general_general<SrcT, DstT>(src, dst, std::forward<Args>(args)...);
+      copy_general_general<SrcT, DstT>(
+          src, dst, stream, std::forward<Args>(args)...);
       return;
   }
 }
 
 template <typename SrcT, typename... Args>
-void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
+void copy(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Stream stream,
+    Args&&... args) {
   switch (dst.dtype()) {
     case bool_:
-      copy<SrcT, bool>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, bool>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint8:
-      copy<SrcT, uint8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, uint8_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint16:
-      copy<SrcT, uint16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, uint16_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint32:
-      copy<SrcT, uint32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, uint32_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint64:
-      copy<SrcT, uint64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, uint64_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int8:
-      copy<SrcT, int8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, int8_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int16:
-      copy<SrcT, int16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, int16_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int32:
-      copy<SrcT, int32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, int32_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int64:
-      copy<SrcT, int64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, int64_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float16:
-      copy<SrcT, float16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, float16_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float32:
-      copy<SrcT, float>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, float>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float64:
-      copy<SrcT, double>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, double>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case bfloat16:
-      copy<SrcT, bfloat16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, bfloat16_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case complex64:
-      copy<SrcT, complex64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<SrcT, complex64_t>(
+          src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
   }
 }
@@ -210,84 +289,71 @@ inline void copy_inplace_dispatch(
     const array& src,
     array& dst,
     CopyType ctype,
+    Stream stream,
     Args&&... args) {
   switch (src.dtype()) {
     case bool_:
-      copy<bool>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<bool>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint8:
-      copy<uint8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<uint8_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint16:
-      copy<uint16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<uint16_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint32:
-      copy<uint32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<uint32_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case uint64:
-      copy<uint64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<uint64_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int8:
-      copy<int8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<int8_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int16:
-      copy<int16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<int16_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int32:
-      copy<int32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<int32_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case int64:
-      copy<int64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<int64_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float16:
-      copy<float16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<float16_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float32:
-      copy<float>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<float>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case float64:
-      copy<double>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<double>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case bfloat16:
-      copy<bfloat16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<bfloat16_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
     case complex64:
-      copy<complex64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      copy<complex64_t>(src, dst, ctype, stream, std::forward<Args>(args)...);
       break;
   }
 }
 
 } // namespace
 
-void copy_inplace(const array& src, array& dst, CopyType ctype) {
-  copy_inplace_dispatch(src, dst, ctype);
+void copy_inplace(const array& src, array& dst, CopyType ctype, Stream stream) {
+  copy_inplace_dispatch(src, dst, ctype, stream);
 }
 
-void copy(const array& src, array& dst, CopyType ctype) {
-  // Allocate the output
-  switch (ctype) {
-    case CopyType::Vector:
-      if (src.is_donatable() && src.itemsize() == dst.itemsize()) {
-        dst.copy_shared_buffer(src);
-      } else {
-        auto size = src.data_size();
-        dst.set_data(
-            allocator::malloc_or_wait(size * dst.itemsize()),
-            size,
-            src.strides(),
-            src.flags());
-      }
-      break;
-    case CopyType::Scalar:
-    case CopyType::General:
-    case CopyType::GeneralGeneral:
-      dst.set_data(allocator::malloc_or_wait(dst.nbytes()));
-      break;
+void copy(const array& src, array& dst, CopyType ctype, Stream stream) {
+  bool donated = set_copy_output_data(src, dst, ctype);
+  if (donated && src.dtype() == dst.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
   }
   if (ctype == CopyType::GeneralGeneral) {
     ctype = CopyType::General;
   }
-  copy_inplace(src, dst, ctype);
+  copy_inplace(src, dst, ctype, stream);
 }
 
 void copy_inplace(
@@ -298,7 +364,10 @@ void copy_inplace(
     const Strides& o_strides,
     int64_t i_offset,
     int64_t o_offset,
-    CopyType ctype) {
+    CopyType ctype,
+    Stream stream,
+    const std::optional<array>& dynamic_i_offset, /* = std::nullopt */
+    const std::optional<array>& dynamic_o_offset /* = std::nullopt */) {
   switch (ctype) {
     case CopyType::General:
     case CopyType::GeneralGeneral:
@@ -306,15 +375,18 @@ void copy_inplace(
           src,
           dst,
           ctype,
+          stream,
           data_shape,
           i_strides,
           o_strides,
           i_offset,
-          o_offset);
+          o_offset,
+          dynamic_i_offset,
+          dynamic_o_offset);
       break;
     case CopyType::Scalar:
     case CopyType::Vector:
-      copy_inplace_dispatch(src, dst, ctype);
+      copy_inplace_dispatch(src, dst, ctype, stream);
   }
 }
 

mlx/backend/cpu/copy.h

@@ -2,14 +2,16 @@
 
 #pragma once
 
+#include <optional>
+
 #include "mlx/array.h"
 #include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/utils.h"
 
 namespace mlx::core {
 
-void copy(const array& src, array& dst, CopyType ctype);
-void copy_inplace(const array& src, array& dst, CopyType ctype);
+void copy(const array& src, array& dst, CopyType ctype, Stream stream);
+void copy_inplace(const array& src, array& dst, CopyType ctype, Stream stream);
 
 void copy_inplace(
     const array& src,
@@ -19,6 +21,9 @@ void copy_inplace(
     const Strides& o_strides,
     int64_t i_offset,
     int64_t o_offset,
-    CopyType ctype);
+    CopyType ctype,
+    Stream stream,
+    const std::optional<array>& dynamic_i_offset = std::nullopt,
+    const std::optional<array>& dynamic_o_offset = std::nullopt);
 
 } // namespace mlx::core

mlx/backend/cpu/distributed.cpp

@@ -0,0 +1,94 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/primitives.h"
+
+namespace mlx::core::distributed {
+
+std::pair<array, bool> ensure_row_contiguous(const array& arr, Stream stream) {
+  if (arr.flags().row_contiguous) {
+    return {arr, false};
+  } else {
+    array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+    copy(arr, arr_copy, CopyType::General, stream);
+    return {arr_copy, true};
+  }
+};
+
+void AllReduce::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto donate_or_copy = [s = stream()](const array& in, array& out) {
+    if (in.flags().row_contiguous) {
+      if (in.is_donatable()) {
+        out.copy_shared_buffer(in);
+      } else {
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+      }
+      return in;
+    } else {
+      array arr_copy(in.shape(), in.dtype(), nullptr, {});
+      copy(in, arr_copy, CopyType::General, s);
+      out.copy_shared_buffer(arr_copy);
+      return arr_copy;
+    }
+  };
+
+  auto in = donate_or_copy(inputs[0], outputs[0]);
+  switch (reduce_type_) {
+    case Sum:
+      distributed::detail::all_sum(group(), in, outputs[0], stream());
+      break;
+    default:
+      throw std::runtime_error("Only all reduce sum is supported for now");
+  }
+}
+
+void AllGather::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
+  outputs[0].set_data(allocator::malloc_or_wait(outputs[0].nbytes()));
+  distributed::detail::all_gather(group(), in, outputs[0], stream());
+  if (copied) {
+    auto& enc = cpu::get_command_encoder(stream());
+    enc.add_temporary(in);
+  }
+}
+
+void Send::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
+  distributed::detail::send(group(), in, dst_, stream());
+  outputs[0].copy_shared_buffer(inputs[0]);
+  if (copied) {
+    auto& enc = cpu::get_command_encoder(stream());
+    enc.add_temporary(in);
+  }
+}
+
+void Recv::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 0);
+  assert(outputs.size() == 1);
+
+  outputs[0].set_data(allocator::malloc_or_wait(outputs[0].nbytes()));
+  distributed::detail::recv(group(), outputs[0], src_, stream());
+}
+
+} // namespace mlx::core::distributed

mlx/backend/cpu/eigh.cpp

@@ -3,6 +3,7 @@
 #include "mlx/allocator.h"
 #include "mlx/array.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/linalg.h"
 #include "mlx/primitives.h"
@@ -11,35 +12,77 @@ namespace mlx::core {
 
 namespace {
 
-void ssyevd(
-    char jobz,
-    char uplo,
-    float* a,
-    int N,
-    float* w,
-    float* work,
-    int lwork,
-    int* iwork,
-    int liwork) {
-  int info;
-  MLX_LAPACK_FUNC(ssyevd)
-  (
-      /* jobz = */ &jobz,
-      /* uplo = */ &uplo,
-      /* n = */ &N,
-      /* a = */ a,
-      /* lda = */ &N,
-      /* w = */ w,
-      /* work = */ work,
-      /* lwork = */ &lwork,
-      /* iwork = */ iwork,
-      /* liwork = */ &liwork,
-      /* info = */ &info);
-  if (info != 0) {
-    std::stringstream msg;
-    msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
-        << info;
-    throw std::runtime_error(msg.str());
+template <typename T>
+void eigh_impl(
+    array& vectors,
+    array& values,
+    const std::string& uplo,
+    bool compute_eigenvectors,
+    Stream stream) {
+  auto vec_ptr = vectors.data<T>();
+  auto eig_ptr = values.data<T>();
+  char jobz = compute_eigenvectors ? 'V' : 'N';
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(vectors);
+  encoder.set_output_array(values);
+  encoder.dispatch([vec_ptr,
+                    eig_ptr,
+                    jobz,
+                    uplo = uplo[0],
+                    N = vectors.shape(-1),
+                    size = vectors.size()]() mutable {
+    // Work query
+    int lwork = -1;
+    int liwork = -1;
+    int info;
+    {
+      T work;
+      int iwork;
+      syevd<T>(
+          &jobz,
+          &uplo,
+          &N,
+          nullptr,
+          &N,
+          nullptr,
+          &work,
+          &lwork,
+          &iwork,
+          &liwork,
+          &info);
+      lwork = static_cast<int>(work);
+      liwork = iwork;
+    }
+
+    auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
+    auto iwork_buf =
+        array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
+    for (size_t i = 0; i < size / (N * N); ++i) {
+      syevd<T>(
+          &jobz,
+          &uplo,
+          &N,
+          vec_ptr,
+          &N,
+          eig_ptr,
+          static_cast<T*>(work_buf.buffer.raw_ptr()),
+          &lwork,
+          static_cast<int*>(iwork_buf.buffer.raw_ptr()),
+          &liwork,
+          &info);
+      vec_ptr += N * N;
+      eig_ptr += N;
+      if (info != 0) {
+        std::stringstream msg;
+        msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
+            << info;
+        throw std::runtime_error(msg.str());
+      }
+    }
+  });
+  if (!compute_eigenvectors) {
+    encoder.add_temporary(vectors);
   }
 }
 
@@ -60,7 +103,8 @@ void Eigh::eval_cpu(
   copy(
       a,
       vectors,
-      a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream());
 
   if (compute_eigenvectors_) {
     // Set the strides and flags so the eigenvectors
@@ -78,41 +122,19 @@ void Eigh::eval_cpu(
         flags.col_contiguous = true;
       }
     }
-    vectors.move_shared_buffer(vectors, strides, flags, vectors.data_size());
+    vectors.copy_shared_buffer(vectors, strides, flags, vectors.data_size());
   }
-
-  auto vec_ptr = vectors.data<float>();
-  auto eig_ptr = values.data<float>();
-
-  char jobz = compute_eigenvectors_ ? 'V' : 'N';
-  auto N = a.shape(-1);
-
-  // Work query
-  int lwork;
-  int liwork;
-  {
-    float work;
-    int iwork;
-    ssyevd(jobz, uplo_[0], nullptr, N, nullptr, &work, -1, &iwork, -1);
-    lwork = static_cast<int>(work);
-    liwork = iwork;
-  }
-
-  auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
-  auto iwork_buf = array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
-  for (size_t i = 0; i < a.size() / (N * N); ++i) {
-    ssyevd(
-        jobz,
-        uplo_[0],
-        vec_ptr,
-        N,
-        eig_ptr,
-        static_cast<float*>(work_buf.buffer.raw_ptr()),
-        lwork,
-        static_cast<int*>(iwork_buf.buffer.raw_ptr()),
-        liwork);
-    vec_ptr += N * N;
-    eig_ptr += N;
+  switch (a.dtype()) {
+    case float32:
+      eigh_impl<float>(vectors, values, uplo_, compute_eigenvectors_, stream());
+      break;
+    case float64:
+      eigh_impl<double>(
+          vectors, values, uplo_, compute_eigenvectors_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Eigh::eval_cpu] only supports float32 or float64.");
   }
 }
 

mlx/backend/cpu/encoder.cpp

@@ -0,0 +1,16 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core::cpu {
+
+CommandEncoder& get_command_encoder(Stream stream) {
+  static std::unordered_map<int, CommandEncoder> encoder_map;
+  auto it = encoder_map.find(stream.index);
+  if (it == encoder_map.end()) {
+    it = encoder_map.emplace(stream.index, stream).first;
+  }
+  return it->second;
+}
+
+} // namespace mlx::core::cpu

mlx/backend/cpu/encoder.h

@@ -0,0 +1,53 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/array.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::cpu {
+
+struct CommandEncoder {
+  CommandEncoder(Stream stream) : stream_(stream) {}
+
+  CommandEncoder(const CommandEncoder&) = delete;
+  CommandEncoder& operator=(const CommandEncoder&) = delete;
+  CommandEncoder(CommandEncoder&&) = delete;
+  CommandEncoder& operator=(CommandEncoder&&) = delete;
+
+  void set_input_array(const array& a) {}
+  void set_output_array(array& a) {}
+
+  // Hold onto a temporary until any already scheduled tasks which use it as
+  // an input are complete.
+  void add_temporary(array arr) {
+    temporaries_.push_back(std::move(arr));
+  }
+
+  void add_temporaries(std::vector<array> arrays) {
+    temporaries_.insert(
+        temporaries_.end(),
+        std::make_move_iterator(arrays.begin()),
+        std::make_move_iterator(arrays.end()));
+  }
+
+  std::vector<array>& temporaries() {
+    return temporaries_;
+  }
+
+  template <class F, class... Args>
+  void dispatch(F&& f, Args&&... args) {
+    auto task = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
+    scheduler::enqueue(stream_, std::move(task));
+  }
+
+ private:
+  Stream stream_;
+  std::vector<array> temporaries_;
+};
+
+CommandEncoder& get_command_encoder(Stream stream);
+
+} // namespace mlx::core::cpu

mlx/backend/cpu/eval.cpp

@@ -0,0 +1,44 @@
+// Copyright © 2025 Apple Inc.
+#include "mlx/backend/cpu/eval.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::cpu {
+
+void eval(array& arr) {
+  auto s = arr.primitive().stream();
+
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
+    }
+    arr.primitive().eval_cpu(arr.inputs(), outputs);
+  }
+
+  std::unordered_set<std::shared_ptr<array::Data>> buffers;
+  for (auto& in : arr.inputs()) {
+    buffers.insert(in.data_shared_ptr());
+  }
+  for (auto& s : arr.siblings()) {
+    buffers.insert(s.data_shared_ptr());
+  }
+  // Remove the output if it was donated to by an input
+  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
+    buffers.erase(it);
+  }
+  auto& encoder = cpu::get_command_encoder(s);
+  scheduler::notify_new_task(s);
+  encoder.dispatch([s,
+                    buffers = std::move(buffers),
+                    temps = std::move(encoder.temporaries())]() {
+    scheduler::notify_task_completion(s);
+  });
+}
+
+} // namespace mlx::core::cpu

mlx/backend/cpu/eval.h

@@ -0,0 +1,12 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::cpu {
+
+void eval(array& arr);
+
+} // namespace mlx::core::cpu

mlx/backend/cpu/fft.cpp

@@ -4,6 +4,7 @@
 
 #include "mlx/3rdparty/pocketfft.h"
 #include "mlx/allocator.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
@@ -38,46 +39,78 @@ void FFT::eval_cpu(const std::vector<array>& inputs, array& out) {
         });
     scale /= nelem;
   }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
   if (in.dtype() == complex64 && out.dtype() == complex64) {
     auto in_ptr =
         reinterpret_cast<const std::complex<float>*>(in.data<complex64_t>());
     auto out_ptr =
         reinterpret_cast<std::complex<float>*>(out.data<complex64_t>());
-    pocketfft::c2c(
-        shape,
-        strides_in,
-        strides_out,
-        axes_,
-        !inverse_,
-        in_ptr,
-        out_ptr,
-        scale);
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::c2c(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
   } else if (in.dtype() == float32 && out.dtype() == complex64) {
     auto in_ptr = in.data<float>();
     auto out_ptr =
         reinterpret_cast<std::complex<float>*>(out.data<complex64_t>());
-    pocketfft::r2c(
-        shape,
-        strides_in,
-        strides_out,
-        axes_,
-        !inverse_,
-        in_ptr,
-        out_ptr,
-        scale);
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::r2c(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
   } else if (in.dtype() == complex64 && out.dtype() == float32) {
     auto in_ptr =
         reinterpret_cast<const std::complex<float>*>(in.data<complex64_t>());
     auto out_ptr = out.data<float>();
-    pocketfft::c2r(
-        shape,
-        strides_in,
-        strides_out,
-        axes_,
-        !inverse_,
-        in_ptr,
-        out_ptr,
-        scale);
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::c2r(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
   } else {
     throw std::runtime_error(
         "[FFT] Received unexpected input and output type combination.");

mlx/backend/cpu/gemm.h

@@ -7,14 +7,20 @@ namespace mlx::core {
 
 template <typename T>
 void matmul(
-    const array& a,
-    const array& b,
-    array& out,
+    const T* a,
+    const T* b,
+    T* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides);
 
 } // namespace mlx::core

mlx/backend/cpu/gemms/bnns.cpp

@@ -9,39 +9,46 @@
 
 namespace mlx::core {
 
-BNNSDataType to_bnns_dtype(Dtype mlx_dtype) {
-  uint32_t size_bits = size_of(mlx_dtype) * 8;
-  switch (kindof(mlx_dtype)) {
-    case Dtype::Kind::b:
-      return BNNSDataTypeBoolean;
-    case Dtype::Kind::u:
-      return BNNSDataType(BNNSDataTypeUIntBit | size_bits);
-    case Dtype::Kind::i:
-      return BNNSDataType(BNNSDataTypeIntBit | size_bits);
-    case Dtype::Kind::f:
-      return BNNSDataType(BNNSDataTypeFloatBit | size_bits);
-    case Dtype::Kind::V:
-      return BNNSDataTypeBFloat16;
-    case Dtype::Kind::c:
-      throw std::invalid_argument("BNNS does not support complex types");
-  }
+template <typename T>
+constexpr BNNSDataType to_bnns_dtype();
+
+template <>
+constexpr BNNSDataType to_bnns_dtype<float>() {
+  return BNNSDataType(BNNSDataTypeFloatBit | 32);
+}
+template <>
+constexpr BNNSDataType to_bnns_dtype<float16_t>() {
+  return BNNSDataType(BNNSDataTypeFloatBit | 16);
 }
 
+template <>
+constexpr BNNSDataType to_bnns_dtype<bfloat16_t>() {
+  return BNNSDataTypeBFloat16;
+}
+
+template <typename T>
 void matmul_bnns(
-    const array& a,
-    const array& b,
-    array& out,
+    const T* a,
+    const T* b,
+    T* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta) {
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
 
-  BNNSDataType bnns_dtype = to_bnns_dtype(out.dtype());
+  BNNSDataType bnns_dtype = to_bnns_dtype<T>();
 
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
@@ -115,14 +122,14 @@ void matmul_bnns(
   auto bnns_filter =
       BNNSFilterCreateLayerBroadcastMatMul(&gemm_params, nullptr);
 
-  for (int i = 0; i < (a.size() / (M * K)); ++i) {
+  for (int i = 0; i < batch_size; ++i) {
     BNNSFilterApplyTwoInput(
         bnns_filter,
-        a.data<uint8_t>() +
-            elem_to_loc(M * K * i, a.shape(), a.strides()) * a.itemsize(),
-        b.data<uint8_t>() +
-            elem_to_loc(K * N * i, b.shape(), b.strides()) * b.itemsize(),
-        out.data<uint8_t>() + M * N * i * out.itemsize());
+        reinterpret_cast<const uint8_t*>(
+            a + elem_to_loc(M * K * i, a_shape, a_strides)),
+        reinterpret_cast<const uint8_t*>(
+            b + elem_to_loc(K * N * i, b_shape, b_strides)),
+        reinterpret_cast<uint8_t*>(out + M * N * i));
   }
 
   BNNSFilterDestroy(bnns_filter);
@@ -131,30 +138,72 @@ void matmul_bnns(
 
 template <>
 void matmul<float16_t>(
-    const array& a,
-    const array& b,
-    array& out,
+    const float16_t* a,
+    const float16_t* b,
+    float16_t* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta) {
-  matmul_bnns(a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  matmul_bnns(
+      a,
+      b,
+      out,
+      a_transposed,
+      b_transposed,
+      lda,
+      ldb,
+      ldc,
+      alpha,
+      beta,
+      batch_size,
+      a_shape,
+      a_strides,
+      b_shape,
+      b_strides);
 }
 
 template <>
 void matmul<bfloat16_t>(
-    const array& a,
-    const array& b,
-    array& out,
+    const bfloat16_t* a,
+    const bfloat16_t* b,
+    bfloat16_t* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta) {
-  matmul_bnns(a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  matmul_bnns(
+      a,
+      b,
+      out,
+      a_transposed,
+      b_transposed,
+      lda,
+      ldb,
+      ldc,
+      alpha,
+      beta,
+      batch_size,
+      a_shape,
+      a_strides,
+      b_shape,
+      b_strides);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/gemms/cblas.cpp

@@ -8,20 +8,27 @@ namespace mlx::core {
 
 template <>
 void matmul<float>(
-    const array& a,
-    const array& b,
-    array& out,
+    const float* a,
+    const float* b,
+    float* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta) {
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
 
-  for (int i = 0; i < (a.size() / (M * K)); ++i) {
+  for (int i = 0; i < batch_size; ++i) {
     cblas_sgemm(
         CblasRowMajor,
         a_transposed ? CblasTrans : CblasNoTrans, // transA
@@ -29,34 +36,40 @@ void matmul<float>(
         M,
         N,
         K,
-        alpha, // alpha
-        a.data<float>() + elem_to_loc(M * K * i, a.shape(), a.strides()),
+        alpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
         lda,
-        b.data<float>() + elem_to_loc(K * N * i, b.shape(), b.strides()),
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
         ldb,
-        beta, // beta
-        out.data<float>() + M * N * i,
-        out.shape(-1) // ldc
-    );
+        beta,
+        out + M * N * i,
+        ldc);
   }
 }
 
 template <>
 void matmul<double>(
-    const array& a,
-    const array& b,
-    array& out,
+    const double* a,
+    const double* b,
+    double* out,
     bool a_transposed,
     bool b_transposed,
     size_t lda,
     size_t ldb,
+    size_t ldc,
     float alpha,
-    float beta) {
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
 
-  for (int i = 0; i < (a.size() / (M * K)); ++i) {
+  for (int i = 0; i < batch_size; ++i) {
     cblas_dgemm(
         CblasRowMajor,
         a_transposed ? CblasTrans : CblasNoTrans, // transA
@@ -64,15 +77,14 @@ void matmul<double>(
         M,
         N,
         K,
-        alpha, // alpha
-        a.data<double>() + elem_to_loc(M * K * i, a.shape(), a.strides()),
+        alpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
         lda,
-        b.data<double>() + elem_to_loc(K * N * i, b.shape(), b.strides()),
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
         ldb,
-        beta, // beta
-        out.data<double>() + M * N * i,
-        out.shape(-1) // ldc
-    );
+        beta,
+        out + M * N * i,
+        ldc);
   }
 }
 

mlx/backend/cpu/gemms/no_bf16.cpp

@@ -6,15 +6,21 @@ namespace mlx::core {
 
 template <>
 void matmul<bfloat16_t>(
-    const array&,
-    const array&,
-    array&,
+    const bfloat16_t*,
+    const bfloat16_t*,
+    bfloat16_t*,
     bool,
     bool,
     size_t,
     size_t,
+    size_t,
     float,
-    float) {
+    float,
+    size_t,
+    const Shape&,
+    const Strides&,
+    const Shape&,
+    const Strides&) {
   throw std::runtime_error("[Matmul::eval_cpu] bfloat16 not supported.");
 }
 

mlx/backend/cpu/gemms/no_fp16.cpp

@@ -6,15 +6,21 @@ namespace mlx::core {
 
 template <>
 void matmul<float16_t>(
-    const array&,
-    const array&,
-    array&,
+    const float16_t*,
+    const float16_t*,
+    float16_t*,
     bool,
     bool,
     size_t,
     size_t,
+    size_t,
     float,
-    float) {
+    float,
+    size_t,
+    const Shape&,
+    const Strides&,
+    const Shape&,
+    const Strides&) {
   throw std::runtime_error("[Matmul::eval_cpu] float16 not supported.");
 }
 

mlx/backend/cpu/hadamard.cpp

@@ -4,16 +4,17 @@
 
 #include "mlx/backend/common/hadamard.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
 // n = 2^k component
 template <typename T>
-void hadamard_n(array& out, int n, int m, float scale) {
-  for (int b = 0; b < out.size() / n; b++) {
+void hadamard_n(T* out, int n, int m, float scale, size_t size) {
+  for (int b = 0; b < size / n; b++) {
     size_t loc = b * n;
-    T* data_ptr = out.data<T>() + loc;
+    T* data_ptr = out + loc;
     int h = 1;
     int n_over_2 = n / 2;
     while (h < n) {
@@ -36,7 +37,7 @@ void hadamard_n(array& out, int n, int m, float scale) {
 
 // m component
 template <typename T>
-void hadamard_m(array& out, int n, int m, float scale) {
+void hadamard_m(T* out, int n, int m, float scale, size_t size) {
   auto h_matrices = hadamard_matrices();
   auto& matrix = h_matrices[m];
   auto start = 1;
@@ -51,9 +52,9 @@ void hadamard_m(array& out, int n, int m, float scale) {
     end = matrix.find('\n', start);
   }
 
-  for (int b = 0; b < out.size() / m / n; b++) {
+  for (int b = 0; b < size / m / n; b++) {
     size_t loc = b * n * m;
-    T* data_ptr = out.data<T>() + loc;
+    T* data_ptr = out + loc;
     for (int i = 0; i < n; i++) {
       std::vector<float> out(m);
       for (int j = 0; j < m; j++) {
@@ -74,12 +75,17 @@ void hadamard_m(array& out, int n, int m, float scale) {
 }
 
 template <typename T>
-void hadamard(array& out, int n, int m, float scale) {
-  float n_scale = m > 1 ? 1.0 : scale;
-  hadamard_n<T>(out, n, m, n_scale);
-  if (m > 1) {
-    hadamard_m<T>(out, n, m, scale);
-  }
+void hadamard(array& out, int n, int m, float scale, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  auto out_ptr = out.data<T>();
+  encoder.dispatch([out_ptr, size = out.size(), n, m, scale]() {
+    float n_scale = m > 1 ? 1.0 : scale;
+    hadamard_n<T>(out_ptr, n, m, n_scale, size);
+    if (m > 1) {
+      hadamard_m<T>(out_ptr, n, m, scale, size);
+    }
+  });
 }
 
 void Hadamard::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -87,18 +93,26 @@ void Hadamard::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
 
   // Copy input to output
-  copy(in, out, CopyType::General);
+  if (in.flags().row_contiguous && in.is_donatable()) {
+    out.copy_shared_buffer(in);
+  } else {
+    copy(
+        in,
+        out,
+        in.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+        stream());
+  }
 
   int axis = out.ndim() - 1;
   auto [n, m] = decompose_hadamard(out.shape(axis));
 
   switch (in.dtype()) {
     case float32:
-      return hadamard<float>(out, n, m, scale_);
+      return hadamard<float>(out, n, m, scale_, stream());
     case float16:
-      return hadamard<float16_t>(out, n, m, scale_);
+      return hadamard<float16_t>(out, n, m, scale_, stream());
     case bfloat16:
-      return hadamard<bfloat16_t>(out, n, m, scale_);
+      return hadamard<bfloat16_t>(out, n, m, scale_, stream());
     default:
       throw std::invalid_argument("[hadamard] Unsupported type.");
   }

mlx/backend/cpu/indexing.cpp

@@ -8,6 +8,7 @@
 
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 
 namespace mlx::core {
 
@@ -27,7 +28,8 @@ void gather(
     const std::vector<array>& inds,
     array& out,
     const std::vector<int>& axes,
-    const Shape& slice_sizes) {
+    const Shape& slice_sizes,
+    Stream stream) {
   // If the array is row contiguous then we can do a contiguous copy given
   // two conditions on the slice size:
   // - Any number of leading ones in the slice sizes are allowed
@@ -73,38 +75,60 @@ void gather(
   size_t ind_size = slice_size == 0 ? 0 : out.size() / slice_size;
   const T* src_ptr = src.data<T>();
   T* dst_ptr = out.data<T>();
-  size_t out_idx = 0;
 
   std::vector<ContiguousIterator> its(inds.begin(), inds.end());
   ContiguousIterator src_it;
   if (!can_copy && src.ndim() > 0) {
     src_it = ContiguousIterator(slice_sizes, src.strides(), src.ndim());
   }
-  for (int idx = 0; idx < ind_size; idx++) {
-    size_t src_idx = 0;
-    for (int ii = 0; ii < inds.size(); ++ii) {
-      auto ax = axes[ii];
-      auto idx_loc = its[ii].loc;
-      its[ii].step();
-      auto idx_val =
-          offset_neg_idx(inds[ii].data<IdxT>()[idx_loc], src.shape(ax));
-      src_idx += (idx_val * src.strides()[ax]);
-    }
 
-    if (slice_size == 1) {
-      dst_ptr[out_idx++] = src_ptr[src_idx];
-    } else if (can_copy) {
-      std::copy(
-          src_ptr + src_idx, src_ptr + src_idx + slice_size, dst_ptr + out_idx);
-      out_idx += slice_size;
-    } else {
-      for (int jj = 0; jj < slice_size; jj++) {
-        dst_ptr[out_idx++] = src_ptr[src_idx + src_it.loc];
-        src_it.step();
-      }
-      src_it.reset();
-    }
+  std::vector<const IdxT*> ind_ptrs;
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  for (auto& idx : inds) {
+    ind_ptrs.push_back(idx.data<IdxT>());
+    encoder.set_input_array(idx);
   }
+  encoder.set_output_array(out);
+  encoder.dispatch([src_ptr,
+                    dst_ptr,
+                    ind_ptrs = std::move(ind_ptrs),
+                    axes,
+                    ind_size,
+                    slice_size,
+                    src_shape = src.shape(),
+                    src_strides = src.strides(),
+                    src_it = std::move(src_it),
+                    its = std::move(its),
+                    can_copy]() mutable {
+    size_t out_idx = 0;
+    for (int idx = 0; idx < ind_size; idx++) {
+      size_t src_idx = 0;
+      for (int ii = 0; ii < ind_ptrs.size(); ++ii) {
+        auto ax = axes[ii];
+        auto idx_loc = its[ii].loc;
+        its[ii].step();
+        auto idx_val = offset_neg_idx(ind_ptrs[ii][idx_loc], src_shape[ax]);
+        src_idx += (idx_val * src_strides[ax]);
+      }
+
+      if (slice_size == 1) {
+        dst_ptr[out_idx++] = src_ptr[src_idx];
+      } else if (can_copy) {
+        std::copy(
+            src_ptr + src_idx,
+            src_ptr + src_idx + slice_size,
+            dst_ptr + out_idx);
+        out_idx += slice_size;
+      } else {
+        for (int jj = 0; jj < slice_size; jj++) {
+          dst_ptr[out_idx++] = src_ptr[src_idx + src_it.loc];
+          src_it.step();
+        }
+        src_it.reset();
+      }
+    }
+  });
 }
 
 template <typename IdxT>
@@ -113,49 +137,50 @@ void dispatch_gather(
     const std::vector<array>& inds,
     array& out,
     const std::vector<int>& axes,
-    const Shape& size) {
+    const Shape& size,
+    Stream stream) {
   switch (out.dtype()) {
     case bool_:
-      gather<bool, IdxT>(src, inds, out, axes, size);
+      gather<bool, IdxT>(src, inds, out, axes, size, stream);
       break;
     case uint8:
-      gather<uint8_t, IdxT>(src, inds, out, axes, size);
+      gather<uint8_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case uint16:
-      gather<uint16_t, IdxT>(src, inds, out, axes, size);
+      gather<uint16_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case uint32:
-      gather<uint32_t, IdxT>(src, inds, out, axes, size);
+      gather<uint32_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case uint64:
-      gather<uint64_t, IdxT>(src, inds, out, axes, size);
+      gather<uint64_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case int8:
-      gather<int8_t, IdxT>(src, inds, out, axes, size);
+      gather<int8_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case int16:
-      gather<int16_t, IdxT>(src, inds, out, axes, size);
+      gather<int16_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case int32:
-      gather<int32_t, IdxT>(src, inds, out, axes, size);
+      gather<int32_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case int64:
-      gather<int64_t, IdxT>(src, inds, out, axes, size);
+      gather<int64_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case float16:
-      gather<float16_t, IdxT>(src, inds, out, axes, size);
+      gather<float16_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case float32:
-      gather<float, IdxT>(src, inds, out, axes, size);
+      gather<float, IdxT>(src, inds, out, axes, size, stream);
       break;
     case float64:
-      gather<double, IdxT>(src, inds, out, axes, size);
+      gather<double, IdxT>(src, inds, out, axes, size, stream);
       break;
     case bfloat16:
-      gather<bfloat16_t, IdxT>(src, inds, out, axes, size);
+      gather<bfloat16_t, IdxT>(src, inds, out, axes, size, stream);
       break;
     case complex64:
-      gather<complex64_t, IdxT>(src, inds, out, axes, size);
+      gather<complex64_t, IdxT>(src, inds, out, axes, size, stream);
       break;
   }
 }
@@ -167,34 +192,34 @@ void Gather::eval_cpu(const std::vector<array>& inputs, array& out) {
   std::vector<array> inds(inputs.begin() + 1, inputs.end());
 
   if (inds.empty()) {
-    dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_);
+    dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_, stream());
     return;
   }
 
   switch (inds[0].dtype()) {
     case uint8:
-      dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case uint16:
-      dispatch_gather<uint16_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<uint16_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case uint32:
-      dispatch_gather<uint32_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<uint32_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case uint64:
-      dispatch_gather<uint64_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<uint64_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case int8:
-      dispatch_gather<int8_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<int8_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case int16:
-      dispatch_gather<int16_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<int16_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case int32:
-      dispatch_gather<int32_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<int32_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     case int64:
-      dispatch_gather<int64_t>(src, inds, out, axes_, slice_sizes_);
+      dispatch_gather<int64_t>(src, inds, out, axes_, slice_sizes_, stream());
       break;
     default:
       throw std::runtime_error(
@@ -207,7 +232,8 @@ void gather_axis(
     const array& src,
     const array& ind,
     array& out,
-    const int axis) {
+    const int axis,
+    Stream stream) {
   auto strides = ind.strides();
   strides.erase(strides.begin() + axis);
   auto shape = ind.shape();
@@ -235,20 +261,39 @@ void gather_axis(
   for (int i = axis + 1; i < ind.ndim(); ++i) {
     size_post *= ind.shape(i);
   }
-  size_t stride_pre = size_post * ind_ax_size;
-  for (size_t i = 0; i < size_pre; i++) {
-    for (size_t k = 0; k < size_post; k++) {
-      for (int j = 0; j < ind_ax_size; ++j) {
-        auto ind_val = offset_neg_idx(
-            ind_ptr[ind_it.loc + j * ind_ax_stride], src_ax_size);
-        dst_ptr[k + j * dst_ax_stride] =
-            src_ptr[src_it.loc + ind_val * src_ax_stride];
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_input_array(ind);
+  encoder.set_output_array(out);
+
+  encoder.dispatch([ind_ptr,
+                    src_ptr,
+                    dst_ptr,
+                    size_pre,
+                    size_post,
+                    ind_ax_size,
+                    src_ax_size,
+                    ind_ax_stride,
+                    src_ax_stride,
+                    dst_ax_stride,
+                    ind_it = std::move(ind_it),
+                    src_it = std::move(src_it)]() mutable {
+    size_t stride_pre = size_post * ind_ax_size;
+    for (size_t i = 0; i < size_pre; i++) {
+      for (size_t k = 0; k < size_post; k++) {
+        for (int j = 0; j < ind_ax_size; ++j) {
+          auto ind_val = offset_neg_idx(
+              ind_ptr[ind_it.loc + j * ind_ax_stride], src_ax_size);
+          dst_ptr[k + j * dst_ax_stride] =
+              src_ptr[src_it.loc + ind_val * src_ax_stride];
+        }
+        ind_it.step();
+        src_it.step();
       }
-      ind_it.step();
-      src_it.step();
+      dst_ptr += stride_pre;
     }
-    dst_ptr += stride_pre;
-  }
+  });
 }
 
 template <typename IdxT>
@@ -256,49 +301,50 @@ void dispatch_gather_axis(
     const array& src,
     const array& inds,
     array& out,
-    const int axis) {
+    const int axis,
+    Stream stream) {
   switch (out.dtype()) {
     case bool_:
-      gather_axis<bool, IdxT>(src, inds, out, axis);
+      gather_axis<bool, IdxT>(src, inds, out, axis, stream);
       break;
     case uint8:
-      gather_axis<uint8_t, IdxT>(src, inds, out, axis);
+      gather_axis<uint8_t, IdxT>(src, inds, out, axis, stream);
       break;
     case uint16:
-      gather_axis<uint16_t, IdxT>(src, inds, out, axis);
+      gather_axis<uint16_t, IdxT>(src, inds, out, axis, stream);
       break;
     case uint32:
-      gather_axis<uint32_t, IdxT>(src, inds, out, axis);
+      gather_axis<uint32_t, IdxT>(src, inds, out, axis, stream);
       break;
     case uint64:
-      gather_axis<uint64_t, IdxT>(src, inds, out, axis);
+      gather_axis<uint64_t, IdxT>(src, inds, out, axis, stream);
       break;
     case int8:
-      gather_axis<int8_t, IdxT>(src, inds, out, axis);
+      gather_axis<int8_t, IdxT>(src, inds, out, axis, stream);
       break;
     case int16:
-      gather_axis<int16_t, IdxT>(src, inds, out, axis);
+      gather_axis<int16_t, IdxT>(src, inds, out, axis, stream);
       break;
     case int32:
-      gather_axis<int32_t, IdxT>(src, inds, out, axis);
+      gather_axis<int32_t, IdxT>(src, inds, out, axis, stream);
       break;
     case int64:
-      gather_axis<int64_t, IdxT>(src, inds, out, axis);
+      gather_axis<int64_t, IdxT>(src, inds, out, axis, stream);
       break;
     case float16:
-      gather_axis<float16_t, IdxT>(src, inds, out, axis);
+      gather_axis<float16_t, IdxT>(src, inds, out, axis, stream);
       break;
     case float32:
-      gather_axis<float, IdxT>(src, inds, out, axis);
+      gather_axis<float, IdxT>(src, inds, out, axis, stream);
       break;
     case float64:
-      gather_axis<double, IdxT>(src, inds, out, axis);
+      gather_axis<double, IdxT>(src, inds, out, axis, stream);
       break;
     case bfloat16:
-      gather_axis<bfloat16_t, IdxT>(src, inds, out, axis);
+      gather_axis<bfloat16_t, IdxT>(src, inds, out, axis, stream);
       break;
     case complex64:
-      gather_axis<complex64_t, IdxT>(src, inds, out, axis);
+      gather_axis<complex64_t, IdxT>(src, inds, out, axis, stream);
       break;
   }
 }
@@ -309,28 +355,28 @@ void GatherAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& inds = inputs[1];
   switch (inds.dtype()) {
     case uint8:
-      dispatch_gather_axis<uint8_t>(src, inds, out, axis_);
+      dispatch_gather_axis<uint8_t>(src, inds, out, axis_, stream());
       break;
     case uint16:
-      dispatch_gather_axis<uint16_t>(src, inds, out, axis_);
+      dispatch_gather_axis<uint16_t>(src, inds, out, axis_, stream());
       break;
     case uint32:
-      dispatch_gather_axis<uint32_t>(src, inds, out, axis_);
+      dispatch_gather_axis<uint32_t>(src, inds, out, axis_, stream());
       break;
     case uint64:
-      dispatch_gather_axis<uint64_t>(src, inds, out, axis_);
+      dispatch_gather_axis<uint64_t>(src, inds, out, axis_, stream());
       break;
     case int8:
-      dispatch_gather_axis<int8_t>(src, inds, out, axis_);
+      dispatch_gather_axis<int8_t>(src, inds, out, axis_, stream());
       break;
     case int16:
-      dispatch_gather_axis<int16_t>(src, inds, out, axis_);
+      dispatch_gather_axis<int16_t>(src, inds, out, axis_, stream());
       break;
     case int32:
-      dispatch_gather_axis<int32_t>(src, inds, out, axis_);
+      dispatch_gather_axis<int32_t>(src, inds, out, axis_, stream());
       break;
     case int64:
-      dispatch_gather_axis<int64_t>(src, inds, out, axis_);
+      dispatch_gather_axis<int64_t>(src, inds, out, axis_, stream());
       break;
     default:
       throw std::runtime_error(
@@ -345,7 +391,8 @@ void scatter(
     array& out,
     const std::vector<array>& inds,
     const std::vector<int>& axes,
-    const OpT& op) {
+    const OpT& op,
+    Stream stream) {
   int nind = inds.size();
   auto inds_ndim = updates.ndim() - out.ndim();
   size_t n_updates = nind ? inds[0].size() : 1;
@@ -361,26 +408,45 @@ void scatter(
   ContiguousIterator update_it(updates);
   ContiguousIterator out_it(update_shape, out.strides(), out.ndim());
 
-  for (int i = 0; i < n_updates; ++i) {
-    size_t out_offset = 0;
-    for (int j = 0; j < nind; ++j) {
-      auto ax = axes[j];
-      auto idx_loc = its[j].loc;
-      its[j].step();
-      auto idx_val =
-          offset_neg_idx(inds[j].data<IdxT>()[idx_loc], out.shape(ax));
-      out_offset += (idx_val * out.strides()[ax]);
-    }
-    update_it.seek(i * update_size);
-    for (int j = 0; j < update_size; ++j) {
-      op(updates.data<InT>()[update_it.loc],
-         out.data<InT>() + out_offset + out_it.loc);
-      update_it.step();
-      out_it.step();
-    }
-    out_it.reset();
-    update_it.reset();
+  std::vector<const IdxT*> ind_ptrs;
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(updates);
+  for (auto& idx : inds) {
+    ind_ptrs.push_back(idx.data<IdxT>());
+    encoder.set_input_array(idx);
   }
+  encoder.set_output_array(out);
+  encoder.dispatch([out_ptr = out.data<InT>(),
+                    upd_ptr = updates.data<InT>(),
+                    ind_ptrs = std::move(ind_ptrs),
+                    axes,
+                    n_updates,
+                    update_size,
+                    op = std::move(op),
+                    out_shape = out.shape(),
+                    out_strides = out.strides(),
+                    out_it = std::move(out_it),
+                    update_it = std::move(update_it),
+                    its = std::move(its)]() mutable {
+    for (int i = 0; i < n_updates; ++i) {
+      size_t out_offset = 0;
+      for (int j = 0; j < ind_ptrs.size(); ++j) {
+        auto ax = axes[j];
+        auto idx_loc = its[j].loc;
+        its[j].step();
+        auto idx_val = offset_neg_idx(ind_ptrs[j][idx_loc], out_shape[ax]);
+        out_offset += (idx_val * out_strides[ax]);
+      }
+      update_it.seek(i * update_size);
+      for (int j = 0; j < update_size; ++j) {
+        op(upd_ptr[update_it.loc], out_ptr + out_offset + out_it.loc);
+        update_it.step();
+        out_it.step();
+      }
+      out_it.reset();
+      update_it.reset();
+    }
+  });
 }
 
 template <typename InT, typename IdxT>
@@ -389,29 +455,53 @@ void dispatch_scatter_inds(
     const std::vector<array>& indices,
     const array& updates,
     const std::vector<int>& axes,
-    Scatter::ReduceType rtype) {
+    Scatter::ReduceType rtype,
+    Stream stream) {
   switch (rtype) {
     case Scatter::None:
       scatter<InT, IdxT>(
-          updates, out, indices, axes, [](auto x, auto* y) { (*y) = x; });
+          updates,
+          out,
+          indices,
+          axes,
+          [](auto x, auto* y) { (*y) = x; },
+          stream);
       break;
     case Scatter::Sum:
       scatter<InT, IdxT>(
-          updates, out, indices, axes, [](auto x, auto* y) { (*y) += x; });
+          updates,
+          out,
+          indices,
+          axes,
+          [](auto x, auto* y) { (*y) += x; },
+          stream);
       break;
     case Scatter::Prod:
       scatter<InT, IdxT>(
-          updates, out, indices, axes, [](auto x, auto* y) { (*y) *= x; });
+          updates,
+          out,
+          indices,
+          axes,
+          [](auto x, auto* y) { (*y) *= x; },
+          stream);
       break;
     case Scatter::Max:
-      scatter<InT, IdxT>(updates, out, indices, axes, [](auto x, auto* y) {
-        (*y) = (*y > x) ? *y : x;
-      });
+      scatter<InT, IdxT>(
+          updates,
+          out,
+          indices,
+          axes,
+          [](auto x, auto* y) { (*y) = (*y > x) ? *y : x; },
+          stream);
       break;
     case Scatter::Min:
-      scatter<InT, IdxT>(updates, out, indices, axes, [](auto x, auto* y) {
-        (*y) = (*y < x) ? *y : x;
-      });
+      scatter<InT, IdxT>(
+          updates,
+          out,
+          indices,
+          axes,
+          [](auto x, auto* y) { (*y) = (*y < x) ? *y : x; },
+          stream);
       break;
   }
 }
@@ -422,36 +512,46 @@ void dispatch_scatter(
     const std::vector<array>& inds,
     const array& updates,
     const std::vector<int>& axes,
-    Scatter::ReduceType rtype) {
+    Scatter::ReduceType rtype,
+    Stream stream) {
   if (inds.empty()) {
-    dispatch_scatter_inds<InT, uint8_t>(out, inds, updates, axes, rtype);
+    dispatch_scatter_inds<InT, uint8_t>(
+        out, inds, updates, axes, rtype, stream);
     return;
   }
 
   switch (inds[0].dtype()) {
     case uint8:
-      dispatch_scatter_inds<InT, uint8_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, uint8_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case uint16:
-      dispatch_scatter_inds<InT, uint16_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, uint16_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case uint32:
-      dispatch_scatter_inds<InT, uint32_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, uint32_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case uint64:
-      dispatch_scatter_inds<InT, uint64_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, uint64_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case int8:
-      dispatch_scatter_inds<InT, int8_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, int8_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case int16:
-      dispatch_scatter_inds<InT, int16_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, int16_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case int32:
-      dispatch_scatter_inds<InT, int32_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, int32_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     case int64:
-      dispatch_scatter_inds<InT, int64_t>(out, inds, updates, axes, rtype);
+      dispatch_scatter_inds<InT, int64_t>(
+          out, inds, updates, axes, rtype, stream);
       break;
     default:
       throw std::runtime_error(
@@ -469,50 +569,63 @@ void Scatter::eval_cpu(const std::vector<array>& inputs, array& out) {
   // Copy src into out (copy allocates memory for out)
   auto ctype =
       src.flags().row_contiguous ? CopyType::Vector : CopyType::General;
-  copy(src, out, ctype);
+  copy(src, out, ctype, stream());
 
   switch (src.dtype()) {
     case bool_:
-      dispatch_scatter<bool>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<bool>(out, inds, updates, axes_, reduce_type_, stream());
       break;
     case uint8:
-      dispatch_scatter<uint8_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<uint8_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case uint16:
-      dispatch_scatter<uint16_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<uint16_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case uint32:
-      dispatch_scatter<uint32_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<uint32_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case uint64:
-      dispatch_scatter<uint64_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<uint64_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case int8:
-      dispatch_scatter<int8_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<int8_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case int16:
-      dispatch_scatter<int16_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<int16_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case int32:
-      dispatch_scatter<int32_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<int32_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case int64:
-      dispatch_scatter<int64_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<int64_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case float16:
-      dispatch_scatter<float16_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<float16_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case float32:
-      dispatch_scatter<float>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<float>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case float64:
-      dispatch_scatter<double>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<double>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case bfloat16:
-      dispatch_scatter<bfloat16_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<bfloat16_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
     case complex64:
-      dispatch_scatter<complex64_t>(out, inds, updates, axes_, reduce_type_);
+      dispatch_scatter<complex64_t>(
+          out, inds, updates, axes_, reduce_type_, stream());
       break;
   }
 }
@@ -523,7 +636,8 @@ void scatter_axis(
     const array idx,
     const array& upd,
     int axis,
-    const OpT& op) {
+    const OpT& op,
+    Stream stream) {
   auto strides = idx.strides();
   strides.erase(strides.begin() + axis);
   auto shape = idx.shape();
@@ -543,6 +657,11 @@ void scatter_axis(
   auto idx_ax_size = idx.shape(axis);
   auto dst_ax_size = out.shape(axis);
 
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(idx);
+  encoder.set_input_array(upd);
+  encoder.set_output_array(out);
+
   size_t size_pre = 1;
   size_t size_post = 1;
   for (int i = 0; i < axis; ++i) {
@@ -551,20 +670,34 @@ void scatter_axis(
   for (int i = axis + 1; i < idx.ndim(); ++i) {
     size_post *= idx.shape(i);
   }
-  size_t stride_pre = size_post * dst_ax_size;
-  for (size_t i = 0; i < size_pre; i++) {
-    for (size_t k = 0; k < size_post; k++) {
-      for (int j = 0; j < idx_ax_size; ++j) {
-        auto ind_val = offset_neg_idx(
-            idx_ptr[idx_it.loc + j * idx_ax_stride], dst_ax_size);
-        op(upd_ptr[upd_it.loc + j * upd_ax_stride],
-           dst_ptr + k + ind_val * dst_ax_stride);
+  encoder.dispatch([idx_ptr,
+                    upd_ptr,
+                    dst_ptr,
+                    size_pre,
+                    size_post,
+                    idx_ax_size,
+                    dst_ax_size,
+                    idx_ax_stride,
+                    upd_ax_stride,
+                    dst_ax_stride,
+                    idx_it = std::move(idx_it),
+                    upd_it = std::move(upd_it),
+                    op = std::move(op)]() mutable {
+    size_t stride_pre = size_post * dst_ax_size;
+    for (size_t i = 0; i < size_pre; i++) {
+      for (size_t k = 0; k < size_post; k++) {
+        for (int j = 0; j < idx_ax_size; ++j) {
+          auto ind_val = offset_neg_idx(
+              idx_ptr[idx_it.loc + j * idx_ax_stride], dst_ax_size);
+          op(upd_ptr[upd_it.loc + j * upd_ax_stride],
+             dst_ptr + k + ind_val * dst_ax_stride);
+        }
+        idx_it.step();
+        upd_it.step();
       }
-      idx_it.step();
-      upd_it.step();
+      dst_ptr += stride_pre;
     }
-    dst_ptr += stride_pre;
-  }
+  });
 }
 
 template <typename InT, typename IdxT>
@@ -573,15 +706,16 @@ void dispatch_scatter_axis_op(
     const array& idx,
     const array& updates,
     int axis,
-    ScatterAxis::ReduceType rtype) {
+    ScatterAxis::ReduceType rtype,
+    Stream stream) {
   switch (rtype) {
     case ScatterAxis::None:
       scatter_axis<InT, IdxT>(
-          out, idx, updates, axis, [](auto x, auto* y) { (*y) = x; });
+          out, idx, updates, axis, [](auto x, auto* y) { (*y) = x; }, stream);
       break;
     case ScatterAxis::Sum:
       scatter_axis<InT, IdxT>(
-          out, idx, updates, axis, [](auto x, auto* y) { (*y) += x; });
+          out, idx, updates, axis, [](auto x, auto* y) { (*y) += x; }, stream);
       break;
   }
 }
@@ -592,31 +726,40 @@ void dispatch_scatter_axis(
     const array& idx,
     const array& updates,
     int axis,
-    ScatterAxis::ReduceType rtype) {
+    ScatterAxis::ReduceType rtype,
+    Stream stream) {
   switch (idx.dtype()) {
     case uint8:
-      dispatch_scatter_axis_op<InT, uint8_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, uint8_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case uint16:
-      dispatch_scatter_axis_op<InT, uint16_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, uint16_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case uint32:
-      dispatch_scatter_axis_op<InT, uint32_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, uint32_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case uint64:
-      dispatch_scatter_axis_op<InT, uint64_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, uint64_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case int8:
-      dispatch_scatter_axis_op<InT, int8_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, int8_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case int16:
-      dispatch_scatter_axis_op<InT, int16_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, int16_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case int32:
-      dispatch_scatter_axis_op<InT, int32_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, int32_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     case int64:
-      dispatch_scatter_axis_op<InT, int64_t>(out, idx, updates, axis, rtype);
+      dispatch_scatter_axis_op<InT, int64_t>(
+          out, idx, updates, axis, rtype, stream);
       break;
     default:
       throw std::runtime_error(
@@ -634,51 +777,64 @@ void ScatterAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
   // Copy src into out (copy allocates memory for out)
   auto ctype =
       src.flags().row_contiguous ? CopyType::Vector : CopyType::General;
-  copy(src, out, ctype);
+  copy(src, out, ctype, stream());
 
   switch (src.dtype()) {
     case bool_:
-      dispatch_scatter_axis<bool>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<bool>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case uint8:
-      dispatch_scatter_axis<uint8_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<uint8_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case uint16:
-      dispatch_scatter_axis<uint16_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<uint16_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case uint32:
-      dispatch_scatter_axis<uint32_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<uint32_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case uint64:
-      dispatch_scatter_axis<uint64_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<uint64_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case int8:
-      dispatch_scatter_axis<int8_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<int8_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case int16:
-      dispatch_scatter_axis<int16_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<int16_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case int32:
-      dispatch_scatter_axis<int32_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<int32_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case int64:
-      dispatch_scatter_axis<int64_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<int64_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case float16:
-      dispatch_scatter_axis<float16_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<float16_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case float32:
-      dispatch_scatter_axis<float>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<float>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case float64:
-      dispatch_scatter_axis<double>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<double>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case bfloat16:
-      dispatch_scatter_axis<bfloat16_t>(out, idx, updates, axis_, reduce_type_);
+      dispatch_scatter_axis<bfloat16_t>(
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
     case complex64:
       dispatch_scatter_axis<complex64_t>(
-          out, idx, updates, axis_, reduce_type_);
+          out, idx, updates, axis_, reduce_type_, stream());
       break;
   }
 }

mlx/backend/cpu/inverse.cpp

@@ -2,47 +2,37 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
-int strtri_wrapper(char uplo, char diag, float* matrix, int N) {
-  int info;
-  MLX_LAPACK_FUNC(strtri)
-  (
-      /* uplo = */ &uplo,
-      /* diag = */ &diag,
-      /* N = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info);
-  return info;
-}
-
 namespace mlx::core {
 
-void general_inv(array& inv, int N, int i) {
+template <typename T>
+void general_inv(T* inv, int N) {
   int info;
   auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
   // Compute LU factorization.
-  sgetrf_(
+  getrf<T>(
       /* m = */ &N,
       /* n = */ &N,
-      /* a = */ inv.data<float>() + N * N * i,
+      /* a = */ inv,
       /* lda = */ &N,
       /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
       /* info = */ &info);
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: LU factorization failed with error code " << info;
+    ss << "[Inverse::eval_cpu] LU factorization failed with error code "
+       << info;
     throw std::runtime_error(ss.str());
   }
 
   static const int lwork_query = -1;
-  float workspace_size = 0;
+  T workspace_size = 0;
 
   // Compute workspace size.
-  sgetri_(
+  getri<T>(
       /* m = */ &N,
       /* a = */ nullptr,
       /* lda = */ &N,
@@ -53,68 +43,118 @@ void general_inv(array& inv, int N, int i) {
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: LU workspace calculation failed with error code "
+    ss << "[Inverse::eval_cpu] LU workspace calculation failed with error code "
        << info;
     throw std::runtime_error(ss.str());
   }
 
   const int lwork = workspace_size;
-  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
 
   // Compute inverse.
-  sgetri_(
+  getri<T>(
       /* m = */ &N,
-      /* a = */ inv.data<float>() + N * N * i,
+      /* a = */ inv,
       /* lda = */ &N,
       /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
-      /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
+      /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
       /* lwork = */ &lwork,
       /* info = */ &info);
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: inversion failed with error code " << info;
+    ss << "[Inverse::eval_cpu] inversion failed with error code " << info;
     throw std::runtime_error(ss.str());
   }
 }
 
-void tri_inv(array& inv, int N, int i, bool upper) {
+template <typename T>
+void tri_inv(T* inv, int N, bool upper) {
   const char uplo = upper ? 'L' : 'U';
   const char diag = 'N';
-  int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
+  int info;
+  trtri<T>(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* info = */ &info);
+
+  // zero out the other triangle
+  if (upper) {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv, inv + i, 0.0f);
+      inv += N;
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv + i + 1, inv + N, 0.0f);
+      inv += N;
+    }
+  }
+
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: triangular inversion failed with error code " << info;
+    ss << "[Inverse::eval_cpu] triangular inversion failed with error code "
+       << info;
     throw std::runtime_error(ss.str());
   }
 }
 
-void inverse_impl(const array& a, array& inv, bool tri, bool upper) {
+template <typename T>
+void inverse_impl(
+    const array& a,
+    array& inv,
+    bool tri,
+    bool upper,
+    Stream stream) {
   // Lapack uses the column-major convention. We take advantage of the following
   // identity to avoid transposing (see
   // https://math.stackexchange.com/a/340234):
   //   (A⁻¹)ᵀ = (Aᵀ)⁻¹
 
   // The inverse is computed in place, so just copy the input to the output.
-  copy(a, inv, a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  copy(
+      a,
+      inv,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
 
   const int N = a.shape(-1);
   const size_t num_matrices = a.size() / (N * N);
 
-  for (int i = 0; i < num_matrices; i++) {
-    if (tri) {
-      tri_inv(inv, N, i, upper);
-    } else {
-      general_inv(inv, N, i);
-    }
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(inv);
+
+  auto inv_ptr = inv.data<T>();
+  if (tri) {
+    encoder.dispatch([inv_ptr, N, num_matrices, upper]() {
+      for (int i = 0; i < num_matrices; i++) {
+        tri_inv<T>(inv_ptr + N * N * i, N, upper);
+      }
+    });
+  } else {
+    encoder.dispatch([inv_ptr, N, num_matrices]() {
+      for (int i = 0; i < num_matrices; i++) {
+        general_inv<T>(inv_ptr + N * N * i, N);
+      }
+    });
   }
 }
 
 void Inverse::eval_cpu(const std::vector<array>& inputs, array& output) {
-  if (inputs[0].dtype() != float32) {
-    throw std::runtime_error("[Inverse::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      inverse_impl<float>(inputs[0], output, tri_, upper_, stream());
+      break;
+    case float64:
+      inverse_impl<double>(inputs[0], output, tri_, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Inverse::eval_cpu] only supports float32 or float64.");
   }
-  inverse_impl(inputs[0], output, tri_, upper_);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/lapack.h

@@ -31,3 +31,22 @@
 #define MLX_LAPACK_FUNC(f) f##_
 
 #endif
+
+#define INSTANTIATE_LAPACK_TYPES(FUNC)                       \
+  template <typename T, typename... Args>                    \
+  void FUNC(Args... args) {                                  \
+    if constexpr (std::is_same_v<T, float>) {                \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...); \
+    } else if constexpr (std::is_same_v<T, double>) {        \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...); \
+    }                                                        \
+  }
+
+INSTANTIATE_LAPACK_TYPES(geqrf)
+INSTANTIATE_LAPACK_TYPES(orgqr)
+INSTANTIATE_LAPACK_TYPES(syevd)
+INSTANTIATE_LAPACK_TYPES(potrf)
+INSTANTIATE_LAPACK_TYPES(gesvdx)
+INSTANTIATE_LAPACK_TYPES(getrf)
+INSTANTIATE_LAPACK_TYPES(getri)
+INSTANTIATE_LAPACK_TYPES(trtri)

mlx/backend/cpu/luf.cpp

@@ -4,18 +4,22 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
-void lu_factor_impl(
+template <typename T>
+void luf_impl(
     const array& a,
     array& lu,
     array& pivots,
-    array& row_indices) {
+    array& row_indices,
+    Stream stream) {
   int M = a.shape(-2);
   int N = a.shape(-1);
+  int K = std::min(M, N);
 
   // Copy a into lu and make it col contiguous
   auto ndim = lu.ndim();
@@ -29,60 +33,89 @@ void lu_factor_impl(
   lu.set_data(
       allocator::malloc_or_wait(lu.nbytes()), lu.nbytes(), strides, flags);
   copy_inplace(
-      a, lu, a.shape(), a.strides(), strides, 0, 0, CopyType::GeneralGeneral);
-
-  auto a_ptr = lu.data<float>();
+      a,
+      lu,
+      a.shape(),
+      a.strides(),
+      strides,
+      0,
+      0,
+      CopyType::GeneralGeneral,
+      stream);
 
+  auto a_ptr = lu.data<T>();
   pivots.set_data(allocator::malloc_or_wait(pivots.nbytes()));
   row_indices.set_data(allocator::malloc_or_wait(row_indices.nbytes()));
   auto pivots_ptr = pivots.data<uint32_t>();
   auto row_indices_ptr = row_indices.data<uint32_t>();
-
-  int info;
   size_t num_matrices = a.size() / (M * N);
-  for (size_t i = 0; i < num_matrices; ++i) {
-    // Compute LU factorization of A
-    MLX_LAPACK_FUNC(sgetrf)
-    (/* m */ &M,
-     /* n */ &N,
-     /* a */ a_ptr,
-     /* lda */ &M,
-     /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
-     /* info */ &info);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(lu);
+  encoder.set_output_array(pivots);
+  encoder.set_output_array(row_indices);
 
-    if (info != 0) {
-      std::stringstream ss;
-      ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-         << ((info > 0) ? " because matrix is singular"
-                        : " because argument had an illegal value");
-      throw std::runtime_error(ss.str());
-    }
+  encoder.dispatch(
+      [a_ptr, pivots_ptr, row_indices_ptr, num_matrices, M, N, K]() mutable {
+        int info;
+        for (size_t i = 0; i < num_matrices; ++i) {
+          // Compute LU factorization of A
+          getrf<T>(
+              /* m */ &M,
+              /* n */ &N,
+              /* a */ a_ptr,
+              /* lda */ &M,
+              /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
+              /* info */ &info);
 
-    // Subtract 1 to get 0-based index
-    for (int j = 0; j < pivots.shape(-1); ++j) {
-      pivots_ptr[j]--;
-      row_indices_ptr[j] = j;
-    }
-    for (int j = pivots.shape(-1) - 1; j >= 0; --j) {
-      auto piv = pivots_ptr[j];
-      auto t1 = row_indices_ptr[piv];
-      auto t2 = row_indices_ptr[j];
-      row_indices_ptr[j] = t1;
-      row_indices_ptr[piv] = t2;
-    }
+          if (info != 0) {
+            std::stringstream ss;
+            ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
+               << ((info > 0) ? " because matrix is singular"
+                              : " because argument had an illegal value");
+            throw std::runtime_error(ss.str());
+          }
 
-    // Advance pointers to the next matrix
-    a_ptr += M * N;
-    pivots_ptr += pivots.shape(-1);
-    row_indices_ptr += pivots.shape(-1);
-  }
+          // Subtract 1 to get 0-based index
+          int j = 0;
+          for (; j < K; ++j) {
+            pivots_ptr[j]--;
+            row_indices_ptr[j] = j;
+          }
+          for (; j < M; ++j) {
+            row_indices_ptr[j] = j;
+          }
+          for (int j = K - 1; j >= 0; --j) {
+            auto piv = pivots_ptr[j];
+            auto t1 = row_indices_ptr[piv];
+            auto t2 = row_indices_ptr[j];
+            row_indices_ptr[j] = t1;
+            row_indices_ptr[piv] = t2;
+          }
+
+          // Advance pointers to the next matrix
+          a_ptr += M * N;
+          pivots_ptr += K;
+          row_indices_ptr += M;
+        }
+      });
 }
 
 void LUF::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
   assert(inputs.size() == 1);
-  lu_factor_impl(inputs[0], outputs[0], outputs[1], outputs[2]);
+  switch (inputs[0].dtype()) {
+    case float32:
+      luf_impl<float>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    case float64:
+      luf_impl<double>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[LUF::eval_cpu] only supports float32 or float64.");
+  }
 }
 
 } // namespace mlx::core

mlx/backend/cpu/make_compiled_preamble.sh

@@ -18,13 +18,16 @@ if [ "$CLANG" = "TRUE" ]; then
 #include <complex>
 #include <cstdint>
 #include <vector>
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+#include <arm_fp16.h>
+#endif
 EOM
-CC_FLAGS="-arch ${ARCH}"
+CC_FLAGS="-arch ${ARCH} -nobuiltininc -nostdinc"
 else
 CC_FLAGS="-std=c++17"
 fi
 
-CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
+CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
 const char* get_kernel_preamble() {

mlx/backend/cpu/masked_mm.cpp

@@ -5,6 +5,7 @@
 #include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
@@ -64,36 +65,36 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b_pre = inputs[1];
 
   auto check_transpose =
-      [](const array& arr, bool do_copy, bool expand_all = false) {
+      [s = stream()](const array& arr, bool do_copy, bool expand_all = false) {
         auto stx = arr.strides()[arr.ndim() - 2];
         auto sty = arr.strides()[arr.ndim() - 1];
         if (!expand_all && stx == arr.shape(-1) && sty == 1) {
           if (do_copy) {
             array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-            copy(arr, arr_copy, CopyType::Vector);
-            return std::make_tuple(false, stx, arr_copy);
+            copy(arr, arr_copy, CopyType::Vector, s);
+            return std::make_tuple(false, stx, arr_copy, true);
           }
-          return std::make_tuple(false, stx, arr);
+          return std::make_tuple(false, stx, arr, false);
         } else if (!expand_all && stx == 1 && sty == arr.shape(-2)) {
           if (do_copy) {
             array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-            copy(arr, arr_copy, CopyType::Vector);
-            return std::make_tuple(true, sty, arr_copy);
+            copy(arr, arr_copy, CopyType::Vector, s);
+            return std::make_tuple(true, sty, arr_copy, true);
           }
-          return std::make_tuple(true, sty, arr);
+          return std::make_tuple(true, sty, arr, false);
         } else {
           array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-          copy(arr, arr_copy, CopyType::General);
+          copy(arr, arr_copy, CopyType::General, s);
           int64_t stx = arr.shape(-1);
-          return std::make_tuple(false, stx, arr_copy);
+          return std::make_tuple(false, stx, arr_copy, true);
         }
       };
 
   bool has_op_mask = inputs.size() > 3;
   bool has_out_mask = inputs.size() == 3 || inputs.size() == 5;
-  auto [a_transposed, lda, a] =
+  auto [a_transposed, lda, a, a_copied] =
       check_transpose(a_pre, has_op_mask, inputs.back().dtype() != bool_);
-  auto [b_transposed, ldb, b] =
+  auto [b_transposed, ldb, b, b_copied] =
       check_transpose(b_pre, has_op_mask, inputs.back().dtype() != bool_);
 
   size_t M = a.shape(-2);
@@ -104,31 +105,39 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
+  auto& encoder = cpu::get_command_encoder(stream());
   if (K == 0) {
-    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(), nbytes = out.nbytes()]() {
+      std::memset(out_ptr, 0, nbytes);
+    });
     return;
   }
 
-  auto mask_array = [](const array& mask,
+  auto mask_array = [](const void* mask,
                        float* data,
                        int block_size,
                        int batch_idx,
                        int X,
                        int Y,
                        size_t X_data_str,
-                       size_t Y_data_str) {
+                       size_t Y_data_str,
+                       const Shape& mask_shape,
+                       const Strides& mask_strides,
+                       bool is_bool) {
+    auto ndim = mask_shape.size();
     auto mask_offset = elem_to_loc(
-        mask.shape(-1) * mask.shape(-2) * batch_idx,
-        mask.shape(),
-        mask.strides());
+        mask_shape[ndim - 1] * mask_shape[ndim - 2] * batch_idx,
+        mask_shape,
+        mask_strides);
 
-    auto X_mask_str = mask.strides()[mask.ndim() - 2];
-    auto Y_mask_str = mask.strides()[mask.ndim() - 1];
+    auto X_mask_str = mask_strides[ndim - 2];
+    auto Y_mask_str = mask_strides[ndim - 1];
 
-    if (mask.dtype() == bool_) {
+    if (is_bool) {
       return mask_matrix(
           data,
-          mask.data<bool>(),
+          static_cast<const bool*>(mask),
           block_size,
           X,
           Y,
@@ -140,7 +149,7 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     } else {
       return mask_matrix(
           data,
-          mask.data<float>(),
+          static_cast<const float*>(mask),
           block_size,
           X,
           Y,
@@ -152,61 +161,155 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     }
   };
 
-  for (int i = 0; i < (out.size() / (M * size_t(N))); ++i) {
-    // Adjust pointer
-    float* ai =
-        a.data<float>() + elem_to_loc(M * K * i, a.shape(), a.strides());
-    float* bi =
-        b.data<float>() + elem_to_loc(K * N * i, b.shape(), b.strides());
-    float* ci = out.data<float>() + M * N * i;
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  const void* a_mask_ptr;
+  const void* b_mask_ptr;
+  const void* out_mask_ptr;
+  Shape a_mask_shape;
+  Shape b_mask_shape;
+  Shape out_mask_shape;
+  Strides a_mask_strides;
+  Strides b_mask_strides;
+  Strides out_mask_strides;
+  bool a_mask_bool;
+  bool b_mask_bool;
+  bool out_mask_bool;
+  if (has_op_mask) {
+    auto& a_mask = inputs[inputs.size() - 2];
+    auto& b_mask = inputs[inputs.size() - 1];
+    a_mask_ptr = a_mask.data<void>();
+    b_mask_ptr = b_mask.data<void>();
+    a_mask_shape = a_mask.shape();
+    b_mask_shape = b_mask.shape();
+    a_mask_strides = a_mask.strides();
+    b_mask_strides = b_mask.strides();
+    a_mask_bool = (a_mask.dtype() == bool_);
+    b_mask_bool = (b_mask.dtype() == bool_);
+    encoder.set_input_array(a_mask);
+    encoder.set_input_array(b_mask);
+  }
+  if (has_out_mask) {
+    auto& out_mask = inputs[2];
+    out_mask_ptr = out_mask.data<void>();
+    out_mask_bool = (out_mask.dtype() == bool_);
+    encoder.set_input_array(out_mask);
+    out_mask_shape = out_mask.shape();
+    out_mask_strides = out_mask.strides();
+  }
+  encoder.set_output_array(out);
+  auto a_ptr = a.data<float>();
+  auto b_ptr = b.data<float>();
+  auto out_ptr = out.data<float>();
+  size_t num_matrices = out.size() / (M * size_t(N));
+  auto ldc = out.shape(-1);
 
-    // Zero out blocks in a and b if needed
-    if (has_op_mask) {
-      auto& a_mask = inputs[inputs.size() - 2];
-      mask_array(
-          a_mask,
-          ai,
-          block_size_,
-          i,
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_mask_ptr,
+                    b_mask_ptr,
+                    out_mask_ptr,
+                    has_op_mask,
+                    has_out_mask,
+                    block_size = block_size_,
+                    num_matrices,
+                    M,
+                    N,
+                    K,
+                    a_transposed = a_transposed,
+                    b_transposed = b_transposed,
+                    lda = lda,
+                    ldb = ldb,
+                    ldc,
+                    a_shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_shape = b.shape(),
+                    b_strides = b.strides(),
+                    a_mask_shape = std::move(a_mask_shape),
+                    b_mask_shape = std::move(b_mask_shape),
+                    out_mask_shape = std::move(out_mask_shape),
+                    a_mask_strides = std::move(a_mask_strides),
+                    b_mask_strides = std::move(b_mask_strides),
+                    out_mask_strides = std::move(out_mask_strides),
+                    mask_array,
+                    a_mask_bool,
+                    b_mask_bool,
+                    out_mask_bool]() {
+    for (int i = 0; i < num_matrices; ++i) {
+      // Adjust pointer
+      float* ai = a_ptr + elem_to_loc(M * K * i, a_shape, a_strides);
+      float* bi = b_ptr + elem_to_loc(K * N * i, b_shape, b_strides);
+      float* ci = out_ptr + M * N * i;
+
+      // Zero out blocks in a and b if needed
+      if (has_op_mask) {
+        mask_array(
+            a_mask_ptr,
+            ai,
+            block_size,
+            i,
+            M,
+            K,
+            a_transposed ? 1 : lda,
+            a_transposed ? lda : 1,
+            a_mask_shape,
+            a_mask_strides,
+            a_mask_bool);
+
+        mask_array(
+            b_mask_ptr,
+            bi,
+            block_size,
+            i,
+            K,
+            N,
+            b_transposed ? 1 : ldb,
+            b_transposed ? ldb : 1,
+            b_mask_shape,
+            b_mask_strides,
+            b_mask_bool);
+      }
+
+      // Do matmul
+      cblas_sgemm(
+          CblasRowMajor,
+          a_transposed ? CblasTrans : CblasNoTrans, // transA
+          b_transposed ? CblasTrans : CblasNoTrans, // transB
           M,
-          K,
-          a_transposed ? 1 : lda,
-          a_transposed ? lda : 1);
-
-      auto& b_mask = inputs[inputs.size() - 1];
-      mask_array(
-          b_mask,
-          bi,
-          block_size_,
-          i,
-          K,
           N,
-          b_transposed ? 1 : ldb,
-          b_transposed ? ldb : 1);
-    }
+          K,
+          1.0, // alpha
+          ai,
+          lda,
+          bi,
+          ldb,
+          0.0, // beta
+          ci,
+          ldc);
 
-    // Do matmul
-    cblas_sgemm(
-        CblasRowMajor,
-        a_transposed ? CblasTrans : CblasNoTrans, // transA
-        b_transposed ? CblasTrans : CblasNoTrans, // transB
-        M,
-        N,
-        K,
-        1.0, // alpha
-        ai,
-        lda,
-        bi,
-        ldb,
-        0.0, // beta
-        ci,
-        out.shape(-1) // ldc
-    );
-
-    // Zero out blocks in out
-    if (has_out_mask) {
-      mask_array(inputs[2], ci, block_size_, i, M, N, N, 1);
+      // Zero out blocks in out
+      if (has_out_mask) {
+        mask_array(
+            out_mask_ptr,
+            ci,
+            block_size,
+            i,
+            M,
+            N,
+            N,
+            1,
+            out_mask_shape,
+            out_mask_strides,
+            out_mask_bool);
+      }
     }
+  });
+  if (a_copied) {
+    encoder.add_temporary(a);
+  }
+  if (b_copied) {
+    encoder.add_temporary(b);
   }
 }
 
@@ -220,7 +323,8 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& a_pre = inputs[0];
   auto& b_pre = inputs[1];
 
-  auto check_transpose = [](const array& arr) {
+  std::vector<array> temps;
+  auto check_transpose = [s = stream(), &temps](const array& arr) {
     auto stx = arr.strides()[arr.ndim() - 2];
     auto sty = arr.strides()[arr.ndim() - 1];
     if (stx == arr.shape(-1) && sty == 1) {
@@ -228,10 +332,10 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     } else if (stx == 1 && sty == arr.shape(-2)) {
       return std::make_tuple(true, sty, arr);
     } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy(arr, temps.back(), CopyType::General, s);
       int64_t stx = arr.shape(-1);
-      return std::make_tuple(false, stx, arr_copy);
+      return std::make_tuple(false, stx, temps.back());
     }
   };
 
@@ -246,8 +350,12 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
+  auto& encoder = cpu::get_command_encoder(stream());
   if (K == 0) {
-    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<float>(), size = out.size()]() {
+      std::fill_n(out_ptr, size, 0);
+    });
     return;
   }
 
@@ -272,29 +380,61 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   const uint32_t* lhs_indices_ptr = lhs_indices.data<uint32_t>();
   const uint32_t* rhs_indices_ptr = rhs_indices.data<uint32_t>();
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(lhs_indices);
+  encoder.set_input_array(rhs_indices);
+  encoder.set_output_array(out);
+  auto ldc = out.shape(-1);
 
-  for (int i = 0; i < batch_size_out; i++) {
-    // Get index
-    uint32_t indx_A = lhs_indices_ptr[elem_to_loc(i, lhs_indices)];
-    uint32_t indx_B = rhs_indices_ptr[elem_to_loc(i, rhs_indices)];
+  encoder.dispatch([a_ptr = a.data<float>(),
+                    b_ptr = b.data<float>(),
+                    out_ptr = out.data<float>(),
+                    M,
+                    N,
+                    K,
+                    lda = lda,
+                    ldb = ldb,
+                    a_transposed = a_transposed,
+                    b_transposed = b_transposed,
+                    ldc,
+                    lhs_indices_ptr,
+                    rhs_indices_ptr,
+                    lhs_indices_shape = lhs_indices.shape(),
+                    lhs_indices_strides = lhs_indices.strides(),
+                    rhs_indices_shape = rhs_indices.shape(),
+                    rhs_indices_strides = rhs_indices.strides(),
+                    batch_size_out,
+                    matrix_stride_out,
+                    batch_shape_A = std::move(batch_shape_A),
+                    batch_shape_B = std::move(batch_shape_B),
+                    batch_strides_A = std::move(batch_strides_A),
+                    batch_strides_B = std::move(batch_strides_B)]() {
+    for (int i = 0; i < batch_size_out; i++) {
+      // Get index
+      uint32_t indx_A = lhs_indices_ptr[elem_to_loc(
+          i, lhs_indices_shape, lhs_indices_strides)];
+      uint32_t indx_B = rhs_indices_ptr[elem_to_loc(
+          i, rhs_indices_shape, rhs_indices_strides)];
 
-    cblas_sgemm(
-        CblasRowMajor,
-        a_transposed ? CblasTrans : CblasNoTrans, // transA
-        b_transposed ? CblasTrans : CblasNoTrans, // transB
-        M,
-        N,
-        K,
-        1.0f, // alpha
-        a.data<float>() + elem_to_loc(indx_A, batch_shape_A, batch_strides_A),
-        lda,
-        b.data<float>() + elem_to_loc(indx_B, batch_shape_B, batch_strides_B),
-        ldb,
-        0.0f, // beta
-        out.data<float>() + matrix_stride_out * i,
-        out.shape(-1) // ldc
-    );
-  }
+      cblas_sgemm(
+          CblasRowMajor,
+          a_transposed ? CblasTrans : CblasNoTrans, // transA
+          b_transposed ? CblasTrans : CblasNoTrans, // transB
+          M,
+          N,
+          K,
+          1.0f, // alpha
+          a_ptr + elem_to_loc(indx_A, batch_shape_A, batch_strides_A),
+          lda,
+          b_ptr + elem_to_loc(indx_B, batch_shape_B, batch_strides_B),
+          ldb,
+          0.0f, // beta
+          out_ptr + matrix_stride_out * i,
+          ldc);
+    }
+  });
+  encoder.add_temporaries(std::move(temps));
 }
 
 } // namespace mlx::core

mlx/backend/cpu/matmul.cpp

@@ -3,18 +3,76 @@
 #include <cstring>
 #include "mlx/array.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/gemm.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
+template <typename T>
+void matmul_dispatch(
+    const array& a,
+    const array& b,
+    array& out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    float alpha,
+    float beta,
+    Stream stream) {
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  T* out_ptr = out.data<T>();
+  size_t ldc = out.shape(-1);
+  size_t batch_size = a.size() / (a.shape(-2) * a.shape(-1));
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_transposed,
+                    b_transposed,
+                    lda,
+                    ldb,
+                    ldc,
+                    alpha,
+                    beta,
+                    batch_size,
+                    a_shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_shape = b.shape(),
+                    b_strides = b.strides()]() {
+    matmul<T>(
+        a_ptr,
+        b_ptr,
+        out_ptr,
+        a_transposed,
+        b_transposed,
+        lda,
+        ldb,
+        ldc,
+        alpha,
+        beta,
+        batch_size,
+        a_shape,
+        a_strides,
+        b_shape,
+        b_strides);
+  });
+}
+
 void matmul_general(
     const array& a_pre,
     const array& b_pre,
     array& out,
+    Stream stream,
     float alpha = 1.0f,
     float beta = 0.0f) {
-  auto check_transpose = [](const array& arr) {
+  std::vector<array> temps;
+  auto check_transpose = [stream, &temps](const array& arr) {
     auto stx = arr.strides()[arr.ndim() - 2];
     auto sty = arr.strides()[arr.ndim() - 1];
     if (stx == arr.shape(-1) && sty == 1) {
@@ -22,10 +80,10 @@ void matmul_general(
     } else if (stx == 1 && sty == arr.shape(-2)) {
       return std::make_tuple(true, sty, arr);
     } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy(arr, temps.back(), CopyType::General, stream);
       stx = arr.shape(-1);
-      return std::make_tuple(false, stx, arr_copy);
+      return std::make_tuple(false, stx, temps.back());
     }
   };
 
@@ -39,28 +97,34 @@ void matmul_general(
   }
 
   if (out.dtype() == float32) {
-    matmul<float>(a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    matmul_dispatch<float>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
   } else if (out.dtype() == float16) {
-    matmul<float16_t>(
-        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    matmul_dispatch<float16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
   } else if (out.dtype() == bfloat16) {
-    matmul<bfloat16_t>(
-        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    matmul_dispatch<bfloat16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
   } else if (out.dtype() == float64) {
-    matmul<double>(
-        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta);
+    matmul_dispatch<double>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
   } else {
     throw std::runtime_error("[Matmul::eval_cpu] Invalid type.");
   }
+  cpu::get_command_encoder(stream).add_temporaries(std::move(temps));
 }
 
 void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
   if (inputs[0].shape(-1) == 0) {
-    std::memset(out.data<void>(), 0, out.nbytes());
+    auto& encoder = cpu::get_command_encoder(stream());
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(), nbytes = out.nbytes()]() {
+      std::memset(out_ptr, 0, nbytes);
+    });
     return;
   }
-  return matmul_general(inputs[0], inputs[1], out);
+  matmul_general(inputs[0], inputs[1], out, stream());
 }
 
 void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -74,9 +138,9 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   CopyType ctype = c.data_size() == 1
       ? CopyType::Scalar
       : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
-  copy(c, out, ctype);
+  copy(c, out, ctype, stream());
 
-  return matmul_general(inputs[0], inputs[1], out, alpha_, beta_);
+  matmul_general(inputs[0], inputs[1], out, stream(), alpha_, beta_);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/primitives.cpp

@@ -7,11 +7,11 @@
 #include <sstream>
 
 #include "mlx/allocator.h"
-#include "mlx/backend/common/load.h"
 #include "mlx/backend/common/slicing.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/arange.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/threefry.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
@@ -22,39 +22,58 @@ void reshape(const array& in, array& out) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
   if (copy_necessary) {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    copy_inplace(in, out, CopyType::General);
+    copy_inplace(in, out, CopyType::General, out.primitive().stream());
   } else {
     shared_buffer_reshape(in, out_strides, out);
   }
 }
 
-int64_t compute_dynamic_offset(
+static std::pair<array, bool> compute_dynamic_offset(
     const array& indices,
     const Strides& strides,
-    const std::vector<int>& axes) {
-  auto compute_offset = [&strides, &axes](const auto* indices) {
-    int64_t offset = 0;
-    for (int i = 0; i < axes.size(); ++i) {
-      offset += indices[i] * strides[axes[i]];
-    }
-    return offset;
-  };
+    const std::vector<int>& axes,
+    Stream stream) {
+  array offset({1}, int64, nullptr, {});
+  bool donate = indices.is_donatable() &&
+      (indices.data_size() * indices.itemsize()) >= offset.itemsize();
+  if (donate) {
+    offset.copy_shared_buffer(indices);
+  } else {
+    offset.set_data(allocator::malloc_or_wait(offset.itemsize()));
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(indices);
+  encoder.set_output_array(offset);
+  auto compute_offset =
+      [strides, axes, offset = offset.data<int64_t>()](const auto* indices) {
+        int64_t offset_ = 0;
+        for (int i = 0; i < axes.size(); ++i) {
+          offset_ += indices[i] * strides[axes[i]];
+        }
+        offset[0] = offset_;
+      };
   switch (indices.dtype()) {
     case int8:
     case uint8:
-      return compute_offset(indices.data<uint8_t>());
+      encoder.dispatch(compute_offset, indices.data<uint8_t>());
+      break;
     case int16:
     case uint16:
-      return compute_offset(indices.data<uint16_t>());
+      encoder.dispatch(compute_offset, indices.data<uint16_t>());
+      break;
     case int32:
     case uint32:
-      return compute_offset(indices.data<uint32_t>());
+      encoder.dispatch(compute_offset, indices.data<uint32_t>());
+      break;
     case int64:
     case uint64:
-      return compute_offset(indices.data<uint64_t>());
+      encoder.dispatch(compute_offset, indices.data<uint64_t>());
+      break;
     default:
       throw std::runtime_error("Invalid indices type.");
   }
+  return {offset, donate};
 }
 
 void AsStrided::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -104,14 +123,59 @@ void Transpose::eval_cpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Arange::eval_cpu(const std::vector<array>& inputs, array& out) {
-  arange(inputs, out, start_, step_);
+  assert(inputs.size() == 0);
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  switch (out.dtype()) {
+    case bool_:
+      throw std::runtime_error("Bool type unsupported for arange.");
+      break;
+    case uint8:
+      arange<uint8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint16:
+      arange<uint16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint32:
+      arange<uint32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint64:
+      arange<uint64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int8:
+      arange<int8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int16:
+      arange<int16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int32:
+      arange<int32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int64:
+      arange<int64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float16:
+      arange<float16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float32:
+      arange<float>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float64:
+      arange<double>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case bfloat16:
+      arange<bfloat16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case complex64:
+      arange<complex64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+  }
 }
 
 void AsType::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General;
-  copy(in, out, ctype);
+  copy(in, out, ctype, stream());
 }
 
 void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -134,7 +198,7 @@ void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
     size_t data_offset = strides[axis_] * sizes[i];
     out_slice.copy_shared_buffer(
         out, strides, flags, out_slice.size(), data_offset);
-    copy_inplace(inputs[i], out_slice, CopyType::GeneralGeneral);
+    copy_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, stream());
   }
 }
 
@@ -145,7 +209,7 @@ void Contiguous::eval_cpu(const std::vector<array>& inputs, array& out) {
       (allow_col_major_ && in.flags().col_contiguous)) {
     out.copy_shared_buffer(in);
   } else {
-    copy(in, out, CopyType::General);
+    copy(in, out, CopyType::General, stream());
   }
 }
 
@@ -169,14 +233,7 @@ void Full::eval_cpu(const std::vector<array>& inputs, array& out) {
   } else {
     ctype = CopyType::General;
   }
-  copy(in, out, ctype);
-}
-
-void Load::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 0);
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  load(out, offset_, reader_, swap_endianness_);
+  copy(in, out, ctype, stream());
 }
 
 void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -192,7 +249,7 @@ void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(val.dtype() == in.dtype() && in.dtype() == out.dtype());
 
   // Fill output with val
-  copy(val, out, CopyType::Scalar);
+  copy(val, out, CopyType::Scalar, stream());
 
   // Find offset for start of input values
   size_t data_offset = 0;
@@ -207,7 +264,7 @@ void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
       out, out.strides(), out.flags(), out_slice.size(), data_offset);
 
   // Copy input values into the slice
-  copy_inplace(in, out_slice, CopyType::GeneralGeneral);
+  copy_inplace(in, out_slice, CopyType::GeneralGeneral, stream());
 }
 
 void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -223,39 +280,49 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   auto kptr = inputs[0].data<uint32_t>();
   auto cptr = out.data<char>();
-  size_t out_skip = (bytes_per_key + 4 - 1) / 4;
-  auto half_size = out_skip / 2;
-  bool even = out_skip % 2 == 0;
-  for (int i = 0; i < num_keys; ++i, cptr += bytes_per_key) {
-    auto ptr = reinterpret_cast<uint32_t*>(cptr);
-    // Get ith key
-    auto kidx = 2 * i;
-    auto k1_elem = elem_to_loc(kidx, keys.shape(), keys.strides());
-    auto k2_elem = elem_to_loc(kidx + 1, keys.shape(), keys.strides());
-    auto key = std::make_pair(kptr[k1_elem], kptr[k2_elem]);
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(inputs[0]);
+  encoder.set_output_array(out);
+  encoder.dispatch([kptr,
+                    cptr,
+                    bytes_per_key,
+                    num_keys,
+                    kshape = keys.shape(),
+                    kstrides = keys.strides()]() mutable {
+    size_t out_skip = (bytes_per_key + 4 - 1) / 4;
+    auto half_size = out_skip / 2;
+    bool even = out_skip % 2 == 0;
+    for (int i = 0; i < num_keys; ++i, cptr += bytes_per_key) {
+      auto ptr = reinterpret_cast<uint32_t*>(cptr);
+      // Get ith key
+      auto kidx = 2 * i;
+      auto k1_elem = elem_to_loc(kidx, kshape, kstrides);
+      auto k2_elem = elem_to_loc(kidx + 1, kshape, kstrides);
+      auto key = std::make_pair(kptr[k1_elem], kptr[k2_elem]);
 
-    std::pair<uintptr_t, uintptr_t> count{0, half_size + !even};
-    for (; count.first + 1 < half_size; count.first++, count.second++) {
-      std::tie(ptr[count.first], ptr[count.second]) =
-          random::threefry2x32_hash(key, count);
-    }
-    if (count.first < half_size) {
-      auto rb = random::threefry2x32_hash(key, count);
-      ptr[count.first++] = rb.first;
-      if (bytes_per_key % 4 > 0) {
-        std::copy(
-            reinterpret_cast<char*>(&rb.second),
-            reinterpret_cast<char*>(&rb.second) + bytes_per_key % 4,
-            cptr + 4 * count.second);
-      } else {
-        ptr[count.second] = rb.second;
+      std::pair<uintptr_t, uintptr_t> count{0, half_size + !even};
+      for (; count.first + 1 < half_size; count.first++, count.second++) {
+        std::tie(ptr[count.first], ptr[count.second]) =
+            random::threefry2x32_hash(key, count);
+      }
+      if (count.first < half_size) {
+        auto rb = random::threefry2x32_hash(key, count);
+        ptr[count.first++] = rb.first;
+        if (bytes_per_key % 4 > 0) {
+          std::copy(
+              reinterpret_cast<char*>(&rb.second),
+              reinterpret_cast<char*>(&rb.second) + bytes_per_key % 4,
+              cptr + 4 * count.second);
+        } else {
+          ptr[count.second] = rb.second;
+        }
+      }
+      if (!even) {
+        count.second = 0;
+        ptr[half_size] = random::threefry2x32_hash(key, count).first;
       }
     }
-    if (!even) {
-      count.second = 0;
-      ptr[half_size] = random::threefry2x32_hash(key, count).first;
-    }
-  }
+  });
 }
 
 void Reshape::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -269,16 +336,23 @@ void DynamicSlice::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
   auto& in = inputs[0];
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  auto i_offset = compute_dynamic_offset(inputs[1], in.strides(), axes_);
+  auto [in_offset, donated] =
+      compute_dynamic_offset(inputs[1], in.strides(), axes_, stream());
   copy_inplace(
       /* const array& src = */ in,
       /* array& dst = */ out,
       /* const Shape& data_shape = */ out.shape(),
       /* const Strides& i_strides = */ in.strides(),
       /* const Strides& o_strides = */ out.strides(),
-      /* int64_t i_offset = */ i_offset,
+      /* int64_t i_offset = */ 0,
       /* int64_t o_offset = */ 0,
-      /* CopyType ctype = */ CopyType::GeneralGeneral);
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ in_offset,
+      /* const std::optional<array>& dynamic_o_offset = */ std::nullopt);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(in_offset));
+  }
 }
 
 void DynamicSliceUpdate::eval_cpu(
@@ -296,9 +370,10 @@ void DynamicSliceUpdate::eval_cpu(
   auto ctype = in.flags().contiguous && in.size() == in.data_size()
       ? CopyType::Vector
       : CopyType::General;
-  copy(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype);
+  copy(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
 
-  auto o_offset = compute_dynamic_offset(inputs[2], out.strides(), axes_);
+  auto [out_offset, donated] =
+      compute_dynamic_offset(inputs[2], out.strides(), axes_, stream());
   copy_inplace(
       /* const array& src = */ upd,
       /* array& dst = */ out,
@@ -306,8 +381,14 @@ void DynamicSliceUpdate::eval_cpu(
       /* const std::vector<stride_t>& i_strides = */ upd.strides(),
       /* const std::vector<stride_t>& o_strides = */ out.strides(),
       /* int64_t i_offset = */ 0,
-      /* int64_t o_offset = */ o_offset,
-      /* CopyType ctype = */ CopyType::GeneralGeneral);
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ std::nullopt,
+      /* const std::optional<array>& dynamic_o_offset = */ out_offset);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(out_offset));
+  }
 }
 
 void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -329,7 +410,7 @@ void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto ctype = in.flags().contiguous && in.size() == in.data_size()
       ? CopyType::Vector
       : CopyType::General;
-  copy(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype);
+  copy(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
 
   // Calculate out strides, initial offset and if copy needs to be made
   auto [data_offset, out_strides] =
@@ -344,7 +425,8 @@ void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
       /* const std::vector<stride_t>& o_strides = */ out_strides,
       /* int64_t i_offset = */ 0,
       /* int64_t o_offset = */ data_offset,
-      /* CopyType ctype = */ CopyType::GeneralGeneral);
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream());
 }
 
 void View::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -372,9 +454,9 @@ void View::eval_cpu(const std::vector<array>& inputs, array& out) {
     if (in.dtype() == bool_) {
       auto in_tmp = array(in.shape(), uint8, nullptr, {});
       in_tmp.copy_shared_buffer(in);
-      copy_inplace(in_tmp, tmp, CopyType::General);
+      copy_inplace(in_tmp, tmp, CopyType::General, stream());
     } else {
-      copy_inplace(in, tmp, CopyType::General);
+      copy_inplace(in, tmp, CopyType::General, stream());
     }
 
     auto flags = out.flags();
@@ -382,7 +464,7 @@ void View::eval_cpu(const std::vector<array>& inputs, array& out) {
     flags.row_contiguous = true;
     auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
     flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
-    out.move_shared_buffer(tmp, out.strides(), flags, out.size());
+    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
   }
 }
 

mlx/backend/cpu/qrf.cpp

@@ -2,53 +2,21 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
 template <typename T>
-struct lpack;
-
-template <>
-struct lpack<float> {
-  static void xgeqrf(
-      const int* m,
-      const int* n,
-      float* a,
-      const int* lda,
-      float* tau,
-      float* work,
-      const int* lwork,
-      int* info) {
-    sgeqrf_(m, n, a, lda, tau, work, lwork, info);
-  }
-  static void xorgqr(
-      const int* m,
-      const int* n,
-      const int* k,
-      float* a,
-      const int* lda,
-      const float* tau,
-      float* work,
-      const int* lwork,
-      int* info) {
-    sorgqr_(m, n, k, a, lda, tau, work, lwork, info);
-  }
-};
-
-template <typename T>
-void qrf_impl(const array& a, array& q, array& r) {
+void qrf_impl(const array& a, array& q, array& r, Stream stream) {
   const int M = a.shape(-2);
   const int N = a.shape(-1);
   const int lda = M;
   size_t num_matrices = a.size() / (M * N);
-  int num_reflectors = std::min(M, N);
-  auto tau =
-      allocator::malloc_or_wait(sizeof(T) * num_matrices * num_reflectors);
 
   // Copy A to inplace input and make it col-contiguous
-  array in(a.shape(), float32, nullptr, {});
+  array in(a.shape(), a.dtype(), nullptr, {});
   auto flags = in.flags();
 
   // Copy the input to be column contiguous
@@ -59,103 +27,123 @@ void qrf_impl(const array& a, array& q, array& r) {
   strides[in.ndim() - 1] = M;
   in.set_data(
       allocator::malloc_or_wait(in.nbytes()), in.nbytes(), strides, flags);
-  copy_inplace(a, in, CopyType::GeneralGeneral);
-
-  T optimal_work;
-  int lwork = -1;
-  int info;
-
-  // Compute workspace size
-  lpack<T>::xgeqrf(
-      &M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
-
-  // Update workspace size
-  lwork = optimal_work;
-  auto work = allocator::malloc_or_wait(sizeof(T) * lwork);
-
-  // Loop over matrices
-  for (int i = 0; i < num_matrices; ++i) {
-    // Solve
-    lpack<T>::xgeqrf(
-        &M,
-        &N,
-        in.data<float>() + M * N * i,
-        &lda,
-        static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
-        static_cast<T*>(work.raw_ptr()),
-        &lwork,
-        &info);
-  }
-  allocator::free(work);
-
+  copy_inplace(a, in, CopyType::GeneralGeneral, stream);
+  auto& encoder = cpu::get_command_encoder(stream);
+  q.set_data(allocator::malloc_or_wait(q.nbytes()));
   r.set_data(allocator::malloc_or_wait(r.nbytes()));
 
-  for (int i = 0; i < num_matrices; ++i) {
-    /// num_reflectors x N
-    for (int j = 0; j < r.shape(-2); ++j) {
-      for (int k = 0; k < j; ++k) {
-        r.data<T>()[i * N * num_reflectors + j * N + k] = 0;
-      }
-      for (int k = j; k < r.shape(-1); ++k) {
-        r.data<T>()[i * N * num_reflectors + j * N + k] =
-            in.data<T>()[i * N * M + j + k * M];
+  auto in_ptr = in.data<T>();
+  auto r_ptr = r.data<T>();
+  auto q_ptr = q.data<T>();
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(q);
+  encoder.set_output_array(r);
+  encoder.dispatch([in_ptr, q_ptr, r_ptr, M, N, lda, num_matrices]() {
+    int num_reflectors = std::min(M, N);
+    auto tau =
+        allocator::malloc_or_wait(sizeof(T) * num_matrices * num_reflectors);
+
+    T optimal_work;
+    int lwork = -1;
+    int info;
+
+    // Compute workspace size
+    geqrf<T>(&M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
+
+    // Update workspace size
+    lwork = optimal_work;
+    auto work = allocator::malloc_or_wait(sizeof(T) * lwork);
+
+    // Loop over matrices
+    for (int i = 0; i < num_matrices; ++i) {
+      // Solve
+      geqrf<T>(
+          &M,
+          &N,
+          in_ptr + M * N * i,
+          &lda,
+          static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+          static_cast<T*>(work.raw_ptr()),
+          &lwork,
+          &info);
+    }
+    allocator::free(work);
+
+    for (int i = 0; i < num_matrices; ++i) {
+      /// num_reflectors x N
+      for (int j = 0; j < num_reflectors; ++j) {
+        for (int k = 0; k < j; ++k) {
+          r_ptr[i * N * num_reflectors + j * N + k] = 0;
+        }
+        for (int k = j; k < N; ++k) {
+          r_ptr[i * N * num_reflectors + j * N + k] =
+              in_ptr[i * N * M + j + k * M];
+        }
       }
     }
-  }
 
-  // Get work size
-  lwork = -1;
-  lpack<T>::xorgqr(
-      &M,
-      &num_reflectors,
-      &num_reflectors,
-      nullptr,
-      &lda,
-      nullptr,
-      &optimal_work,
-      &lwork,
-      &info);
-  lwork = optimal_work;
-  work = allocator::malloc_or_wait(sizeof(T) * lwork);
-
-  // Loop over matrices
-  for (int i = 0; i < num_matrices; ++i) {
-    // Compute Q
-    lpack<T>::xorgqr(
+    // Get work size
+    lwork = -1;
+    orgqr<T>(
         &M,
         &num_reflectors,
         &num_reflectors,
-        in.data<float>() + M * N * i,
+        nullptr,
         &lda,
-        static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
-        static_cast<T*>(work.raw_ptr()),
+        nullptr,
+        &optimal_work,
         &lwork,
         &info);
-  }
+    lwork = optimal_work;
+    work = allocator::malloc_or_wait(sizeof(T) * lwork);
 
-  q.set_data(allocator::malloc_or_wait(q.nbytes()));
-  for (int i = 0; i < num_matrices; ++i) {
-    // M x num_reflectors
-    for (int j = 0; j < q.shape(-2); ++j) {
-      for (int k = 0; k < q.shape(-1); ++k) {
-        q.data<T>()[i * M * num_reflectors + j * num_reflectors + k] =
-            in.data<T>()[i * N * M + j + k * M];
+    // Loop over matrices
+    for (int i = 0; i < num_matrices; ++i) {
+      // Compute Q
+      orgqr<T>(
+          &M,
+          &num_reflectors,
+          &num_reflectors,
+          in_ptr + M * N * i,
+          &lda,
+          static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+          static_cast<T*>(work.raw_ptr()),
+          &lwork,
+          &info);
+    }
+
+    for (int i = 0; i < num_matrices; ++i) {
+      // M x num_reflectors
+      for (int j = 0; j < M; ++j) {
+        for (int k = 0; k < num_reflectors; ++k) {
+          q_ptr[i * M * num_reflectors + j * num_reflectors + k] =
+              in_ptr[i * N * M + j + k * M];
+        }
       }
     }
-  }
 
-  // Cleanup
-  allocator::free(work);
-  allocator::free(tau);
+    // Cleanup
+    allocator::free(work);
+    allocator::free(tau);
+  });
+  encoder.add_temporary(in);
 }
 
 void QRF::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  if (!(inputs[0].dtype() == float32)) {
-    throw std::runtime_error("[QRF::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      qrf_impl<float>(inputs[0], outputs[0], outputs[1], stream());
+      break;
+    case float64:
+      qrf_impl<double>(inputs[0], outputs[0], outputs[1], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[QRF::eval_cpu] only supports float32 or float64.");
   }
-  qrf_impl<float>(inputs[0], outputs[0], outputs[1]);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/quantized.cpp

@@ -3,6 +3,7 @@
 #include <cassert>
 
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 #include "mlx/fast_primitives.h"
 #include "mlx/primitives.h"
@@ -164,8 +165,8 @@ simd::Simd<uint32_t, S> extract_bits_simd(const uint32_t* w) {
   } else if constexpr (bits == 8 && S == 8) {
     constexpr std::array<uint32_t, 8> shifts_ = {{0, 8, 16, 24, 0, 8, 16, 24}};
     auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
-    auto l = simd::Simd<uint32_t, 4>(*w++);
-    auto r = simd::Simd<uint32_t, 4>(*w);
+    auto l = simd::Simd<uint32_t, S / 2>(*w++);
+    auto r = simd::Simd<uint32_t, S / 2>(*w);
     wi = simd::Simd<uint32_t, S>(l, r);
     wi = wi >> shifts;
     wi = wi & bitmask;
@@ -316,6 +317,76 @@ void _qmm_dispatch_typed(
   }
 }
 
+template <typename T>
+void _qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int bits,
+    int group_size,
+    bool transposed_w,
+    Stream stream) {
+  int K = x.shape(-1);
+  int M = x.ndim() > 1 ? x.shape(-2) : 1;
+  int N = out.shape(-1);
+  int w_els = w.ndim() > 2 ? w.shape(-1) * w.shape(-2) : 0;
+  int g_els = w.ndim() > 2 ? scales.shape(-1) * scales.shape(-2) : 0;
+  int batch_size = x.size() / (K * M);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(biases);
+  encoder.set_output_array(out);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+
+  encoder.dispatch([out_ptr,
+                    x_ptr,
+                    w_ptr,
+                    scales_ptr,
+                    biases_ptr,
+                    x_shape = x.shape(),
+                    x_strides = x.strides(),
+                    w_shape = w.shape(),
+                    w_strides = w.strides(),
+                    scales_shape = scales.shape(),
+                    scales_strides = scales.strides(),
+                    biases_shape = biases.shape(),
+                    biases_strides = biases.strides(),
+                    w_els,
+                    g_els,
+                    batch_size,
+                    M,
+                    N,
+                    K,
+                    bits,
+                    group_size,
+                    transposed_w] {
+    for (int i = 0; i < batch_size; i++) {
+      _qmm_dispatch_typed<T>(
+          out_ptr + i * M * N,
+          x_ptr + elem_to_loc(i * M * K, x_shape, x_strides),
+          w_ptr + elem_to_loc(i * w_els, w_shape, w_strides),
+          scales_ptr + elem_to_loc(i * g_els, scales_shape, scales_strides),
+          biases_ptr + elem_to_loc(i * g_els, biases_shape, biases_strides),
+          M,
+          N,
+          K,
+          bits,
+          group_size,
+          transposed_w);
+    }
+  });
+}
+
 void _qmm_dispatch(
     array& out,
     const array& x,
@@ -324,64 +395,111 @@ void _qmm_dispatch(
     const array& biases,
     int bits,
     int group_size,
-    bool transposed_w) {
+    bool transposed_w,
+    Stream stream) {
+  switch (x.dtype()) {
+    case float32:
+      _qmm_dispatch_typed<float>(
+          out, x, w, scales, biases, bits, group_size, transposed_w, stream);
+      break;
+    case float16:
+      _qmm_dispatch_typed<float16_t>(
+          out, x, w, scales, biases, bits, group_size, transposed_w, stream);
+      break;
+    case bfloat16:
+      _qmm_dispatch_typed<bfloat16_t>(
+          out, x, w, scales, biases, bits, group_size, transposed_w, stream);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+
+template <typename T>
+void _bs_qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int bits,
+    int group_size,
+    bool transposed_w,
+    Stream stream) {
   int K = x.shape(-1);
-  int M = x.ndim() > 1 ? x.shape(-2) : 1;
+  int M = x.shape(-2);
   int N = out.shape(-1);
 
-  int w_els = w.ndim() > 2 ? w.shape(-1) * w.shape(-2) : 0;
-  int g_els = w.ndim() > 2 ? scales.shape(-1) * scales.shape(-2) : 0;
+  int w_els = w.shape(-1) * w.shape(-2);
+  int g_els = scales.shape(-1) * scales.shape(-2);
 
-  int batch_size = x.size() / (K * M);
-  for (int i = 0; i < batch_size; i++) {
-    switch (x.dtype()) {
-      case float32:
-        _qmm_dispatch_typed<float>(
-            out.data<float>() + i * M * N,
-            x.data<float>() + elem_to_loc(i * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(i * w_els, w),
-            scales.data<float>() + elem_to_loc(i * g_els, scales),
-            biases.data<float>() + elem_to_loc(i * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case float16:
-        _qmm_dispatch_typed<float16_t>(
-            out.data<float16_t>() + i * M * N,
-            x.data<float16_t>() + elem_to_loc(i * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(i * w_els, w),
-            scales.data<float16_t>() + elem_to_loc(i * g_els, scales),
-            biases.data<float16_t>() + elem_to_loc(i * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case bfloat16:
-        _qmm_dispatch_typed<bfloat16_t>(
-            out.data<bfloat16_t>() + i * M * N,
-            x.data<bfloat16_t>() + elem_to_loc(i * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(i * w_els, w),
-            scales.data<bfloat16_t>() + elem_to_loc(i * g_els, scales),
-            biases.data<bfloat16_t>() + elem_to_loc(i * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      default:
-        throw std::invalid_argument(
-            "[quantized_matmul] only floating types are supported");
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(biases);
+  encoder.set_input_array(lhs_indices);
+  encoder.set_input_array(rhs_indices);
+  encoder.set_output_array(out);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+  auto lhs_indices_ptr = lhs_indices.data<uint32_t>();
+  auto rhs_indices_ptr = rhs_indices.data<uint32_t>();
+
+  encoder.dispatch([out_ptr,
+                    x_ptr,
+                    w_ptr,
+                    scales_ptr,
+                    biases_ptr,
+                    lhs_indices_ptr,
+                    rhs_indices_ptr,
+                    x_shape = x.shape(),
+                    x_strides = x.strides(),
+                    w_shape = w.shape(),
+                    w_strides = w.strides(),
+                    scales_shape = scales.shape(),
+                    scales_strides = scales.strides(),
+                    biases_shape = biases.shape(),
+                    biases_strides = biases.strides(),
+                    lhs_indices_shape = lhs_indices.shape(),
+                    lhs_indices_strides = lhs_indices.strides(),
+                    rhs_indices_shape = rhs_indices.shape(),
+                    rhs_indices_strides = rhs_indices.strides(),
+                    w_els,
+                    g_els,
+                    indices_size = lhs_indices.size(),
+                    M,
+                    N,
+                    K,
+                    bits,
+                    group_size,
+                    transposed_w]() {
+    for (int i = 0; i < indices_size; i++) {
+      int x_idx = lhs_indices_ptr[elem_to_loc(
+          i, lhs_indices_shape, lhs_indices_strides)];
+      int w_idx = rhs_indices_ptr[elem_to_loc(
+          i, rhs_indices_shape, rhs_indices_strides)];
+      _qmm_dispatch_typed<T>(
+          out_ptr + i * M * N,
+          x_ptr + elem_to_loc(x_idx * M * K, x_shape, x_strides),
+          w_ptr + elem_to_loc(w_idx * w_els, w_shape, w_strides),
+          scales_ptr + elem_to_loc(w_idx * g_els, scales_shape, scales_strides),
+          biases_ptr + elem_to_loc(w_idx * g_els, biases_shape, biases_strides),
+          M,
+          N,
+          K,
+          bits,
+          group_size,
+          transposed_w);
     }
-  }
+  });
 }
 
 void _bs_qmm_dispatch(
@@ -394,68 +512,54 @@ void _bs_qmm_dispatch(
     const array& rhs_indices,
     int bits,
     int group_size,
-    bool transposed_w) {
-  int K = x.shape(-1);
-  int M = x.shape(-2);
-  int N = out.shape(-1);
-
-  int w_els = w.shape(-1) * w.shape(-2);
-  int g_els = scales.shape(-1) * scales.shape(-2);
-
-  const uint32_t* lhs_indices_data = lhs_indices.data<uint32_t>();
-  const uint32_t* rhs_indices_data = rhs_indices.data<uint32_t>();
-
-  for (int i = 0; i < lhs_indices.size(); i++) {
-    int x_idx = lhs_indices_data[elem_to_loc(i, lhs_indices)];
-    int w_idx = rhs_indices_data[elem_to_loc(i, rhs_indices)];
-
-    switch (x.dtype()) {
-      case float32:
-        _qmm_dispatch_typed<float>(
-            out.data<float>() + i * M * N,
-            x.data<float>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<float>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<float>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case float16:
-        _qmm_dispatch_typed<float16_t>(
-            out.data<float16_t>() + i * M * N,
-            x.data<float16_t>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<float16_t>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<float16_t>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case bfloat16:
-        _qmm_dispatch_typed<bfloat16_t>(
-            out.data<bfloat16_t>() + i * M * N,
-            x.data<bfloat16_t>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<bfloat16_t>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<bfloat16_t>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      default:
-        throw std::invalid_argument(
-            "[quantized_matmul] only floating types are supported");
-    }
+    bool transposed_w,
+    Stream stream) {
+  switch (x.dtype()) {
+    case float32:
+      _bs_qmm_dispatch_typed<float>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w,
+          stream);
+      break;
+    case float16:
+      _bs_qmm_dispatch_typed<float16_t>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w,
+          stream);
+      break;
+    case bfloat16:
+      _bs_qmm_dispatch_typed<bfloat16_t>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w,
+          stream);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
   }
 }
 
@@ -469,13 +573,14 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& scales_pre = inputs[2];
   auto& biases_pre = inputs[3];
 
-  auto ensure_row_contiguous = [](const array& arr) {
+  std::vector<array> temps;
+  auto ensure_row_contiguous = [s = stream(), &temps](const array& arr) {
     if (arr.flags().row_contiguous) {
       return arr;
     } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
-      return arr_copy;
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy(arr, temps.back(), CopyType::General, s);
+      return temps.back();
     }
   };
 
@@ -485,7 +590,10 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto biases = ensure_row_contiguous(biases_pre);
 
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  _qmm_dispatch(out, x, w, scales, biases, group_size_, bits_, transpose_);
+  _qmm_dispatch(
+      out, x, w, scales, biases, group_size_, bits_, transpose_, stream());
+  auto& enc = cpu::get_command_encoder(stream());
+  enc.add_temporaries(std::move(temps));
 }
 
 void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -498,15 +606,17 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& lhs_indices = inputs[4];
   auto& rhs_indices = inputs[5];
 
-  auto ensure_row_contiguous_last_dims = [](const array& arr) {
+  std::vector<array> temps;
+  auto ensure_row_contiguous_last_dims = [s = stream(),
+                                          &temps](const array& arr) {
     auto stride_0 = arr.strides()[arr.ndim() - 2];
     auto stride_1 = arr.strides()[arr.ndim() - 1];
     if (stride_0 == arr.shape(-1) && stride_1 == 1) {
       return arr;
     } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
-      return arr_copy;
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy(arr, temps.back(), CopyType::General, s);
+      return temps.back();
     }
   };
 
@@ -526,60 +636,57 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
       rhs_indices,
       group_size_,
       bits_,
-      transpose_);
+      transpose_,
+      stream());
+  auto& enc = cpu::get_command_encoder(stream());
+  enc.add_temporaries(std::move(temps));
 }
 
 template <typename T, typename U>
 void quantize(
-    const array& w_,
-    array& out_,
-    array& scales_,
-    array& biases_,
+    const T* w,
+    U* out,
+    T* scales,
+    T* biases,
     int bits,
-    int group_size) {
-  const T* w = w_.data<T>();
+    int group_size,
+    size_t w_size) {
+  float n_bins = (1 << bits) - 1;
+  float eps = 1e-7;
 
-  auto out = out_.data<U>();
-  T* scales = scales_.data<T>();
-  T* biases = biases_.data<T>();
-
-  T n_bins = (1 << bits) - 1;
-  T eps = 1e-7;
   bool power_of_2_bits = is_power_of_2(bits);
   int el_per_int = bits == 3 ? 8 : bits == 6 ? 4 : 32 / bits;
   // For 3/6 bits we read 3 uint8s at a time instead of 1 uint32
   int bytes_per_pack = power_of_2_bits ? 1 : 3;
   int int_per_group = group_size * bytes_per_pack / el_per_int;
-  size_t n_groups = w_.size() / group_size;
+  size_t n_groups = w_size / group_size;
 
   for (size_t i = 0; i < n_groups; ++i) {
     size_t w_idx = i * group_size;
-    T w_min = std::numeric_limits<float>::infinity();
-    T w_max = -w_min;
+    float w_min = std::numeric_limits<float>::infinity();
+    float w_max = -w_min;
     for (int j = 0; j < group_size; ++j) {
-      w_max = std::max(w_max, w[w_idx + j]);
-      w_min = std::min(w_min, w[w_idx + j]);
+      w_max = std::max(w_max, (float)w[w_idx + j]);
+      w_min = std::min(w_min, (float)w[w_idx + j]);
     }
     bool mask = std::abs(w_min) > std::abs(w_max);
-    T scale = std::max(T((w_max - w_min) / n_bins), eps);
+    float scale = std::max((w_max - w_min) / n_bins, eps);
     scale = mask ? scale : -scale;
 
-    auto edge = mask ? w_min : w_max;
-    auto q0 = std::rint(edge / scale);
-    if (q0 == 0) {
-      scales[i] = scale;
-      biases[i] = 0;
-    } else {
-      scales[i] = edge / q0;
-      biases[i] = edge;
+    float edge = mask ? w_min : w_max;
+    float q0 = std::rint(edge / scale);
+    float bias = 0;
+    if (q0 != 0) {
+      scale = edge / q0;
+      bias = edge;
     }
     size_t out_idx = i * int_per_group;
     for (int j = 0; j < int_per_group / bytes_per_pack; ++j) {
       uint32_t out_el = 0;
       for (int k = 0; k < el_per_int; ++k) {
-        T w_el = w[w_idx + j * el_per_int + k];
-        w_el = std::rint((w_el - biases[i]) / scales[i]);
-        w_el = std::min(std::max(w_el, T(0)), n_bins);
+        float w_el = w[w_idx + j * el_per_int + k];
+        w_el = std::rint((w_el - bias) / scale);
+        w_el = std::min(std::max(w_el, 0.0f), n_bins);
         out_el |= static_cast<uint32_t>(w_el) << (k * bits);
       }
       if (power_of_2_bits) {
@@ -590,23 +697,55 @@ void quantize(
         out[out_idx + bytes_per_pack * j + 2] = (out_el & 0xff0000) >> 16;
       }
     }
+    scales[i] = static_cast<T>(scale);
+    biases[i] = static_cast<T>(bias);
   }
 }
 
+template <typename T, typename U>
+void dispatch_quantize(
+    const array& w,
+    array& out,
+    array& scales,
+    array& biases,
+    int bits,
+    int group_size,
+    Stream stream) {
+  auto w_ptr = w.data<T>();
+  auto out_ptr = out.data<U>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(biases);
+  encoder.set_output_array(out);
+  encoder.dispatch([w_ptr,
+                    out_ptr,
+                    scales_ptr,
+                    biases_ptr,
+                    bits,
+                    group_size,
+                    w_size = w.size()]() {
+    quantize<T, U>(
+        w_ptr, out_ptr, scales_ptr, biases_ptr, bits, group_size, w_size);
+  });
+}
+
 void fast::AffineQuantize::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  auto ensure_row_contiguous = [](const array& arr) {
+  auto ensure_row_contiguous = [s = stream()](const array& arr) {
     if (arr.flags().row_contiguous) {
-      return arr;
+      return std::make_pair(arr, false);
     } else {
       array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
-      return arr_copy;
+      copy(arr, arr_copy, CopyType::General, s);
+      return std::make_pair(arr_copy, true);
     }
   };
-  auto w = ensure_row_contiguous(inputs[0]);
 
+  auto [w, copied] = ensure_row_contiguous(inputs[0]);
   auto& out = outputs[0];
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
@@ -616,27 +755,35 @@ void fast::AffineQuantize::eval_cpu(
   biases.set_data(allocator::malloc_or_wait(biases.nbytes()));
   if (w.dtype() == float16) {
     if (is_power_of_2(bits_)) {
-      quantize<float16_t, uint32_t>(w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<float16_t, uint32_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     } else {
-      quantize<float16_t, uint8_t>(w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<float16_t, uint8_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     }
   } else if (w.dtype() == bfloat16) {
     if (is_power_of_2(bits_)) {
-      quantize<bfloat16_t, uint32_t>(
-          w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<bfloat16_t, uint32_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     } else {
-      quantize<bfloat16_t, uint8_t>(w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<bfloat16_t, uint8_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     }
   } else if (w.dtype() == float32) {
     if (is_power_of_2(bits_)) {
-      quantize<float, uint32_t>(w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<float, uint32_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     } else {
-      quantize<float, uint8_t>(w, out, scales, biases, bits_, group_size_);
+      dispatch_quantize<float, uint8_t>(
+          w, out, scales, biases, bits_, group_size_, stream());
     }
   } else {
     throw std::runtime_error(
         "[fast::AffineQuantize::eval_cpu] Only supports floating point inputs");
   }
+  if (copied) {
+    cpu::get_command_encoder(stream()).add_temporary(w);
+  }
 }
 
 } // namespace mlx::core

mlx/backend/cpu/reduce.cpp

@@ -5,6 +5,7 @@
 #include <limits>
 
 #include "mlx/backend/common/reduce.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 #include "mlx/primitives.h"
 
@@ -140,25 +141,33 @@ void reduction_op(
     array& out,
     const std::vector<int>& axes,
     U init,
-    Op op) {
+    Stream stream) {
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
   ReductionPlan plan = get_reduction_plan(x, axes);
 
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(x);
+  encoder.set_output_array(out);
+
+  auto in_ptr = x.data<T>();
+  auto out_ptr = out.data<U>();
   if (plan.type == ContiguousAllReduce) {
-    U* out_ptr = out.data<U>();
-    *out_ptr = init;
-    contiguous_reduce(x.data<T>(), out_ptr, x.size(), op, init);
+    encoder.dispatch([in_ptr, out_ptr, init, size = x.size()]() {
+      *out_ptr = init;
+      contiguous_reduce(in_ptr, out_ptr, size, Op{}, init);
+    });
     return;
   }
 
   if (plan.type == ContiguousReduce && plan.shape.size() == 1) {
     int reduction_size = plan.shape[0];
-    const T* x_ptr = x.data<T>();
-    U* out_ptr = out.data<U>();
-    for (int i = 0; i < out.size(); i++, out_ptr++, x_ptr += reduction_size) {
-      *out_ptr = init;
-      contiguous_reduce(x_ptr, out_ptr, reduction_size, op, init);
-    }
+    encoder.dispatch(
+        [in_ptr, out_ptr, init, reduction_size, size = out.size()]() mutable {
+          for (int i = 0; i < size; i++, out_ptr++, in_ptr += reduction_size) {
+            *out_ptr = init;
+            contiguous_reduce(in_ptr, out_ptr, reduction_size, Op{}, init);
+          }
+        });
     return;
   }
 
@@ -166,34 +175,43 @@ void reduction_op(
     int reduction_size = plan.shape.back();
     plan.shape.pop_back();
     plan.strides.pop_back();
-    const T* x_ptr = x.data<T>();
-    U* out_ptr = out.data<U>();
     // Unrolling the following loop (and implementing it in order for
     // ContiguousReduce) should hold extra performance boost.
     auto [shape, strides] = shapes_without_reduction_axes(x, axes);
-    if (plan.shape.size() == 0) {
-      for (int i = 0; i < out.size(); i++, out_ptr++) {
-        int offset = elem_to_loc(i, shape, strides);
-        *out_ptr = init;
-        contiguous_reduce(x_ptr + offset, out_ptr, reduction_size, op, init);
+
+    encoder.dispatch([in_ptr,
+                      out_ptr,
+                      init,
+                      reduction_size,
+                      size = out.size(),
+                      plan = std::move(plan),
+                      shape = std::move(shape),
+                      strides = std::move(strides)]() mutable {
+      if (plan.shape.size() == 0) {
+        for (int i = 0; i < size; i++, out_ptr++) {
+          int offset = elem_to_loc(i, shape, strides);
+          *out_ptr = init;
+          contiguous_reduce(
+              in_ptr + offset, out_ptr, reduction_size, Op{}, init);
+        }
+      } else {
+        for (int i = 0; i < size; i++, out_ptr++) {
+          int offset = elem_to_loc(i, shape, strides);
+          *out_ptr = init;
+          nd_loop(
+              [&](int extra_offset) {
+                contiguous_reduce(
+                    in_ptr + offset + extra_offset,
+                    out_ptr,
+                    reduction_size,
+                    Op{},
+                    init);
+              },
+              plan.shape,
+              plan.strides);
+        }
       }
-    } else {
-      for (int i = 0; i < out.size(); i++, out_ptr++) {
-        int offset = elem_to_loc(i, shape, strides);
-        *out_ptr = init;
-        nd_loop(
-            [&](int extra_offset) {
-              contiguous_reduce(
-                  x_ptr + offset + extra_offset,
-                  out_ptr,
-                  reduction_size,
-                  op,
-                  init);
-            },
-            plan.shape,
-            plan.strides);
-      }
-    }
+    });
     return;
   }
 
@@ -202,14 +220,20 @@ void reduction_op(
     size_t reduction_stride = plan.strides.back();
     plan.shape.pop_back();
     plan.strides.pop_back();
-    const T* x_ptr = x.data<T>();
-    U* out_ptr = out.data<U>();
-    for (int i = 0; i < out.size(); i += reduction_stride) {
-      std::fill_n(out_ptr, reduction_stride, init);
-      strided_reduce(x_ptr, out_ptr, reduction_size, reduction_stride, op);
-      x_ptr += reduction_stride * reduction_size;
-      out_ptr += reduction_stride;
-    }
+
+    encoder.dispatch([in_ptr,
+                      out_ptr,
+                      init,
+                      reduction_size,
+                      reduction_stride,
+                      size = out.size()]() mutable {
+      for (int i = 0; i < size; i += reduction_stride) {
+        std::fill_n(out_ptr, reduction_stride, init);
+        strided_reduce(in_ptr, out_ptr, reduction_size, reduction_stride, Op{});
+        in_ptr += reduction_stride * reduction_size;
+        out_ptr += reduction_stride;
+      }
+    });
     return;
   }
 
@@ -219,53 +243,69 @@ void reduction_op(
     size_t reduction_stride = plan.strides.back();
     plan.shape.pop_back();
     plan.strides.pop_back();
-    const T* x_ptr = x.data<T>();
-    U* out_ptr = out.data<U>();
     auto [shape, strides] = shapes_without_reduction_axes(x, axes);
-    if (plan.shape.size() == 0) {
-      for (int i = 0; i < out.size(); i += reduction_stride) {
-        int offset = elem_to_loc(i, shape, strides);
-        std::fill_n(out_ptr, reduction_stride, init);
-        strided_reduce(
-            x_ptr + offset, out_ptr, reduction_size, reduction_stride, op);
-        out_ptr += reduction_stride;
+
+    encoder.dispatch([in_ptr,
+                      out_ptr,
+                      init,
+                      reduction_size,
+                      reduction_stride,
+                      size = out.size(),
+                      plan = std::move(plan),
+                      shape = std::move(shape),
+                      strides = std::move(strides)]() mutable {
+      if (plan.shape.size() == 0) {
+        for (int i = 0; i < size; i += reduction_stride) {
+          int offset = elem_to_loc(i, shape, strides);
+          std::fill_n(out_ptr, reduction_stride, init);
+          strided_reduce(
+              in_ptr + offset, out_ptr, reduction_size, reduction_stride, Op{});
+          out_ptr += reduction_stride;
+        }
+      } else {
+        for (int i = 0; i < size; i += reduction_stride) {
+          int offset = elem_to_loc(i, shape, strides);
+          std::fill_n(out_ptr, reduction_stride, init);
+          nd_loop(
+              [&](int extra_offset) {
+                strided_reduce(
+                    in_ptr + offset + extra_offset,
+                    out_ptr,
+                    reduction_size,
+                    reduction_stride,
+                    Op{});
+              },
+              plan.shape,
+              plan.strides);
+          out_ptr += reduction_stride;
+        }
       }
-    } else {
-      for (int i = 0; i < out.size(); i += reduction_stride) {
-        int offset = elem_to_loc(i, shape, strides);
-        std::fill_n(out_ptr, reduction_stride, init);
-        nd_loop(
-            [&](int extra_offset) {
-              strided_reduce(
-                  x_ptr + offset + extra_offset,
-                  out_ptr,
-                  reduction_size,
-                  reduction_stride,
-                  op);
-            },
-            plan.shape,
-            plan.strides);
-        out_ptr += reduction_stride;
-      }
-    }
+    });
     return;
   }
 
   if (plan.type == GeneralReduce) {
-    const T* x_ptr = x.data<T>();
-    U* out_ptr = out.data<U>();
     auto [shape, strides] = shapes_without_reduction_axes(x, axes);
-    for (int i = 0; i < out.size(); i++, out_ptr++) {
-      int offset = elem_to_loc(i, shape, strides);
-      U val = init;
-      nd_loop(
-          [&](int extra_offset) {
-            val = op(val, *(x_ptr + offset + extra_offset));
-          },
-          plan.shape,
-          plan.strides);
-      *out_ptr = val;
-    }
+
+    encoder.dispatch([in_ptr,
+                      out_ptr,
+                      init,
+                      size = out.size(),
+                      plan = std::move(plan),
+                      shape = std::move(shape),
+                      strides = std::move(strides)]() mutable {
+      for (int i = 0; i < size; i++, out_ptr++) {
+        int offset = elem_to_loc(i, shape, strides);
+        U val = init;
+        nd_loop(
+            [&](int extra_offset) {
+              val = Op{}(val, *(in_ptr + offset + extra_offset));
+            },
+            plan.shape,
+            plan.strides);
+        *out_ptr = val;
+      }
+    });
   }
 }
 
@@ -394,11 +434,12 @@ void reduce_dispatch_and_or(
     const array& in,
     array& out,
     Reduce::ReduceType rtype,
-    const std::vector<int>& axes) {
+    const std::vector<int>& axes,
+    Stream stream) {
   if (rtype == Reduce::And) {
-    reduction_op<InT, bool>(in, out, axes, true, AndReduce());
+    reduction_op<InT, bool, AndReduce>(in, out, axes, true, stream);
   } else {
-    reduction_op<InT, bool>(in, out, axes, false, OrReduce());
+    reduction_op<InT, bool, OrReduce>(in, out, axes, false, stream);
   }
 }
 
@@ -407,18 +448,19 @@ void reduce_dispatch_sum_prod(
     const array& in,
     array& out,
     Reduce::ReduceType rtype,
-    const std::vector<int>& axes) {
+    const std::vector<int>& axes,
+    Stream stream) {
   if (rtype == Reduce::Sum) {
     if constexpr (std::is_integral_v<InT> && sizeof(InT) <= 4) {
-      reduction_op<InT, int32_t>(in, out, axes, 0, SumReduce());
+      reduction_op<InT, int32_t, SumReduce>(in, out, axes, 0, stream);
     } else {
-      reduction_op<InT, InT>(in, out, axes, 0, SumReduce());
+      reduction_op<InT, InT, SumReduce>(in, out, axes, 0, stream);
     }
   } else {
     if constexpr (std::is_integral_v<InT> && sizeof(InT) <= 4) {
-      reduction_op<InT, int32_t>(in, out, axes, 1, ProdReduce());
+      reduction_op<InT, int32_t, ProdReduce>(in, out, axes, 1, stream);
     } else {
-      reduction_op<InT, InT>(in, out, axes, 1, ProdReduce());
+      reduction_op<InT, InT, ProdReduce>(in, out, axes, 1, stream);
     }
   }
 }
@@ -428,13 +470,14 @@ void reduce_dispatch_min_max(
     const array& in,
     array& out,
     Reduce::ReduceType rtype,
-    const std::vector<int>& axes) {
+    const std::vector<int>& axes,
+    Stream stream) {
   if (rtype == Reduce::Max) {
     auto init = Limits<InT>::min;
-    reduction_op<InT, InT>(in, out, axes, init, MaxReduce());
+    reduction_op<InT, InT, MaxReduce>(in, out, axes, init, stream);
   } else {
     auto init = Limits<InT>::max;
-    reduction_op<InT, InT>(in, out, axes, init, MinReduce());
+    reduction_op<InT, InT, MinReduce>(in, out, axes, init, stream);
   }
 }
 
@@ -448,24 +491,28 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
         case bool_:
         case uint8:
         case int8:
-          reduce_dispatch_and_or<int8_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_and_or<int8_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int16:
         case uint16:
         case float16:
         case bfloat16:
-          reduce_dispatch_and_or<int16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_and_or<int16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case uint32:
         case int32:
         case float32:
-          reduce_dispatch_and_or<int32_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_and_or<int32_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case uint64:
         case int64:
         case float64:
         case complex64:
-          reduce_dispatch_and_or<int64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_and_or<int64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
       }
       break;
@@ -476,34 +523,43 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
         case bool_:
         case uint8:
         case int8:
-          reduce_dispatch_sum_prod<int8_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<int8_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int16:
         case uint16:
-          reduce_dispatch_sum_prod<int16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<int16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int32:
         case uint32:
-          reduce_dispatch_sum_prod<int32_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<int32_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int64:
         case uint64:
-          reduce_dispatch_sum_prod<int64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<int64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float16:
-          reduce_dispatch_sum_prod<float16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<float16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case bfloat16:
-          reduce_dispatch_sum_prod<bfloat16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<bfloat16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float32:
-          reduce_dispatch_sum_prod<float>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<float>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float64:
-          reduce_dispatch_sum_prod<double>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<double>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case complex64:
-          reduce_dispatch_sum_prod<complex64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_sum_prod<complex64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
       }
       break;
@@ -512,46 +568,59 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
     case Reduce::Min: {
       switch (in.dtype()) {
         case bool_:
-          reduce_dispatch_min_max<bool>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<bool>(in, out, reduce_type_, axes_, stream());
           break;
         case uint8:
-          reduce_dispatch_min_max<uint8_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint8_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case uint16:
-          reduce_dispatch_min_max<uint16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case uint32:
-          reduce_dispatch_min_max<uint32_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint32_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case uint64:
-          reduce_dispatch_min_max<uint64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int8:
-          reduce_dispatch_min_max<uint8_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint8_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int16:
-          reduce_dispatch_min_max<uint16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<uint16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int32:
-          reduce_dispatch_min_max<int32_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<int32_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case int64:
-          reduce_dispatch_min_max<int64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<int64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float16:
-          reduce_dispatch_min_max<float16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<float16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float32:
-          reduce_dispatch_min_max<float>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<float>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case float64:
-          reduce_dispatch_min_max<double>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<double>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case bfloat16:
-          reduce_dispatch_min_max<bfloat16_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<bfloat16_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
         case complex64:
-          reduce_dispatch_min_max<complex64_t>(in, out, reduce_type_, axes_);
+          reduce_dispatch_min_max<complex64_t>(
+              in, out, reduce_type_, axes_, stream());
           break;
       }
       break;

mlx/backend/cpu/scan.cpp

@@ -4,6 +4,7 @@
 
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 #include "mlx/primitives.h"
 
@@ -153,37 +154,44 @@ void strided_scan(
 
 template <typename T, typename U, typename Op>
 void scan_op(
-    const array& input,
-    array& output,
+    const array& in,
+    array& out,
     int axis,
     bool reverse,
     bool inclusive,
     const Op& op,
-    U init) {
-  output.set_data(allocator::malloc_or_wait(output.nbytes()));
+    U init,
+    Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
 
-  if (input.flags().row_contiguous) {
-    if (input.strides()[axis] == 1) {
-      contiguous_scan(
-          input.data<T>(),
-          output.data<U>(),
-          input.size() / input.shape(axis),
-          input.shape(axis),
-          reverse,
-          inclusive,
-          op,
-          init);
+  if (in.flags().row_contiguous) {
+    if (in.strides()[axis] == 1) {
+      encoder.dispatch([in_ptr = in.data<T>(),
+                        out_ptr = out.data<U>(),
+                        count = in.size() / in.shape(axis),
+                        stride = in.shape(axis),
+                        reverse,
+                        inclusive,
+                        op = std::move(op),
+                        init]() {
+        contiguous_scan(
+            in_ptr, out_ptr, count, stride, reverse, inclusive, op, init);
+      });
     } else {
-      strided_scan(
-          input.data<T>(),
-          output.data<U>(),
-          input.size() / input.shape(axis) / input.strides()[axis],
-          input.shape(axis),
-          input.strides()[axis],
-          reverse,
-          inclusive,
-          op,
-          init);
+      encoder.dispatch([in_ptr = in.data<T>(),
+                        out_ptr = out.data<U>(),
+                        count = in.size() / in.shape(axis) / in.strides()[axis],
+                        size = in.shape(axis),
+                        stride = in.strides()[axis],
+                        reverse,
+                        inclusive,
+                        op = std::move(op),
+                        init]() {
+        strided_scan(
+            in_ptr, out_ptr, count, size, stride, reverse, inclusive, op, init);
+      });
     }
   } else {
     throw std::runtime_error("Scan op supports only contiguous inputs");
@@ -193,38 +201,39 @@ void scan_op(
 template <typename T, typename U>
 void scan_dispatch(
     Scan::ReduceType rtype,
-    const array& input,
-    array& output,
+    const array& in,
+    array& out,
     int axis,
     bool reverse,
-    bool inclusive) {
+    bool inclusive,
+    Stream stream) {
   switch (rtype) {
     case Scan::Sum: {
       auto op = [](U y, T x) { return y + x; };
       auto init = static_cast<U>(0);
-      scan_op<T, U>(input, output, axis, reverse, inclusive, op, init);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init, stream);
       break;
     }
     case Scan::Prod: {
       auto op = [](U y, T x) { return y * x; };
       auto init = static_cast<U>(1);
-      scan_op<T, U>(input, output, axis, reverse, inclusive, op, init);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init, stream);
       break;
     }
     case Scan::Min: {
       auto op = [](U y, T x) { return x < y ? x : y; };
-      auto init = (issubdtype(input.dtype(), floating))
+      auto init = (issubdtype(in.dtype(), floating))
           ? static_cast<U>(std::numeric_limits<float>::infinity())
           : std::numeric_limits<U>::max();
-      scan_op<T, U>(input, output, axis, reverse, inclusive, op, init);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init, stream);
       break;
     }
     case Scan::Max: {
       auto op = [](U y, T x) { return x < y ? y : x; };
-      auto init = (issubdtype(input.dtype(), floating))
+      auto init = (issubdtype(in.dtype(), floating))
           ? static_cast<U>(-std::numeric_limits<float>::infinity())
           : std::numeric_limits<U>::min();
-      scan_op<T, U>(input, output, axis, reverse, inclusive, op, init);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init, stream);
       break;
     }
   }
@@ -237,11 +246,14 @@ void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   // Ensure contiguity
   auto in = inputs[0];
+  bool copied = false;
   if (!in.flags().row_contiguous) {
     array arr_copy(in.shape(), in.dtype(), nullptr, {});
-    copy(in, arr_copy, CopyType::General);
+    copy(in, arr_copy, CopyType::General, stream());
     in = arr_copy;
+    copied = true;
   }
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   switch (in.dtype()) {
     case bool_: {
@@ -252,65 +264,68 @@ void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
       //       floats perhaps we should add the full all-to-all dispatch.
       if (reduce_type_ == Scan::Sum && out.dtype() == int32) {
         scan_dispatch<bool, int32_t>(
-            reduce_type_, in, out, axis_, reverse_, inclusive_);
+            reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       } else {
         scan_dispatch<bool, bool>(
-            reduce_type_, in, out, axis_, reverse_, inclusive_);
+            reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       }
       break;
     }
     case uint8:
       scan_dispatch<uint8_t, uint8_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case uint16:
       scan_dispatch<uint16_t, uint16_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case uint32:
       scan_dispatch<uint32_t, uint32_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case uint64:
       scan_dispatch<uint64_t, uint64_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case int8:
       scan_dispatch<int8_t, int8_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case int16:
       scan_dispatch<int16_t, int16_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case int32:
       scan_dispatch<int32_t, int32_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case int64:
       scan_dispatch<int64_t, int64_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case float16:
       scan_dispatch<float16_t, float16_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case float32:
       scan_dispatch<float, float>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case float64:
       scan_dispatch<double, double>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case bfloat16:
       scan_dispatch<bfloat16_t, bfloat16_t>(
-          reduce_type_, in, out, axis_, reverse_, inclusive_);
+          reduce_type_, in, out, axis_, reverse_, inclusive_, stream());
       break;
     case complex64:
       throw std::runtime_error("Scan ops do not support complex types yet");
       break;
   }
+  if (copied) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(in));
+  }
 }
 
 } // namespace mlx::core

mlx/backend/cpu/simd/math.h

@@ -186,7 +186,7 @@ Simd<T, N> erfinv(Simd<T, N> a_) {
       return a * rhs(t);
     }
   } else {
-    return a * select(t > thresh, lhs(t), rhs(t));
+    return a * select(abs(t) > thresh, lhs(t), rhs(t));
   }
 }
 

mlx/backend/cpu/softmax.cpp

@@ -4,6 +4,7 @@
 #include <cmath>
 
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 #include "mlx/primitives.h"
 #include "mlx/types/limits.h"
@@ -15,92 +16,100 @@ namespace {
 using namespace mlx::core::simd;
 
 template <typename T, typename AccT>
-void softmax(const array& in, array& out) {
-  constexpr bool same_t = std::is_same_v<T, AccT>;
-  constexpr int N = std::min(max_size<AccT>, max_size<T>);
+void softmax(const array& in, array& out, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
 
   const T* in_ptr = in.data<T>();
   T* out_ptr = out.data<T>();
+
   int M = in.shape().back();
   int L = in.data_size() / M;
-  const T* current_in_ptr;
-  T* current_out_ptr;
 
-  for (int i = 0; i < L; i++, in_ptr += M, out_ptr += M) {
-    // Find the maximum
-    current_in_ptr = in_ptr;
-    Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
-    size_t s = M;
-    while (s >= N) {
-      Simd<AccT, N> vals = load<T, N>(current_in_ptr);
-      vmaximum = maximum(vals, vmaximum);
-      current_in_ptr += N;
-      s -= N;
-    }
+  encoder.dispatch([in_ptr, out_ptr, M, L]() mutable {
+    constexpr bool same_t = std::is_same_v<T, AccT>;
+    constexpr int N = std::min(max_size<AccT>, max_size<T>);
 
-    AccT maximum = max(vmaximum);
-    while (s-- > 0) {
-      maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
-      current_in_ptr++;
-    }
+    const T* current_in_ptr;
+    T* current_out_ptr;
 
-    // Compute the normalizer and the exponentials
-    Simd<AccT, N> vnormalizer(0.0);
-    current_out_ptr = out_ptr;
-    current_in_ptr = in_ptr;
-    s = M;
-    while (s >= N) {
-      Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
-      vexp = exp(vexp - maximum);
-      if constexpr (same_t) {
-        store(current_out_ptr, vexp);
-      }
-      vnormalizer = vnormalizer + vexp;
-      current_in_ptr += N;
-      current_out_ptr += N;
-      s -= N;
-    }
-    AccT normalizer = sum(vnormalizer);
-    while (s-- > 0) {
-      AccT _exp = std::exp(*current_in_ptr - maximum);
-      if constexpr (same_t) {
-        *current_out_ptr = _exp;
-      }
-      normalizer += _exp;
-      current_in_ptr++;
-      current_out_ptr++;
-    }
-    normalizer = 1 / normalizer;
-
-    // Normalize
-    current_out_ptr = out_ptr;
-    current_in_ptr = in_ptr;
-    s = M;
-    while (s >= N) {
-      if constexpr (same_t) {
-        store(
-            current_out_ptr,
-            Simd<T, N>(load<T, N>(current_out_ptr) * normalizer));
-      } else {
-        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
-        vexp = exp(vexp - maximum) * normalizer;
-        store(current_out_ptr, Simd<T, N>(vexp));
+    for (int i = 0; i < L; i++, in_ptr += M, out_ptr += M) {
+      // Find the maximum
+      current_in_ptr = in_ptr;
+      Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
+      size_t s = M;
+      while (s >= N) {
+        Simd<AccT, N> vals = load<T, N>(current_in_ptr);
+        vmaximum = maximum(vals, vmaximum);
         current_in_ptr += N;
+        s -= N;
       }
-      current_out_ptr += N;
-      s -= N;
-    }
-    while (s-- > 0) {
-      if constexpr (same_t) {
-        *current_out_ptr *= normalizer;
-      } else {
-        AccT _exp = std::exp(*current_in_ptr - maximum);
-        *current_out_ptr = static_cast<T>(_exp * normalizer);
+
+      AccT maximum = max(vmaximum);
+      while (s-- > 0) {
+        maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
         current_in_ptr++;
       }
-      current_out_ptr++;
+
+      // Compute the normalizer and the exponentials
+      Simd<AccT, N> vnormalizer(0.0);
+      current_out_ptr = out_ptr;
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+        vexp = exp(vexp - maximum);
+        if constexpr (same_t) {
+          store(current_out_ptr, vexp);
+        }
+        vnormalizer = vnormalizer + vexp;
+        current_in_ptr += N;
+        current_out_ptr += N;
+        s -= N;
+      }
+      AccT normalizer = sum(vnormalizer);
+      while (s-- > 0) {
+        AccT _exp = std::exp(*current_in_ptr - maximum);
+        if constexpr (same_t) {
+          *current_out_ptr = _exp;
+        }
+        normalizer += _exp;
+        current_in_ptr++;
+        current_out_ptr++;
+      }
+      normalizer = 1 / normalizer;
+
+      // Normalize
+      current_out_ptr = out_ptr;
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        if constexpr (same_t) {
+          store(
+              current_out_ptr,
+              Simd<T, N>(load<T, N>(current_out_ptr) * normalizer));
+        } else {
+          Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+          vexp = exp(vexp - maximum) * normalizer;
+          store(current_out_ptr, Simd<T, N>(vexp));
+          current_in_ptr += N;
+        }
+        current_out_ptr += N;
+        s -= N;
+      }
+      while (s-- > 0) {
+        if constexpr (same_t) {
+          *current_out_ptr *= normalizer;
+        } else {
+          AccT _exp = std::exp(*current_in_ptr - maximum);
+          *current_out_ptr = static_cast<T>(_exp * normalizer);
+          current_in_ptr++;
+        }
+        current_out_ptr++;
+      }
     }
-  }
+  });
 }
 
 } // namespace
@@ -109,30 +118,32 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
 
   // Make sure that the last dimension is contiguous
-  auto check_input = [](array x) {
+  auto set_output = [s = stream(), &out](const array& x) {
     bool no_copy = x.strides()[x.ndim() - 1] == 1;
     if (x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
       no_copy &= (s == 0 || s == x.shape().back());
     }
     if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
       return x;
     } else {
       array x_copy(x.shape(), x.dtype(), nullptr, {});
-      copy(x, x_copy, CopyType::General);
+      copy(x, x_copy, CopyType::General, s);
+      out.copy_shared_buffer(x_copy);
       return x_copy;
     }
   };
-  array in = check_input(std::move(inputs[0]));
-  if (in.is_donatable()) {
-    out.copy_shared_buffer(in);
-  } else {
-    out.set_data(
-        allocator::malloc_or_wait(in.data_size() * in.itemsize()),
-        in.data_size(),
-        in.strides(),
-        in.flags());
-  }
+
+  auto in = set_output(inputs[0]);
 
   switch (in.dtype()) {
     case bool_:
@@ -148,24 +159,24 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
           "Softmax is defined only for floating point types");
       break;
     case float32:
-      softmax<float, float>(in, out);
+      softmax<float, float>(in, out, stream());
       break;
     case float16:
       if (precise_) {
-        softmax<float16_t, float>(in, out);
+        softmax<float16_t, float>(in, out, stream());
       } else {
-        softmax<float16_t, float16_t>(in, out);
+        softmax<float16_t, float16_t>(in, out, stream());
       }
       break;
     case bfloat16:
       if (precise_) {
-        softmax<bfloat16_t, float>(in, out);
+        softmax<bfloat16_t, float>(in, out, stream());
       } else {
-        softmax<bfloat16_t, bfloat16_t>(in, out);
+        softmax<bfloat16_t, bfloat16_t>(in, out, stream());
       }
       break;
     case float64:
-      softmax<double, double>(in, out);
+      softmax<double, double>(in, out, stream());
       break;
     case complex64:
       throw std::invalid_argument(

mlx/backend/cpu/sort.cpp

@@ -7,6 +7,7 @@
 
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 
 #include "mlx/primitives.h"
 
@@ -103,11 +104,11 @@ struct StridedIterator {
   T* ptr_;
 };
 
-template <typename T, typename IdxT = uint32_t>
-void sort(const array& in, array& out, int axis) {
+template <typename T>
+void sort(const array& in, array& out, int axis, Stream stream) {
   // Copy input to output
   CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General;
-  copy(in, out, ctype);
+  copy(in, out, ctype, stream);
 
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + in.ndim() : axis;
@@ -126,19 +127,27 @@ void sort(const array& in, array& out, int axis) {
   // Perform sorting in place
   ContiguousIterator src_it(
       remaining_shape, remaining_strides, remaining_shape.size());
-  for (int i = 0; i < n_rows; i++) {
-    T* data_ptr = out.data<T>() + src_it.loc;
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([out_ptr = out.data<T>(),
+                    src_it = std::move(src_it),
+                    n_rows,
+                    axis_size,
+                    axis_stride]() mutable {
+    for (int i = 0; i < n_rows; i++) {
+      T* data_ptr = out_ptr + src_it.loc;
 
-    StridedIterator st(data_ptr, axis_stride, 0);
-    StridedIterator ed(data_ptr, axis_stride, axis_size);
+      StridedIterator st(data_ptr, axis_stride, 0);
+      StridedIterator ed(data_ptr, axis_stride, axis_size);
 
-    std::stable_sort(st, ed);
-    src_it.step();
-  }
+      std::stable_sort(st, ed);
+      src_it.step();
+    }
+  });
 }
 
 template <typename T, typename IdxT = uint32_t>
-void argsort(const array& in, array& out, int axis) {
+void argsort(const array& in, array& out, int axis, Stream stream) {
   // Allocate output
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
@@ -167,35 +176,48 @@ void argsort(const array& in, array& out, int axis) {
       in_remaining_shape, in_remaining_strides, in_remaining_shape.size());
   ContiguousIterator out_it(
       out_remaining_shape, out_remaining_strides, out_remaining_shape.size());
-  for (int i = 0; i < n_rows; i++) {
-    const T* data_ptr = in.data<T>() + in_it.loc;
-    IdxT* idx_ptr = out.data<IdxT>() + out_it.loc;
-    in_it.step();
-    out_it.step();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_input_array(out);
+  encoder.dispatch([in_ptr = in.data<T>(),
+                    out_ptr = out.data<IdxT>(),
+                    in_it = std::move(in_it),
+                    out_it = std::move(out_it),
+                    n_rows,
+                    axis_size,
+                    in_stride,
+                    out_stride]() mutable {
+    for (int i = 0; i < n_rows; i++) {
+      const T* data_ptr = in_ptr + in_it.loc;
+      IdxT* idx_ptr = out_ptr + out_it.loc;
 
-    StridedIterator st_(idx_ptr, out_stride, 0);
-    StridedIterator ed_(idx_ptr, out_stride, axis_size);
+      in_it.step();
+      out_it.step();
 
-    // Initialize with iota
-    std::iota(st_, ed_, IdxT(0));
+      StridedIterator st_(idx_ptr, out_stride, 0);
+      StridedIterator ed_(idx_ptr, out_stride, axis_size);
 
-    // Sort according to vals
-    StridedIterator st(idx_ptr, out_stride, 0);
-    StridedIterator ed(idx_ptr, out_stride, axis_size);
+      // Initialize with iota
+      std::iota(st_, ed_, IdxT(0));
 
-    std::stable_sort(st, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
-      auto v1 = data_ptr[a * in_stride];
-      auto v2 = data_ptr[b * in_stride];
-      return v1 < v2 || (v1 == v2 && a < b);
-    });
-  }
+      // Sort according to vals
+      StridedIterator st(idx_ptr, out_stride, 0);
+      StridedIterator ed(idx_ptr, out_stride, axis_size);
+
+      std::stable_sort(st, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
+        auto v1 = data_ptr[a * in_stride];
+        auto v2 = data_ptr[b * in_stride];
+        return v1 < v2 || (v1 == v2 && a < b);
+      });
+    }
+  });
 }
 
-template <typename T, typename IdxT = uint32_t>
-void partition(const array& in, array& out, int axis, int kth) {
+template <typename T>
+void partition(const array& in, array& out, int axis, int kth, Stream stream) {
   // Copy input to output
   CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General;
-  copy(in, out, ctype);
+  copy(in, out, ctype, stream);
 
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + in.ndim() : axis;
@@ -216,20 +238,34 @@ void partition(const array& in, array& out, int axis, int kth) {
   // Perform partition in place
   ContiguousIterator src_it(
       remaining_shape, remaining_strides, remaining_shape.size());
-  for (int i = 0; i < n_rows; i++) {
-    T* data_ptr = out.data<T>() + src_it.loc;
-    src_it.step();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([out_ptr = out.data<T>(),
+                    src_it = std::move(src_it),
+                    n_rows,
+                    axis_size,
+                    axis_stride,
+                    kth]() mutable {
+    for (int i = 0; i < n_rows; i++) {
+      T* data_ptr = out_ptr + src_it.loc;
+      src_it.step();
 
-    StridedIterator st(data_ptr, axis_stride, 0);
-    StridedIterator md(data_ptr, axis_stride, kth);
-    StridedIterator ed(data_ptr, axis_stride, axis_size);
+      StridedIterator st(data_ptr, axis_stride, 0);
+      StridedIterator md(data_ptr, axis_stride, kth);
+      StridedIterator ed(data_ptr, axis_stride, axis_size);
 
-    std::nth_element(st, md, ed);
-  }
+      std::nth_element(st, md, ed);
+    }
+  });
 }
 
 template <typename T, typename IdxT = uint32_t>
-void argpartition(const array& in, array& out, int axis, int kth) {
+void argpartition(
+    const array& in,
+    array& out,
+    int axis,
+    int kth,
+    Stream stream) {
   // Allocate output
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
@@ -260,29 +296,43 @@ void argpartition(const array& in, array& out, int axis, int kth) {
       in_remaining_shape, in_remaining_strides, in_remaining_shape.size());
   ContiguousIterator out_it(
       out_remaining_shape, out_remaining_strides, out_remaining_shape.size());
-  for (int i = 0; i < n_rows; i++) {
-    const T* data_ptr = in.data<T>() + in_it.loc;
-    IdxT* idx_ptr = out.data<IdxT>() + out_it.loc;
-    in_it.step();
-    out_it.step();
 
-    StridedIterator st_(idx_ptr, out_stride, 0);
-    StridedIterator ed_(idx_ptr, out_stride, axis_size);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_input_array(out);
+  encoder.dispatch([in_ptr = in.data<T>(),
+                    out_ptr = out.data<IdxT>(),
+                    in_it = std::move(in_it),
+                    out_it = std::move(out_it),
+                    n_rows,
+                    axis_size,
+                    in_stride,
+                    out_stride,
+                    kth]() mutable {
+    for (int i = 0; i < n_rows; i++) {
+      const T* data_ptr = in_ptr + in_it.loc;
+      IdxT* idx_ptr = out_ptr + out_it.loc;
+      in_it.step();
+      out_it.step();
 
-    // Initialize with iota
-    std::iota(st_, ed_, IdxT(0));
+      StridedIterator st_(idx_ptr, out_stride, 0);
+      StridedIterator ed_(idx_ptr, out_stride, axis_size);
 
-    // Sort according to vals
-    StridedIterator st(idx_ptr, out_stride, 0);
-    StridedIterator md(idx_ptr, out_stride, kth);
-    StridedIterator ed(idx_ptr, out_stride, axis_size);
+      // Initialize with iota
+      std::iota(st_, ed_, IdxT(0));
 
-    std::nth_element(st, md, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
-      auto v1 = data_ptr[a * in_stride];
-      auto v2 = data_ptr[b * in_stride];
-      return v1 < v2 || (v1 == v2 && a < b);
-    });
-  }
+      // Sort according to vals
+      StridedIterator st(idx_ptr, out_stride, 0);
+      StridedIterator md(idx_ptr, out_stride, kth);
+      StridedIterator ed(idx_ptr, out_stride, axis_size);
+
+      std::nth_element(st, md, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
+        auto v1 = data_ptr[a * in_stride];
+        auto v2 = data_ptr[b * in_stride];
+        return v1 < v2 || (v1 == v2 && a < b);
+      });
+    }
+  });
 }
 
 } // namespace
@@ -293,33 +343,33 @@ void ArgSort::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   switch (in.dtype()) {
     case bool_:
-      return argsort<bool>(in, out, axis_);
+      return argsort<bool>(in, out, axis_, stream());
     case uint8:
-      return argsort<uint8_t>(in, out, axis_);
+      return argsort<uint8_t>(in, out, axis_, stream());
     case uint16:
-      return argsort<uint16_t>(in, out, axis_);
+      return argsort<uint16_t>(in, out, axis_, stream());
     case uint32:
-      return argsort<uint32_t>(in, out, axis_);
+      return argsort<uint32_t>(in, out, axis_, stream());
     case uint64:
-      return argsort<uint64_t>(in, out, axis_);
+      return argsort<uint64_t>(in, out, axis_, stream());
     case int8:
-      return argsort<int8_t>(in, out, axis_);
+      return argsort<int8_t>(in, out, axis_, stream());
     case int16:
-      return argsort<int16_t>(in, out, axis_);
+      return argsort<int16_t>(in, out, axis_, stream());
     case int32:
-      return argsort<int32_t>(in, out, axis_);
+      return argsort<int32_t>(in, out, axis_, stream());
     case int64:
-      return argsort<int64_t>(in, out, axis_);
+      return argsort<int64_t>(in, out, axis_, stream());
     case float32:
-      return argsort<float>(in, out, axis_);
+      return argsort<float>(in, out, axis_, stream());
     case float64:
-      return argsort<double>(in, out, axis_);
+      return argsort<double>(in, out, axis_, stream());
     case float16:
-      return argsort<float16_t>(in, out, axis_);
+      return argsort<float16_t>(in, out, axis_, stream());
     case bfloat16:
-      return argsort<bfloat16_t>(in, out, axis_);
+      return argsort<bfloat16_t>(in, out, axis_, stream());
     case complex64:
-      return argsort<complex64_t>(in, out, axis_);
+      return argsort<complex64_t>(in, out, axis_, stream());
   }
 }
 
@@ -329,33 +379,33 @@ void Sort::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   switch (in.dtype()) {
     case bool_:
-      return sort<bool>(in, out, axis_);
+      return sort<bool>(in, out, axis_, stream());
     case uint8:
-      return sort<uint8_t>(in, out, axis_);
+      return sort<uint8_t>(in, out, axis_, stream());
     case uint16:
-      return sort<uint16_t>(in, out, axis_);
+      return sort<uint16_t>(in, out, axis_, stream());
     case uint32:
-      return sort<uint32_t>(in, out, axis_);
+      return sort<uint32_t>(in, out, axis_, stream());
     case uint64:
-      return sort<uint64_t>(in, out, axis_);
+      return sort<uint64_t>(in, out, axis_, stream());
     case int8:
-      return sort<int8_t>(in, out, axis_);
+      return sort<int8_t>(in, out, axis_, stream());
     case int16:
-      return sort<int16_t>(in, out, axis_);
+      return sort<int16_t>(in, out, axis_, stream());
     case int32:
-      return sort<int32_t>(in, out, axis_);
+      return sort<int32_t>(in, out, axis_, stream());
     case int64:
-      return sort<int64_t>(in, out, axis_);
+      return sort<int64_t>(in, out, axis_, stream());
     case float32:
-      return sort<float>(in, out, axis_);
+      return sort<float>(in, out, axis_, stream());
     case float64:
-      return sort<double>(in, out, axis_);
+      return sort<double>(in, out, axis_, stream());
     case float16:
-      return sort<float16_t>(in, out, axis_);
+      return sort<float16_t>(in, out, axis_, stream());
     case bfloat16:
-      return sort<bfloat16_t>(in, out, axis_);
+      return sort<bfloat16_t>(in, out, axis_, stream());
     case complex64:
-      return sort<complex64_t>(in, out, axis_);
+      return sort<complex64_t>(in, out, axis_, stream());
   }
 }
 
@@ -365,33 +415,33 @@ void ArgPartition::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   switch (in.dtype()) {
     case bool_:
-      return argpartition<bool>(in, out, axis_, kth_);
+      return argpartition<bool>(in, out, axis_, kth_, stream());
     case uint8:
-      return argpartition<uint8_t>(in, out, axis_, kth_);
+      return argpartition<uint8_t>(in, out, axis_, kth_, stream());
     case uint16:
-      return argpartition<uint16_t>(in, out, axis_, kth_);
+      return argpartition<uint16_t>(in, out, axis_, kth_, stream());
     case uint32:
-      return argpartition<uint32_t>(in, out, axis_, kth_);
+      return argpartition<uint32_t>(in, out, axis_, kth_, stream());
     case uint64:
-      return argpartition<uint64_t>(in, out, axis_, kth_);
+      return argpartition<uint64_t>(in, out, axis_, kth_, stream());
     case int8:
-      return argpartition<int8_t>(in, out, axis_, kth_);
+      return argpartition<int8_t>(in, out, axis_, kth_, stream());
     case int16:
-      return argpartition<int16_t>(in, out, axis_, kth_);
+      return argpartition<int16_t>(in, out, axis_, kth_, stream());
     case int32:
-      return argpartition<int32_t>(in, out, axis_, kth_);
+      return argpartition<int32_t>(in, out, axis_, kth_, stream());
     case int64:
-      return argpartition<int64_t>(in, out, axis_, kth_);
+      return argpartition<int64_t>(in, out, axis_, kth_, stream());
     case float32:
-      return argpartition<float>(in, out, axis_, kth_);
+      return argpartition<float>(in, out, axis_, kth_, stream());
     case float64:
-      return argpartition<double>(in, out, axis_, kth_);
+      return argpartition<double>(in, out, axis_, kth_, stream());
     case float16:
-      return argpartition<float16_t>(in, out, axis_, kth_);
+      return argpartition<float16_t>(in, out, axis_, kth_, stream());
     case bfloat16:
-      return argpartition<bfloat16_t>(in, out, axis_, kth_);
+      return argpartition<bfloat16_t>(in, out, axis_, kth_, stream());
     case complex64:
-      return argpartition<complex64_t>(in, out, axis_, kth_);
+      return argpartition<complex64_t>(in, out, axis_, kth_, stream());
   }
 }
 
@@ -401,33 +451,33 @@ void Partition::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   switch (in.dtype()) {
     case bool_:
-      return partition<bool>(in, out, axis_, kth_);
+      return partition<bool>(in, out, axis_, kth_, stream());
     case uint8:
-      return partition<uint8_t>(in, out, axis_, kth_);
+      return partition<uint8_t>(in, out, axis_, kth_, stream());
     case uint16:
-      return partition<uint16_t>(in, out, axis_, kth_);
+      return partition<uint16_t>(in, out, axis_, kth_, stream());
     case uint32:
-      return partition<uint32_t>(in, out, axis_, kth_);
+      return partition<uint32_t>(in, out, axis_, kth_, stream());
     case uint64:
-      return partition<uint64_t>(in, out, axis_, kth_);
+      return partition<uint64_t>(in, out, axis_, kth_, stream());
     case int8:
-      return partition<int8_t>(in, out, axis_, kth_);
+      return partition<int8_t>(in, out, axis_, kth_, stream());
     case int16:
-      return partition<int16_t>(in, out, axis_, kth_);
+      return partition<int16_t>(in, out, axis_, kth_, stream());
     case int32:
-      return partition<int32_t>(in, out, axis_, kth_);
+      return partition<int32_t>(in, out, axis_, kth_, stream());
     case int64:
-      return partition<int64_t>(in, out, axis_, kth_);
+      return partition<int64_t>(in, out, axis_, kth_, stream());
     case float32:
-      return partition<float>(in, out, axis_, kth_);
+      return partition<float>(in, out, axis_, kth_, stream());
     case float64:
-      return partition<double>(in, out, axis_, kth_);
+      return partition<double>(in, out, axis_, kth_, stream());
     case float16:
-      return partition<float16_t>(in, out, axis_, kth_);
+      return partition<float16_t>(in, out, axis_, kth_, stream());
     case bfloat16:
-      return partition<bfloat16_t>(in, out, axis_, kth_);
+      return partition<bfloat16_t>(in, out, axis_, kth_, stream());
     case complex64:
-      return partition<complex64_t>(in, out, axis_, kth_);
+      return partition<complex64_t>(in, out, axis_, kth_, stream());
   }
 }
 

mlx/backend/cpu/svd.cpp

@@ -2,12 +2,18 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
-void svd_impl(const array& a, array& u, array& s, array& vt) {
+template <typename T>
+void svd_impl(
+    const array& a,
+    std::vector<array>& outputs,
+    bool compute_uv,
+    Stream stream) {
   // Lapack uses the column-major convention. To avoid having to transpose
   // the input and then transpose the outputs, we swap the indices/sizes of the
   // matrices and take advantage of the following identity (see
@@ -21,129 +27,179 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
   const int N = a.shape(-1);
   const int K = std::min(M, N);
 
-  // A of shape M x N. The leading dimension is N since lapack receives Aᵀ.
-  const int lda = N;
-  // U of shape M x M. (N x N in lapack).
-  const int ldu = N;
-  // Vᵀ of shape N x N. (M x M in lapack).
-  const int ldvt = M;
-
   size_t num_matrices = a.size() / (M * N);
 
   // lapack clobbers the input, so we have to make a copy.
-  array in(a.shape(), float32, nullptr, {});
-  copy(a, in, a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  array in(a.shape(), a.dtype(), nullptr, {});
+  copy(
+      a,
+      in,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
 
   // Allocate outputs.
-  u.set_data(allocator::malloc_or_wait(u.nbytes()));
-  s.set_data(allocator::malloc_or_wait(s.nbytes()));
-  vt.set_data(allocator::malloc_or_wait(vt.nbytes()));
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  auto in_ptr = in.data<T>();
+  T* u_ptr;
+  T* s_ptr;
+  T* vt_ptr;
 
-  static constexpr auto job_u = "V";
-  static constexpr auto job_vt = "V";
-  static constexpr auto range = "A";
+  if (compute_uv) {
+    array& u = outputs[0];
+    array& s = outputs[1];
+    array& vt = outputs[2];
 
-  // Will contain the number of singular values after the call has returned.
-  int ns = 0;
-  float workspace_dimension = 0;
+    u.set_data(allocator::malloc_or_wait(u.nbytes()));
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
+    vt.set_data(allocator::malloc_or_wait(vt.nbytes()));
 
-  // Will contain the indices of eigenvectors that failed to converge (not used
-  // here but required by lapack).
-  auto iwork = array::Data{allocator::malloc_or_wait(sizeof(int) * 12 * K)};
+    encoder.set_output_array(u);
+    encoder.set_output_array(s);
+    encoder.set_output_array(vt);
 
-  static const int lwork_query = -1;
+    s_ptr = s.data<T>();
+    u_ptr = u.data<T>();
+    vt_ptr = vt.data<T>();
+  } else {
+    array& s = outputs[0];
 
-  static const int ignored_int = 0;
-  static const float ignored_float = 0;
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
 
-  int info;
+    encoder.set_output_array(s);
 
-  // Compute workspace size.
-  MLX_LAPACK_FUNC(sgesvdx)
-  (
-      /* jobu = */ job_u,
-      /* jobvt = */ job_vt,
-      /* range = */ range,
-      // M and N are swapped since lapack expects column-major.
-      /* m = */ &N,
-      /* n = */ &M,
-      /* a = */ nullptr,
-      /* lda = */ &lda,
-      /* vl = */ &ignored_float,
-      /* vu = */ &ignored_float,
-      /* il = */ &ignored_int,
-      /* iu = */ &ignored_int,
-      /* ns = */ &ns,
-      /* s = */ nullptr,
-      /* u = */ nullptr,
-      /* ldu = */ &ldu,
-      /* vt = */ nullptr,
-      /* ldvt = */ &ldvt,
-      /* work = */ &workspace_dimension,
-      /* lwork = */ &lwork_query,
-      /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
-      /* info = */ &info);
-
-  if (info != 0) {
-    std::stringstream ss;
-    ss << "svd_impl: sgesvdx_ workspace calculation failed with code " << info;
-    throw std::runtime_error(ss.str());
+    s_ptr = s.data<T>();
+    u_ptr = nullptr;
+    vt_ptr = nullptr;
   }
 
-  const int lwork = workspace_dimension;
-  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+  encoder.dispatch([in_ptr, u_ptr, s_ptr, vt_ptr, M, N, K, num_matrices]() {
+    // A of shape M x N. The leading dimension is N since lapack receives Aᵀ.
+    const int lda = N;
+    // U of shape M x M. (N x N in lapack).
+    const int ldu = N;
+    // Vᵀ of shape N x N. (M x M in lapack).
+    const int ldvt = M;
 
-  // Loop over matrices.
-  for (int i = 0; i < num_matrices; i++) {
-    MLX_LAPACK_FUNC(sgesvdx)
-    (
+    auto job_u = (u_ptr) ? "V" : "N";
+    auto job_vt = (u_ptr) ? "V" : "N";
+    static constexpr auto range = "A";
+
+    // Will contain the number of singular values after the call has returned.
+    int ns = 0;
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    auto iwork = array::Data{allocator::malloc_or_wait(sizeof(int) * 12 * K)};
+
+    static const int lwork_query = -1;
+
+    static const int ignored_int = 0;
+    static const T ignored_float = 0;
+
+    int info;
+
+    // Compute workspace size.
+    gesvdx<T>(
         /* jobu = */ job_u,
         /* jobvt = */ job_vt,
         /* range = */ range,
         // M and N are swapped since lapack expects column-major.
         /* m = */ &N,
         /* n = */ &M,
-        /* a = */ in.data<float>() + M * N * i,
+        /* a = */ nullptr,
         /* lda = */ &lda,
         /* vl = */ &ignored_float,
         /* vu = */ &ignored_float,
         /* il = */ &ignored_int,
         /* iu = */ &ignored_int,
         /* ns = */ &ns,
-        /* s = */ s.data<float>() + K * i,
-        // According to the identity above, lapack will write Vᵀᵀ as U.
-        /* u = */ vt.data<float>() + N * N * i,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
         /* ldu = */ &ldu,
-        // According to the identity above, lapack will write Uᵀ as Vᵀ.
-        /* vt = */ u.data<float>() + M * M * i,
+        /* vt = */ nullptr,
         /* ldvt = */ &ldvt,
-        /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
-        /* lwork = */ &lwork,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
         /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
         /* info = */ &info);
 
     if (info != 0) {
       std::stringstream ss;
-      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
       throw std::runtime_error(ss.str());
     }
 
-    if (ns != K) {
-      std::stringstream ss;
-      ss << "svd_impl: expected " << K << " singular values, but " << ns
-         << " were computed.";
-      throw std::runtime_error(ss.str());
+    const int lwork = workspace_dimension;
+    auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
+
+    // Loop over matrices.
+    for (int i = 0; i < num_matrices; i++) {
+      gesvdx<T>(
+          /* jobu = */ job_u,
+          /* jobvt = */ job_vt,
+          /* range = */ range,
+          // M and N are swapped since lapack expects column-major.
+          /* m = */ &N,
+          /* n = */ &M,
+          /* a = */ in_ptr + M * N * i,
+          /* lda = */ &lda,
+          /* vl = */ &ignored_float,
+          /* vu = */ &ignored_float,
+          /* il = */ &ignored_int,
+          /* iu = */ &ignored_int,
+          /* ns = */ &ns,
+          /* s = */ s_ptr + K * i,
+          // According to the identity above, lapack will write Vᵀᵀ as U.
+          /* u = */ vt_ptr ? vt_ptr + N * N * i : nullptr,
+          /* ldu = */ &ldu,
+          // According to the identity above, lapack will write Uᵀ as Vᵀ.
+          /* vt = */ u_ptr ? u_ptr + M * M * i : nullptr,
+          /* ldvt = */ &ldvt,
+          /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
+          /* lwork = */ &lwork,
+          /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
+          /* info = */ &info);
+
+      if (info != 0) {
+        std::stringstream ss;
+        ss << "svd_impl: sgesvdx_ failed with code " << info;
+        throw std::runtime_error(ss.str());
+      }
+
+      if (ns != K) {
+        std::stringstream ss;
+        ss << "svd_impl: expected " << K << " singular values, but " << ns
+           << " were computed.";
+        throw std::runtime_error(ss.str());
+      }
     }
-  }
+  });
+  encoder.add_temporary(in);
 }
 
+template <typename T>
+void compute_svd(
+    const array& a,
+    bool compute_uv,
+    std::vector<array>& outputs,
+    Stream stream) {}
+
 void SVD::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  if (!(inputs[0].dtype() == float32)) {
-    throw std::runtime_error("[SVD::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      svd_impl<float>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    case float64:
+      svd_impl<double>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[SVD::eval_cpu] only supports float32 or float64.");
   }
-  svd_impl(inputs[0], outputs[0], outputs[1], outputs[2]);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/ternary.h

@@ -5,6 +5,8 @@
 #include "mlx/array.h"
 #include "mlx/backend/common/ternary.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
 
 namespace mlx::core {
 
@@ -53,22 +55,18 @@ void ternary_op_dims(
 
 template <typename T1, typename T2, typename T3, typename U, typename Op>
 void ternary_op_dispatch_dims(
-    const array& a,
-    const array& b,
-    const array& c,
-    array& out,
-    Op op) {
-  auto [shape, strides] = collapse_contiguous_dims(
-      a.shape(), {a.strides(), b.strides(), c.strides(), out.strides()});
+    const T1* a_ptr,
+    const T2* b_ptr,
+    const T3* c_ptr,
+    U* out_ptr,
+    Op op,
+    size_t size,
+    Shape& shape,
+    std::vector<Strides>& strides) {
   const auto& a_strides = strides[0];
   const auto& b_strides = strides[1];
   const auto& c_strides = strides[2];
   const auto& out_strides = strides[3];
-
-  const T1* a_ptr = a.data<T1>();
-  const T2* b_ptr = b.data<T2>();
-  const T3* c_ptr = c.data<T3>();
-  U* out_ptr = out.data<T3>();
   int ndim = shape.size();
   switch (ndim) {
     case 1:
@@ -105,7 +103,7 @@ void ternary_op_dispatch_dims(
   ContiguousIterator b_it(shape, b_strides, ndim - 2);
   ContiguousIterator c_it(shape, c_strides, ndim - 2);
   auto stride = out_strides[ndim - 3];
-  for (size_t elem = 0; elem < a.size(); elem += stride) {
+  for (size_t elem = 0; elem < size; elem += stride) {
     ternary_op_dims<T1, T2, T3, U, Op, 2>(
         a_ptr + a_it.loc,
         b_ptr + b_it.loc,
@@ -134,23 +132,53 @@ void ternary_op(
   TernaryOpType topt = get_ternary_op_type(a, b, c);
   set_ternary_op_output_data(a, b, c, out, topt);
 
-  // The full computation is scalar-scalar-scalar so we call the base op once.
+  auto& encoder = cpu::get_command_encoder(out.primitive().stream());
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+  U* out_ptr = out.data<U>();
+
   if (topt == TernaryOpType::ScalarScalarScalar) {
-    *(out.data<U>()) = op(*a.data<T1>(), *b.data<T2>(), *c.data<T3>());
+    encoder.dispatch(
+        [a_ptr, b_ptr, c_ptr, out_ptr, op = std::move(op)]() mutable {
+          *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+        });
   } else if (topt == TernaryOpType::VectorVectorVector) {
-    const T1* a_ptr = a.data<T1>();
-    const T2* b_ptr = b.data<T2>();
-    const T3* c_ptr = c.data<T3>();
-    U* out_ptr = out.data<U>();
-    for (size_t i = 0; i < out.size(); ++i) {
-      *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
-      a_ptr++;
-      b_ptr++;
-      c_ptr++;
-      out_ptr++;
-    }
+    encoder.dispatch([a_ptr,
+                      b_ptr,
+                      c_ptr,
+                      out_ptr,
+                      op = std::move(op),
+                      size = out.size()]() mutable {
+      for (size_t i = 0; i < size; ++i) {
+        *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+        a_ptr++;
+        b_ptr++;
+        c_ptr++;
+        out_ptr++;
+      }
+    });
   } else {
-    ternary_op_dispatch_dims<T1, T2, T3, U>(a, b, c, out, op);
+    auto [shape, strides] = collapse_contiguous_dims(
+        a.shape(), {a.strides(), b.strides(), c.strides(), out.strides()});
+    encoder.dispatch(
+
+        [a_ptr,
+         b_ptr,
+         c_ptr,
+         out_ptr,
+         op = std::move(op),
+         size = out.size(),
+         shape = std::move(shape),
+         strides = std::move(strides)]() mutable {
+          ternary_op_dispatch_dims<T1, T2, T3, U>(
+              a_ptr, b_ptr, c_ptr, out_ptr, op, size, shape, strides);
+        });
   }
 }
 

mlx/backend/cpu/unary.h

@@ -5,67 +5,83 @@
 #include "mlx/allocator.h"
 #include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {
 
 void set_unary_output_data(const array& in, array& out) {
-  if (is_donatable(in, out)) {
-    out.copy_shared_buffer(in);
+  if (in.flags().contiguous) {
+    if (is_donatable(in, out)) {
+      out.copy_shared_buffer(in);
+    } else {
+      auto size = in.data_size();
+      out.set_data(
+          allocator::malloc_or_wait(size * out.itemsize()),
+          size,
+          in.strides(),
+          in.flags());
+    }
   } else {
-    auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
 }
 
 template <typename T, typename U = T, typename Op>
-void unary_op(const T* a, U* out, Op op, size_t shape, size_t stride) {
+void unary_op(const T* a, U* out, size_t shape, size_t stride) {
   for (size_t i = 0; i < shape; i += 1) {
-    out[i] = op(*a);
+    out[i] = Op{}(*a);
     a += stride;
   }
 }
 
 template <typename T, typename U = T, typename Op>
-void unary_op(const array& a, array& out, Op op) {
-  const T* a_ptr = a.data<T>();
-  if (a.flags().contiguous) {
-    set_unary_output_data(a, out);
-    U* dst = out.data<U>();
-    constexpr int N = simd::max_size<T>;
-    size_t size = a.data_size();
-    while (size >= N) {
-      simd::store(dst, op(simd::load<T, N>(a_ptr)));
-      size -= N;
-      a_ptr += N;
-      dst += N;
+void unary_op(const array& a, array& out, Op) {
+  set_unary_output_data(a, out);
+  const T* src = a.data<T>();
+  U* dst = out.data<U>();
+  auto& encoder = cpu::get_command_encoder(out.primitive().stream());
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+
+  encoder.dispatch([src,
+                    dst,
+                    contig = a.flags().contiguous,
+                    data_size = a.data_size(),
+                    size = a.size(),
+                    shapes = a.shape(),
+                    strides = a.strides()]() mutable {
+    auto ndim = shapes.size();
+    if (contig) {
+      constexpr int N = simd::max_size<T>;
+      while (data_size >= N) {
+        simd::store(dst, Op{}(simd::load<T, N>(src)));
+        data_size -= N;
+        src += N;
+        dst += N;
+      }
+      while (data_size > 0) {
+        *dst = Op{}(*src);
+        data_size--;
+        dst++;
+        src++;
+      }
+    } else {
+      size_t shape = ndim > 0 ? shapes.back() : 1;
+      size_t stride = ndim > 0 ? strides.back() : 1;
+      if (ndim <= 1) {
+        unary_op<T, U, Op>(src, dst, shape, stride);
+        return;
+      }
+      auto it = ContiguousIterator(shapes, strides, ndim - 1);
+      for (size_t elem = 0; elem < size; elem += shape) {
+        unary_op<T, U, Op>(src + it.loc, dst + elem, shape, stride);
+        it.step();
+      }
     }
-    while (size > 0) {
-      *dst = op(*a_ptr);
-      size--;
-      dst++;
-      a_ptr++;
-    }
-  } else {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    U* dst = out.data<U>();
-    size_t shape = a.ndim() > 0 ? a.shape(-1) : 1;
-    size_t stride = a.ndim() > 0 ? a.strides(-1) : 1;
-    if (a.ndim() <= 1) {
-      unary_op(a_ptr, dst, op, shape, stride);
-      return;
-    }
-    ContiguousIterator it(a.shape(), a.strides(), a.ndim() - 1);
-    for (size_t elem = 0; elem < a.size(); elem += shape) {
-      unary_op(a_ptr + it.loc, dst + elem, op, shape, stride);
-      it.step();
-    }
-  }
+  });
 }
 
 template <typename Op>

mlx/backend/metal/allocator.cpp

@@ -10,6 +10,9 @@
 
 namespace mlx::core {
 
+constexpr size_t resource_options =
+    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked;
+
 namespace allocator {
 
 Allocator& allocator() {
@@ -150,15 +153,34 @@ MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
       residency_set_(device_),
       buffer_cache_(device_) {
-  auto memsize = std::get<size_t>(device_info()["memory_size"]);
+  auto pool = metal::new_scoped_memory_pool();
+  auto memsize = std::get<size_t>(device_info().at("memory_size"));
   auto max_rec_size =
-      std::get<size_t>(device_info()["max_recommended_working_set_size"]);
-  resource_limit_ = std::get<size_t>(device_info()["resource_limit"]);
+      std::get<size_t>(device_info().at("max_recommended_working_set_size"));
+  resource_limit_ = std::get<size_t>(device_info().at("resource_limit"));
   block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
   gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
   max_pool_size_ = block_limit_;
   device(mlx::core::Device::gpu)
       .set_residency_set(residency_set_.mtl_residency_set());
+  bool is_vm = std::get<std::string>(device_info().at("device_name")) ==
+      "Apple Paravirtual device";
+  if (is_vm) {
+    return;
+  }
+  auto heap_desc = MTL::HeapDescriptor::alloc()->init();
+  heap_desc->setResourceOptions(resource_options);
+  heap_desc->setSize(heap_size_);
+  heap_ = device_->newHeap(heap_desc);
+  heap_desc->release();
+  residency_set_.insert(heap_);
+}
+
+MetalAllocator::~MetalAllocator() {
+  auto pool = metal::new_scoped_memory_pool();
+  if (heap_) {
+    heap_->release();
+  }
 }
 
 size_t MetalAllocator::set_cache_limit(size_t limit) {
@@ -226,8 +248,6 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
     }
 
     // Allocate new buffer if needed
-    size_t res_opt = MTL::ResourceStorageModeShared;
-    res_opt |= MTL::ResourceHazardTrackingModeUntracked;
     if (num_resources_ >= resource_limit_) {
       std::ostringstream msg;
       msg << "[metal::malloc] Resource limit (" << resource_limit_
@@ -235,7 +255,12 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
       throw std::runtime_error(msg.str());
     }
     lk.unlock();
-    buf = device_->newBuffer(size, res_opt);
+    if (size < small_size_ && heap_) {
+      buf = heap_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      buf = device_->newBuffer(size, resource_options);
+    }
     lk.lock();
     if (buf) {
       num_resources_++;
@@ -246,13 +271,15 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   peak_memory_ = std::max(peak_memory_, active_memory_);
 
   // Maintain the cache below the requested limit
-  if (get_cache_memory() >= max_pool_size_) {
+  if (get_cache_memory() > max_pool_size_) {
     auto pool = metal::new_scoped_memory_pool();
     num_resources_ -= buffer_cache_.release_cached_buffers(
         get_cache_memory() - max_pool_size_);
   }
 
-  residency_set_.insert(buf);
+  if (!buf->heap()) {
+    residency_set_.insert(buf);
+  }
 
   return Buffer{static_cast<void*>(buf)};
 }
@@ -269,7 +296,9 @@ void MetalAllocator::free(Buffer buffer) {
     return;
   }
   std::unique_lock lk(mutex_);
-  residency_set_.erase(buf);
+  if (!buf->heap()) {
+    residency_set_.erase(buf);
+  }
   active_memory_ -= buf->length();
   if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
@@ -301,7 +330,7 @@ size_t set_memory_limit(size_t limit, bool relaxed /* = true */) {
 }
 size_t set_wired_limit(size_t limit) {
   if (limit >
-      std::get<size_t>(device_info()["max_recommended_working_set_size"])) {
+      std::get<size_t>(device_info().at("max_recommended_working_set_size"))) {
     throw std::invalid_argument(
         "[metal::set_wired_limit] Setting a wired limit larger than "
         "the maximum working set size is not allowed.");

mlx/backend/metal/allocator.h

@@ -43,6 +43,7 @@ class BufferCache {
   void remove_from_list(BufferHolder* to_remove);
 
   MTL::Device* device_;
+  MTL::Heap* heap_{nullptr};
 
   std::multimap<size_t, BufferHolder*> buffer_pool_;
   BufferHolder* head_;
@@ -78,7 +79,15 @@ class MetalAllocator : public allocator::Allocator {
 
  private:
   MTL::Device* device_;
+
+  // The size of allocations which go on the heap until it is full. This size
+  // is chosen because it is the actual minimum size of a buffer allocated from
+  // the heap, a heap can have at most heap.size() / 256 buffers.
+  static constexpr int small_size_ = 256;
+  static constexpr int heap_size_ = 1 << 20;
+  MTL::Heap* heap_;
   MetalAllocator();
+  ~MetalAllocator();
   friend MetalAllocator& allocator();
 
   // Caching allocator

mlx/backend/metal/binary.cpp

@@ -102,16 +102,9 @@ void binary_op_gpu_inplace(
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
 
-  // - If a is donated it goes to the first output
-  // - If b is donated it goes to the first output if a was not donated
-  //   otherwise it goes to the second output.
-  // - If there is only one output only one of a and b will be donated.
-  bool donate_a = a.data_shared_ptr() == nullptr;
-  bool donate_b = b.data_shared_ptr() == nullptr;
   int arg_idx = 0;
-  compute_encoder.set_input_array(donate_a ? outputs[0] : a, arg_idx++);
-  compute_encoder.set_input_array(
-      donate_b ? (donate_a ? outputs[1] : outputs[0]) : b, arg_idx++);
+  compute_encoder.set_input_array(a, arg_idx++);
+  compute_encoder.set_input_array(b, arg_idx++);
   compute_encoder.set_output_array(outputs[0], arg_idx++);
   if (outputs.size() == 2) {
     compute_encoder.set_output_array(outputs[1], arg_idx++);
@@ -164,8 +157,8 @@ void binary_op_gpu(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, outputs[0], bopt, true);
-  set_binary_op_output_data(a, b, outputs[1], bopt, true);
+  set_binary_op_output_data(a, b, outputs[0], bopt);
+  set_binary_op_output_data(a, b, outputs[1], bopt);
   binary_op_gpu_inplace(inputs, outputs, op, s);
 }
 
@@ -195,7 +188,7 @@ void binary_op_gpu(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt, true);
+  set_binary_op_output_data(a, b, out, bopt);
   binary_op_gpu_inplace(inputs, out, op, s);
 }
 

mlx/backend/metal/compiled.cpp

@@ -457,7 +457,7 @@ void Compiled::eval_gpu(
   }
 
   compiled_allocate_outputs(
-      inputs, outputs, inputs_, constant_ids_, contiguous, true);
+      inputs, outputs, inputs_, constant_ids_, contiguous);
 
   // Put the outputs in
   for (auto& x : outputs) {

mlx/backend/metal/copy.cpp

@@ -14,25 +14,11 @@ namespace mlx::core {
 constexpr int MAX_COPY_SPECIALIZED_DIMS = 3;
 
 void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
-  if (ctype == CopyType::Vector) {
-    // If the input is donateable, we are doing a vector copy and the types
-    // have the same size, then the input buffer can hold the output.
-    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
-      out.move_shared_buffer(in);
-      // If the output has the same type as the input then there is nothing to
-      // copy, just use the buffer.
-      if (in.dtype() == out.dtype()) {
-        return;
-      }
-    } else {
-      out.set_data(
-          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
-          in.data_size(),
-          in.strides(),
-          in.flags());
-    }
-  } else {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  bool donated = set_copy_output_data(in, out, ctype);
+  if (donated && in.dtype() == out.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
   }
   if (ctype == CopyType::GeneralGeneral) {
     ctype = CopyType::General;

mlx/backend/metal/device.cpp

@@ -168,9 +168,10 @@ void CommandEncoder::set_output_array(
   register_output_array(a);
 }
 
-void CommandEncoder::register_output_array(array& a) {
+void CommandEncoder::register_output_array(const array& a) {
   all_outputs_.insert(a.buffer().ptr());
-  auto buf = static_cast<MTL::Resource*>(a.buffer().ptr());
+
+  auto buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
   if (concurrent_) {
     concurrent_outputs_.insert(buf);
   } else {
@@ -249,14 +250,12 @@ Device::~Device() {
   for (auto& l : library_map_) {
     l.second->release();
   }
+  stream_map_.clear();
   device_->release();
 }
 
 void Device::new_queue(int index) {
   auto thread_pool = metal::new_scoped_memory_pool();
-
-  // Multiple threads can ask the device for queues
-  // We lock this as a critical section for safety
   auto q = device_->newCommandQueue(MAX_BUFFERS_PER_QUEUE);
   debug_set_stream_queue_label(q, index);
   if (!q) {
@@ -269,6 +268,10 @@ void Device::new_queue(int index) {
   }
 }
 
+MTL::CommandQueue* Device::get_queue(Stream stream) {
+  return get_stream_(stream.index).queue;
+}
+
 bool Device::command_buffer_needs_commit(int index) {
   auto& stream = get_stream_(index);
   if (stream.buffer_ops > max_ops_per_buffer_ ||
@@ -690,12 +693,13 @@ void new_stream(Stream stream) {
   }
 }
 
-std::unordered_map<std::string, std::variant<std::string, size_t>>
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
 device_info() {
   auto init_device_info = []()
       -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
     auto pool = new_scoped_memory_pool();
     auto raw_device = device(default_device()).mtl_device();
+    auto name = std::string(raw_device->name()->utf8String());
     auto arch = std::string(raw_device->architecture()->name()->utf8String());
 
     size_t memsize = 0;
@@ -709,6 +713,7 @@ device_info() {
     }
 
     return {
+        {"device_name", name},
         {"architecture", arch},
         {"max_buffer_length", raw_device->maxBufferLength()},
         {"max_recommended_working_set_size",

mlx/backend/metal/device.h

@@ -62,7 +62,7 @@ struct CommandEncoder {
 
   void set_input_array(const array& a, int idx, int64_t offset = 0);
   void set_output_array(array& a, int idx, int64_t offset = 0);
-  void register_output_array(array& a);
+  void register_output_array(const array& a);
   void dispatch_threadgroups(MTL::Size grid_dims, MTL::Size group_dims);
   void dispatch_threads(MTL::Size grid_dims, MTL::Size group_dims);
   void maybeInsertBarrier();
@@ -178,6 +178,9 @@ class Device {
   }
 
   void new_queue(int index);
+
+  MTL::CommandQueue* get_queue(Stream stream);
+
   MTL::CommandBuffer* get_command_buffer(int index);
   bool command_buffer_needs_commit(int index);
   void commit_command_buffer(int index);

mlx/backend/metal/distributed.cpp

@@ -4,149 +4,30 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/device.h"
-#include "mlx/backend/metal/event.h"
-#include "mlx/backend/metal/fence.h"
 #include "mlx/backend/metal/utils.h"
 #include "mlx/distributed/ops.h"
 #include "mlx/distributed/primitives.h"
+#include "mlx/fence.h"
 #include "mlx/scheduler.h"
 
 namespace mlx::core::distributed {
 
-void signal_and_wait(const Event& e_signal, const Event& e_wait) {
-  if (e_signal.valid()) {
-    encode_signal(e_signal);
-  }
-  encode_wait(e_wait);
+void AllReduce::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllReduce::eval_gpu] has no GPU implementation.");
 }
 
-void AllReduce::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-  assert(outputs.size() == 1);
-
-  auto& in = inputs[0];
-
-  Fence f{stream()};
-
-  if (in.event().valid()) {
-    f.update_gpu(in);
-  }
-
-  auto& out = outputs[0];
-  if (in.is_donatable()) {
-    out.move_shared_buffer(in);
-  } else {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  }
-  f.wait_gpu(out);
-
-  auto task = [in = in,
-               out = unsafe_weak_copy(out),
-               f = std::move(f),
-               reduce_type = reduce_type_,
-               group = group()]() mutable {
-    if (in.event().valid()) {
-      f.wait();
-    }
-    switch (reduce_type) {
-      case Sum:
-        distributed::detail::all_sum(
-            group, in.data_shared_ptr() == nullptr ? out : in, out);
-        break;
-      default:
-        throw std::runtime_error("Only all reduce sum is supported for now");
-    }
-    f.update();
-  };
-  scheduler::enqueue(detail::communication_stream(), std::move(task));
+void AllGather::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllGather::eval_gpu] has no GPU implementation.");
 }
 
-void AllGather::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-  assert(outputs.size() == 1);
-  auto& in = inputs[0];
-  auto& out = outputs[0];
-
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  Fence f{stream()};
-
-  if (in.event().valid()) {
-    f.update_gpu(in);
-  }
-  f.wait_gpu(out);
-
-  auto task = [in = in,
-               out = unsafe_weak_copy(out),
-               f = std::move(f),
-               group = group()]() mutable {
-    if (in.event().valid()) {
-      f.wait();
-    }
-    distributed::detail::all_gather(group, in, out);
-    f.update();
-  };
-  scheduler::enqueue(detail::communication_stream(), std::move(task));
+void Send::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Send::eval_gpu] has no GPU implementation.");
 }
 
-void Send::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-  assert(outputs.size() == 1);
-
-  auto& in = inputs[0];
-
-  // Encode a signal event for the input
-  Fence f{stream()};
-  if (in.event().valid()) {
-    f.update_gpu(in);
-  }
-
-  auto& out = outputs[0];
-  move_or_copy(in, out);
-
-  // Schedule an async send on the comm stream
-  auto task = [in = in,
-               out = unsafe_weak_copy(out),
-               f = std::move(f),
-               group = group(),
-               dst = dst_]() mutable {
-    if (in.event().valid()) {
-      f.wait();
-    }
-    distributed::detail::send(group, out, dst);
-  };
-  scheduler::enqueue(detail::communication_stream(), std::move(task));
-}
-
-void Recv::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 0);
-  assert(outputs.size() == 1);
-
-  auto& out = outputs[0];
-
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  Fence f{stream()};
-  f.wait_gpu(out);
-
-  // Schedule an async recv on the comm stream
-  auto task = [out = unsafe_weak_copy(out),
-               f = std::move(f),
-               group = group(),
-               src = src_]() mutable {
-    distributed::detail::recv(group, out, src);
-    f.update();
-  };
-  scheduler::enqueue(detail::communication_stream(), std::move(task));
+void Recv::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Recv::eval_gpu] has no GPU implementation.");
 }
 
 } // namespace mlx::core::distributed

mlx/backend/metal/event.cpp

@@ -3,52 +3,58 @@
 #include "mlx/event.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/metal_impl.h"
+#include "mlx/scheduler.h"
 
 namespace mlx::core {
 
-void encode_wait(Event e) {
-  auto& d = metal::device(e.stream().device);
-  d.end_encoding(e.stream().index);
-  auto command_buffer = d.get_command_buffer(e.stream().index);
-  command_buffer->encodeWait(
-      static_cast<MTL::Event*>(e.raw_event().get()), e.value());
-  command_buffer->addCompletedHandler(
-      [e = std::move(e)](MTL::CommandBuffer* cbuf) {});
-}
-
-void encode_signal(Event e) {
-  auto& d = metal::device(e.stream().device);
-  d.end_encoding(e.stream().index);
-  auto command_buffer = d.get_command_buffer(e.stream().index);
-  command_buffer->encodeSignalEvent(
-      static_cast<MTL::Event*>(e.raw_event().get()), e.value());
-  command_buffer->addCompletedHandler(
-      [e = std::move(e)](MTL::CommandBuffer* cbuf) {});
-}
-
-Event::Event(const Stream& stream) : stream_(stream) {
+Event::Event(Stream stream) : stream_(stream) {
   auto dtor = [](void* ptr) {
     auto p = metal::new_scoped_memory_pool();
     static_cast<MTL::SharedEvent*>(ptr)->release();
   };
   auto p = metal::new_scoped_memory_pool();
   event_ = std::shared_ptr<void>(
-      metal::device(stream.device).mtl_device()->newSharedEvent(), dtor);
+      metal::device(Device::gpu).mtl_device()->newSharedEvent(), dtor);
 }
 
 void Event::wait() {
-  if (!static_cast<MTL::SharedEvent*>(raw_event().get())
+  if (!static_cast<MTL::SharedEvent*>(event_.get())
            ->waitUntilSignaledValue(value(), -1)) {
     throw std::runtime_error("[Event::wait] Timed out");
   }
 }
 
 void Event::signal() {
-  static_cast<MTL::SharedEvent*>(raw_event().get())->setSignaledValue(value());
+  static_cast<MTL::SharedEvent*>(event_.get())->setSignaledValue(value());
+}
+
+void Event::wait(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable { wait(); });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
+}
+
+void Event::signal(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable { signal(); });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
 }
 
 bool Event::is_signaled() const {
-  return static_cast<MTL::SharedEvent*>(raw_event().get())->signaledValue() >=
+  return static_cast<MTL::SharedEvent*>(event_.get())->signaledValue() >=
       value();
 }
 

mlx/backend/metal/event.h

@@ -1,10 +0,0 @@
-// Copyright © 2024 Apple Inc.
-#pragma once
-
-namespace mlx::core {
-
-void encode_wait(Event e);
-
-void encode_signal(Event e);
-
-} // namespace mlx::core

mlx/backend/metal/fence.cpp

@@ -1,49 +1,108 @@
 // Copyright © 2024 Apple Inc.
+#include <csignal>
 
-#include "mlx/backend/metal/fence.h"
 #include "mlx/backend/metal/device.h"
-#include "mlx/backend/metal/event.h"
 #include "mlx/backend/metal/metal_impl.h"
+#include "mlx/fence.h"
+#include "mlx/scheduler.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {
 
-Fence::Fence(const Stream& stream) : stream_(stream) {
-  auto d = metal::device(stream.device).mtl_device();
-  if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
-    use_fast_ = false;
-  } else if (__builtin_available(macOS 15, iOS 18, *)) {
-    use_fast_ = env::metal_fast_synch();
-  }
-  if (!use_fast_) {
-    // Wraps Metal SharedEvent
-    auto dtor = [](void* ptr) {
-      auto p = metal::new_scoped_memory_pool();
-      static_cast<MTL::SharedEvent*>(ptr)->release();
-    };
-    auto p = metal::new_scoped_memory_pool();
-    fence_ = std::shared_ptr<void>(
-        metal::device(stream.device).mtl_device()->newSharedEvent(), dtor);
-  } else {
+void signal_handler(int signum);
+
+MTL::Buffer* signal_buffer() {
+  auto init = []() {
+    signal(SIGTERM, signal_handler);
     auto dtor = [](void* buf) {
       allocator::free(static_cast<MTL::Buffer*>(buf));
     };
-    auto buf = allocator::malloc_or_wait(sizeof(uint32_t)).ptr();
-    fence_ = std::shared_ptr<void>(buf, dtor);
-    cpu_value()[0] = 0;
-  }
+    auto buf = std::shared_ptr<void>(
+        allocator::malloc_or_wait(sizeof(uint32_t)).ptr(), dtor);
+    static_cast<uint32_t*>(
+        static_cast<MTL::Buffer*>(buf.get())->contents())[0] = 0;
+    return buf;
+  };
+  static std::shared_ptr<void> buf = init();
+  return static_cast<MTL::Buffer*>(buf.get());
 }
 
-void Fence::wait_gpu(array& x) {
-  gpu_count_++;
-  auto& d = metal::device(stream_.device);
-  auto idx = stream_.index;
+void signal_handler(int signum) {
+  auto buf = signal_buffer();
+  static_cast<std::atomic_uint*>(buf->contents())[0] = 1;
+  signal(signum, SIG_DFL);
+  raise(signum);
+}
 
-  if (!use_fast_) {
+struct FenceImpl {
+  FenceImpl() {
+    auto d = metal::device(Device::gpu).mtl_device();
+    if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
+      use_fast = false;
+    } else if (__builtin_available(macOS 15, iOS 18, *)) {
+      use_fast = env::metal_fast_synch();
+    }
+
+    if (!use_fast) {
+      auto p = metal::new_scoped_memory_pool();
+      fence = static_cast<void*>(d->newSharedEvent());
+    } else {
+      auto buf = allocator::malloc_or_wait(sizeof(uint32_t)).ptr();
+      fence = static_cast<void*>(buf);
+      cpu_value()[0] = 0;
+    }
+  }
+
+  ~FenceImpl() {
+    if (!use_fast) {
+      // Wraps Metal SharedEvent
+      auto p = metal::new_scoped_memory_pool();
+      static_cast<MTL::SharedEvent*>(fence)->release();
+    } else {
+      allocator::free(static_cast<MTL::Buffer*>(fence));
+    }
+  }
+  bool use_fast{false};
+  uint32_t count{0};
+  void* fence;
+
+  std::atomic_uint* cpu_value() {
+    return static_cast<std::atomic_uint*>(
+        static_cast<MTL::Buffer*>(fence)->contents());
+  }
+};
+
+Fence::Fence(Stream) {
+  auto dtor = [](void* ptr) { delete static_cast<FenceImpl*>(ptr); };
+  fence_ = std::shared_ptr<void>(new FenceImpl{}, dtor);
+}
+
+void Fence::wait(Stream stream, const array& x) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        if (!static_cast<MTL::SharedEvent*>(f.fence)->waitUntilSignaledValue(
+                count, -1)) {
+          throw std::runtime_error("[Fence::wait] Timed out");
+        }
+        return;
+      }
+      while (f.cpu_value()[0] < count) {
+      }
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+
+  if (!f.use_fast) {
     d.end_encoding(idx);
     auto command_buffer = d.get_command_buffer(idx);
-    command_buffer->encodeWait(
-        static_cast<MTL::Event*>(fence_.get()), gpu_count_);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(f.fence), f.count);
     command_buffer->addCompletedHandler(
         [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
     return;
@@ -51,7 +110,7 @@ void Fence::wait_gpu(array& x) {
 
   auto& compute_encoder = d.get_command_encoder(idx);
 
-  // Register the output to ensure that no kernels which depends on the
+  // Register outputs to ensure that no kernels which depends on the
   // output starts before this one is done
   compute_encoder.register_output_array(x);
 
@@ -59,36 +118,50 @@ void Fence::wait_gpu(array& x) {
   MTL::Size kernel_dims = MTL::Size(1, 1, 1);
   compute_encoder.set_compute_pipeline_state(kernel);
 
-  auto buf = static_cast<MTL::Buffer*>(fence_.get());
+  auto buf = static_cast<MTL::Buffer*>(f.fence);
   compute_encoder.set_buffer(buf, 0);
-  compute_encoder.set_bytes(gpu_count_, 1);
+  compute_encoder.set_bytes(f.count, 1);
+  compute_encoder.set_buffer(signal_buffer(), 2);
   compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
 
   d.get_command_buffer(idx)->addCompletedHandler(
-      [fence = fence_](MTL::CommandBuffer* cbuf) {});
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
 }
 
-void Fence::update_gpu(const array& x) {
-  gpu_count_++;
-  auto& d = metal::device(stream_.device);
-  auto idx = stream_.index;
+void Fence::update(Stream stream, const array& x) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+  f.count++;
 
-  if (!use_fast_) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        static_cast<MTL::SharedEvent*>(f.fence)->setSignaledValue(count);
+        return;
+      }
+
+      f.cpu_value()[0] = count;
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+  if (!f.use_fast) {
     d.end_encoding(idx);
     auto command_buffer = d.get_command_buffer(idx);
     command_buffer->encodeSignalEvent(
-        static_cast<MTL::Event*>(fence_.get()), gpu_count_);
+        static_cast<MTL::Event*>(f.fence), f.count);
     command_buffer->addCompletedHandler(
         [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
     return;
   }
 
-  // Launch input visibility kernel
+  // Launch input visibility kernels
   auto& compute_encoder = d.get_command_encoder(idx);
   auto kernel = d.get_kernel("input_coherent");
   uint32_t nthreads =
       (x.data_size() * x.itemsize() + sizeof(uint32_t) - 1) / sizeof(uint32_t);
-
   MTL::Size group_dims = MTL::Size(1024, 1, 1);
   MTL::Size grid_dims = MTL::Size((nthreads + 1024 - 1) / 1024, 1, 1);
   compute_encoder.set_compute_pipeline_state(kernel);
@@ -96,7 +169,7 @@ void Fence::update_gpu(const array& x) {
   compute_encoder.set_bytes(nthreads, 1);
   compute_encoder.dispatch_threadgroups(group_dims, grid_dims);
 
-  // Barrier on previous kernel
+  // Barrier on previous kernels
   compute_encoder.barrier();
 
   // Launch value update kernel
@@ -104,42 +177,13 @@ void Fence::update_gpu(const array& x) {
   MTL::Size kernel_dims = MTL::Size(1, 1, 1);
   compute_encoder.set_compute_pipeline_state(kernel);
 
-  auto buf = static_cast<MTL::Buffer*>(fence_.get());
+  auto buf = static_cast<MTL::Buffer*>(f.fence);
   compute_encoder.set_buffer(buf, 0);
-  compute_encoder.set_bytes(gpu_count_, 1);
+  compute_encoder.set_bytes(f.count, 1);
   compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
 
   d.get_command_buffer(idx)->addCompletedHandler(
-      [fence = fence_](MTL::CommandBuffer* cbuf) {});
-}
-
-void Fence::wait() {
-  cpu_count_++;
-  if (!use_fast_) {
-    if (!static_cast<MTL::SharedEvent*>(fence_.get())
-             ->waitUntilSignaledValue(cpu_count_, -1)) {
-      throw std::runtime_error("[Event::wait] Timed out");
-    }
-    return;
-  }
-  while (cpu_value()[0] < cpu_count_) {
-  }
-}
-
-void Fence::update() {
-  cpu_count_++;
-
-  if (!use_fast_) {
-    static_cast<MTL::SharedEvent*>(fence_.get())->setSignaledValue(cpu_count_);
-    return;
-  }
-
-  cpu_value()[0] = cpu_count_;
-}
-
-std::atomic_uint* Fence::cpu_value() {
-  return static_cast<std::atomic_uint*>(
-      static_cast<MTL::Buffer*>(fence_.get())->contents());
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
 }
 
 } // namespace mlx::core

mlx/backend/metal/fence.h

@@ -1,40 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include "mlx/backend/metal/device.h"
-
-namespace mlx::core {
-
-/* A fence to be used for synchronizing work between the CPU and GPU
- *
- * Calls to `update_gpu` should be paired with calls to `wait`. This ensures
- * that the array passed to `update_gpu` is computed and visible to the CPU
- * after the call to `wait` returns.
- *
- * Calls to `update` should be paired with calls to `wait_gpu`. This ensures
- * that the array passed to `wait_gpu` will not be read by the GPU until the CPU
- * has called `update`.
- *
- * The fence supports slow (default) and fast mode. Fast mode requires setting
- * the environment variable `MLX_METAL_FAST_SYNCH=1`. Fast mode also requires
- * Metal 3.2+ (macOS 15+, iOS 18+).
- */
-class Fence {
- public:
-  Fence(const Stream& stream);
-
-  void update_gpu(const array& x);
-  void wait_gpu(array& x);
-
-  void wait();
-  void update();
-
- private:
-  Stream stream_;
-  std::shared_ptr<void> fence_;
-  uint32_t cpu_count_{0};
-  uint32_t gpu_count_{0};
-  bool use_fast_;
-  std::atomic_uint* cpu_value();
-};
-
-} // namespace mlx::core

mlx/backend/metal/hadamard.cpp

@@ -82,7 +82,7 @@ void Hadamard::eval_gpu(const std::vector<array>& inputs, array& out) {
   const array& in_contiguous = check_input(in);
 
   if (in_contiguous.is_donatable()) {
-    out.move_shared_buffer(in_contiguous);
+    out.copy_shared_buffer(in_contiguous);
   } else {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }

mlx/backend/metal/kernels/CMakeLists.txt

@@ -4,11 +4,12 @@ set(BASE_HEADERS
     bf16_math.h
     complex.h
     defines.h
+    erf.h
     expm1f.h
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
-  set(METAL_FLAGS -Wall -Wextra -fno-fast-math)
+  set(METAL_FLAGS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()

mlx/backend/metal/kernels/fence.metal

@@ -39,13 +39,14 @@ constexpr constant metal::thread_scope thread_scope_system =
 // single thread kernel to spin wait for timestamp value
 [[kernel]] void fence_wait(
     volatile coherent(system) device uint* timestamp [[buffer(0)]],
-    constant uint& value [[buffer(1)]]) {
+    constant uint& value [[buffer(1)]],
+    volatile coherent(system) device uint* sig_handler [[buffer(2)]]) {
   while (1) {
     metal::atomic_thread_fence(
         metal::mem_flags::mem_device,
         metal::memory_order_seq_cst,
         metal::thread_scope_system);
-    if (timestamp[0] >= value) {
+    if (timestamp[0] >= value || sig_handler[0] > 0) {
       break;
     }
   }

mlx/backend/metal/kernels/layer_norm.metal

@@ -7,6 +7,8 @@
 
 using namespace metal;
 
+constant bool has_w [[function_constant(20)]];
+
 template <typename T, int N_READS = RMS_N_READS>
 [[kernel]] void layer_norm_single_row(
     const device T* x,
@@ -327,7 +329,9 @@ template <typename T, int N_READS = RMS_N_READS>
       gx[i] = static_cast<T>(
           normalizer * (thread_w[i] * thread_g[i] - meanwg) -
           thread_x[i] * meanwgxc * normalizer2);
-      gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      }
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
@@ -336,7 +340,9 @@ template <typename T, int N_READS = RMS_N_READS>
         gx[i] = static_cast<T>(
             normalizer * (thread_w[i] * thread_g[i] - meanwg) -
             thread_x[i] * meanwgxc * normalizer2);
-        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+        }
       }
     }
   }
@@ -465,7 +471,9 @@ template <typename T, int N_READS = RMS_N_READS>
         float gi = g[i + r];
         gx[i + r] = static_cast<T>(
             normalizer * (wi * gi - meanwg) - xi * meanwgxc * normalizer2);
-        gw[i + r] = static_cast<T>(gi * xi);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi);
+        }
       }
     } else {
       for (int i = 0; i < N_READS; i++) {
@@ -475,7 +483,9 @@ template <typename T, int N_READS = RMS_N_READS>
           float gi = g[i + r];
           gx[i + r] = static_cast<T>(
               normalizer * (wi * gi - meanwg) - xi * meanwgxc * normalizer2);
-          gw[i + r] = static_cast<T>(gi * xi);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi);
+          }
         }
       }
     }

mlx/backend/metal/kernels/quantized.h

@@ -2015,9 +2015,9 @@ template <typename T, const int group_size, const int bits>
     device T* biases [[buffer(3)]],
     uint2 index [[thread_position_in_grid]],
     uint2 grid_dim [[threads_per_grid]]) {
-  constexpr T eps = T(1e-7);
+  constexpr float eps = 1e-7;
   constexpr int simd_size = 32;
-  constexpr T n_bins = (1 << bits) - 1;
+  constexpr float n_bins = (1 << bits) - 1;
   constexpr int packs_per_int = bits == 3 ? 8 : bits == 6 ? 4 : 8 / bits;
   constexpr int values_per_reduce = group_size / simd_size;
   constexpr int writes_per_reduce = packs_per_int / values_per_reduce;
@@ -2036,13 +2036,13 @@ template <typename T, const int group_size, const int bits>
       ? offset * writes_per_pack
       : offset * bytes_per_pack / writes_per_reduce;
 
-  T w_thread[values_per_reduce];
-  T w_min = Limits<T>::max;
-  T w_max = 0;
+  float w_thread[values_per_reduce];
+  float w_min = Limits<T>::max;
+  float w_max = 0;
 
 #pragma clang loop unroll(full)
   for (int i = 0; i < values_per_reduce; i++) {
-    T val = w[in_index + i];
+    float val = w[in_index + i];
     w_thread[i] = val;
     w_min = min(w_min, val);
     w_max = max(w_max, val);
@@ -2051,20 +2051,20 @@ template <typename T, const int group_size, const int bits>
   w_min = simd_min(w_min);
   w_max = simd_max(w_max);
 
-  T scale = max((w_max - w_min) / n_bins, eps);
+  float scale = max((w_max - w_min) / n_bins, eps);
   bool side = abs(w_min) > abs(w_max);
   scale = side ? scale : -scale;
-  T edge = side ? w_min : w_max;
-  T q0 = round(edge / scale);
+  float edge = side ? w_min : w_max;
+  float q0 = round(edge / scale);
   bool at_zero = q0 == 0.0f;
   scale = at_zero ? scale : edge / q0;
-  T bias = at_zero ? T(0) : edge;
+  float bias = at_zero ? 0 : edge;
 
   // Write out the scales and biases
   size_t gindex = in_index / group_size;
   if (in_index % group_size == 0) {
-    scales[gindex] = scale;
-    biases[gindex] = bias;
+    scales[gindex] = static_cast<T>(scale);
+    biases[gindex] = static_cast<T>(bias);
   }
 
   // We accumulate 3 bytes worth for 3/6 bit so we need a uint32_t

mlx/backend/metal/kernels/rms_norm.metal

@@ -7,6 +7,8 @@
 
 using namespace metal;
 
+constant bool has_w [[function_constant(20)]];
+
 template <typename T, int N_READS = RMS_N_READS>
 [[kernel]] void rms_single_row(
     const device T* x,
@@ -243,7 +245,9 @@ template <typename T, int N_READS = RMS_N_READS>
       gx[i] = static_cast<T>(
           thread_g[i] * thread_w[i] * normalizer -
           thread_x[i] * meangwx * normalizer3);
-      gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      }
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
@@ -251,7 +255,9 @@ template <typename T, int N_READS = RMS_N_READS>
         gx[i] = static_cast<T>(
             thread_g[i] * thread_w[i] * normalizer -
             thread_x[i] * meangwx * normalizer3);
-        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        }
       }
     }
   }
@@ -351,7 +357,9 @@ template <typename T, int N_READS = RMS_N_READS>
 
         gx[i + r] =
             static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
-        gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        }
       }
     } else {
       for (int i = 0; i < N_READS; i++) {
@@ -362,7 +370,9 @@ template <typename T, int N_READS = RMS_N_READS>
 
           gx[i + r] =
               static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
-          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          }
         }
       }
     }

mlx/backend/metal/kernels/sdpa_vector.h

@@ -5,6 +5,7 @@
 using namespace metal;
 
 constant bool has_mask [[function_constant(20)]];
+constant bool query_transposed [[function_constant(21)]];
 
 template <typename T, int D, int V = D>
 [[kernel]] void sdpa_vector(
@@ -18,9 +19,11 @@ template <typename T, int D, int V = D>
     const constant size_t& v_stride,
     const constant float& scale,
     const device bool* mask [[function_constant(has_mask)]],
-    const constant int& mask_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_kv_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride [[function_constant(has_mask)]],
     const constant int& mask_head_stride [[function_constant(has_mask)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 32;
@@ -41,15 +44,21 @@ template <typename T, int D, int V = D>
   threadgroup U sum_exp_scores[BN];
 
   // Adjust positions
-  const int head_idx = tid.y;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
   const int kv_head_idx = head_idx / gqa_factor;
-  queries += head_idx * D + simd_lid * qk_per_thread;
+  const int o_offset = tpg.x * q_seq_idx + head_idx;
+  const int q_offset =
+      query_transposed ? o_offset : head_idx * tpg.y + q_seq_idx;
+  queries += q_offset * D + simd_lid * qk_per_thread;
   keys += kv_head_idx * k_stride + simd_gid * D + simd_lid * qk_per_thread;
   values += kv_head_idx * v_stride + simd_gid * V + simd_lid * v_per_thread;
   if (has_mask) {
-    mask += head_idx * mask_head_stride + simd_gid * mask_seq_stride;
+    mask += head_idx * mask_head_stride + simd_gid * mask_kv_seq_stride +
+        q_seq_idx * mask_q_seq_stride;
   }
-  out += head_idx * V + simd_gid * v_per_thread;
+
+  out += o_offset * V + simd_gid * v_per_thread;
 
   // Read the query and 0 the output accumulator
   for (int i = 0; i < qk_per_thread; i++) {
@@ -95,7 +104,7 @@ template <typename T, int D, int V = D>
     keys += inner_k_stride;
     values += inner_v_stride;
     if (has_mask) {
-      mask += BN * mask_seq_stride;
+      mask += BN * mask_kv_seq_stride;
     }
   }
 
@@ -142,9 +151,11 @@ template <typename T, int D, int V = D>
     const constant size_t& v_stride,
     const constant float& scale,
     const device bool* mask [[function_constant(has_mask)]],
-    const constant int& mask_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_kv_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride [[function_constant(has_mask)]],
     const constant int& mask_head_stride [[function_constant(has_mask)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 8;
@@ -167,20 +178,26 @@ template <typename T, int D, int V = D>
 
   // Adjust positions
   const int block_idx = tid.z;
-  const int head_idx = tid.y;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int o_offset = tpg.x * q_seq_idx + head_idx;
+  const int q_offset =
+      query_transposed ? o_offset : head_idx * tpg.y + q_seq_idx;
   const int kv_head_idx = head_idx / gqa_factor;
-  queries += head_idx * D + simd_lid * qk_per_thread;
+
+  queries += q_offset * D + simd_lid * qk_per_thread;
   keys += kv_head_idx * k_stride + (block_idx * BN + simd_gid) * D +
       simd_lid * qk_per_thread;
   values += kv_head_idx * v_stride + (block_idx * BN + simd_gid) * V +
       simd_lid * v_per_thread;
-  out += head_idx * blocks * V + block_idx * V + simd_lid * v_per_thread;
+  out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
   if (has_mask) {
     mask += head_idx * mask_head_stride +
-        (block_idx * BN + simd_gid) * mask_seq_stride;
+        (block_idx * BN + simd_gid) * mask_kv_seq_stride +
+        q_seq_idx * mask_q_seq_stride;
   }
-  sums += head_idx * blocks + block_idx;
-  maxs += head_idx * blocks + block_idx;
+  sums += o_offset * blocks + block_idx;
+  maxs += o_offset * blocks + block_idx;
 
   // Read the query and 0 the output accumulator
   for (int i = 0; i < qk_per_thread; i++) {
@@ -226,7 +243,7 @@ template <typename T, int D, int V = D>
     keys += blocks * inner_k_stride;
     values += blocks * inner_v_stride;
     if (has_mask) {
-      mask += BN * blocks * mask_seq_stride;
+      mask += BN * blocks * mask_kv_seq_stride;
     }
   }
 
@@ -275,6 +292,7 @@ template <typename T, int D>
     const device float* maxs [[buffer(2)]],
     device T* out [[buffer(3)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 32;
@@ -288,11 +306,14 @@ template <typename T, int D>
   threadgroup U outputs[BN * BD];
 
   // Adjust positions
-  const int head_idx = tid.y;
-  partials += head_idx * blocks * D + simd_gid * D + simd_lid * elem_per_thread;
-  sums += head_idx * blocks;
-  maxs += head_idx * blocks;
-  out += head_idx * D + simd_gid * elem_per_thread;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int n_heads = tpg.x;
+  const int q_offset = n_heads * q_seq_idx + head_idx;
+  partials += q_offset * blocks * D + simd_gid * D + simd_lid * elem_per_thread;
+  sums += q_offset * blocks;
+  maxs += q_offset * blocks;
+  out += q_offset * D + simd_gid * elem_per_thread;
 
   // First everybody reads the max and sum_exp
   U max_score = maxs[simd_lid];

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h

@@ -50,7 +50,7 @@ struct SubOp {
 struct ExpSubOp {
   template <typename T>
   METAL_FUNC static constexpr T apply(T x, T y) {
-    return fast::exp(x - y);
+    return fast::exp2(x - y);
   }
 };
 
@@ -103,17 +103,24 @@ template <
       tidl.x * BQ * params->O_strides[2]; // Seqeunce
 
   // Prepare threadgroup memory
-  constexpr short padQ = 0; // 16 / sizeof(T);
-  constexpr short padK = 0; // 16 / sizeof(T);
-  constexpr short padV = 0; // 16 / sizeof(T);
+  constexpr short padQ = 16 / sizeof(T);
+  constexpr short padK = 16 / sizeof(T);
+  constexpr short padV = 16 / sizeof(T);
 
   constexpr short LDQ_tgp = BD + padQ;
   constexpr short LDK_tgp = BK + padK;
   constexpr short LDV_tgp = BD + padV;
 
-  threadgroup T Qs[BQ * (BD + padQ)];
-  threadgroup T Ks[(BK + padK) * BD];
-  threadgroup T Vs[BK * (BD + padV)];
+  constexpr short tgp_mem_0 = (BK + padK) * (BD);
+  constexpr short tgp_mem_1 = BK * (BD + padV);
+  constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1;
+
+  threadgroup T Q_smem[BQ * (BD + padQ)];
+  threadgroup T KV_smem[tgp_mem_s];
+
+  threadgroup T* Qs = Q_smem;
+  threadgroup T* Ks = KV_smem;
+  threadgroup T* Vs = KV_smem;
 
   // Prepare block loaders
   using QBlockLoader = BlockLoaderT<
@@ -151,7 +158,7 @@ template <
   VBlockLoader loader_v(
       V, params->V_strides[2], Vs, simd_group_id, simd_lane_id);
 
-  TransformScale<T> ts(static_cast<T>(params->scale));
+  TransformScale<T> ts(static_cast<T>(params->scale * 1.44269504089));
 
   // Prepare MMA tiles
   constexpr short kFragSize = 8; // MMAFrag size
@@ -174,7 +181,7 @@ template <
   MMATile<AccumType, TQ, 1, MMAFrag_acc_t> Qtile;
   MMATile<AccumType, 1, TK, MMAFrag_acc_t> Ktile;
   MMATile<AccumType, TQ, TK, MMAFrag_acc_t> Stile;
-  MMATile<AccumType, TK, TD, MMAFrag_acc_t> Vtile;
+  MMATile<AccumType, 1, 1, MMAFrag_acc_t> Vtile;
   MMATile<AccumType, TQ, TD, MMAFrag_acc_t> Otile;
 
   Otile.clear();
@@ -224,11 +231,12 @@ template <
       loader_k.load_unsafe();
     }
 
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
     // Do S = Q @ K.T
     Stile.clear();
 
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
     for (short dd = 0; dd < TD; dd++) {
       simdgroup_barrier(mem_flags::mem_none);
 
@@ -264,7 +272,7 @@ template <
       }
     }
 
-    simdgroup_barrier(mem_flags::mem_none);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     // Load V blocks
     if (!align_K && kb == (params->NK_aligned)) {
@@ -292,7 +300,7 @@ template <
     // Factor exp(rowmax(Si) - rowmax(Si-1))
     STEEL_PRAGMA_UNROLL
     for (short i = 0; i < kRowsPT; ++i) {
-      factor[i] = fast::exp(max_score[i] - new_max[i]);
+      factor[i] = fast::exp2(max_score[i] - new_max[i]);
     }
 
     // Save max for next iteration
@@ -316,12 +324,35 @@ template <
 
     // Load V into registers
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    Vtile.template load<T, 1, 1, LDV_tgp, 1>(&Vs[Vs_offset]);
 
-    simdgroup_barrier(mem_flags::mem_none);
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short id = 0; id < TD; id++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
 
-    // Do O = S @ V
-    tile_matmad(Otile, Stile, Vtile, Otile);
+          const short kk = ik * kFragSize;
+          const short dd = id * kFragSize;
+
+          Vtile.template load<T, 1, 1, LDV_tgp, 1>(
+              &Vs[Vs_offset + kk * LDV_tgp + dd]);
+
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          MMAFrag_acc_t::mma(
+              Otile.frag_at(iq, id),
+              Stile.frag_at(iq, ik),
+              Vtile.frag_at(0, 0),
+              Otile.frag_at(iq, id));
+        }
+      }
+    }
 
     // Prepare for next iteration
     loader_k.next();

mlx/backend/metal/kernels/steel/attn/mma.h

@@ -62,6 +62,12 @@ struct BaseMMAFrag<T, 8, 8> {
   typedef metal::vec<T, kElemRows> row_frag_type;
   typedef metal::vec<T, kElemCols> col_frag_type;
 
+  template <typename U>
+  using dtype_mat_t = typename metal::simdgroup_matrix<U, kFragRows, kFragCols>;
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
   METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id
                                                [[thread_index_in_simdgroup]]) {
     const short qid = simd_lane_id / 4;
@@ -158,30 +164,32 @@ struct BaseMMAFrag<T, 8, 8> {
     }
   }
 
+  template <typename Atype, typename Btype, typename Ctype>
   METAL_FUNC static constexpr void mma(
       thread frag_type& D,
-      thread frag_type& A,
-      thread frag_type& B,
-      thread frag_type& C) {
+      thread dtype_frag_t<Atype>& A,
+      thread dtype_frag_t<Btype>& B,
+      thread dtype_frag_t<Ctype>& C) {
     mat_type D_mat;
-    mat_type A_mat;
-    mat_type B_mat;
-    mat_type C_mat;
+    dtype_mat_t<Atype> A_mat;
+    dtype_mat_t<Btype> B_mat;
+    dtype_mat_t<Ctype> C_mat;
 
-    reinterpret_cast<thread frag_type&>(A_mat.thread_elements()) = A;
-    reinterpret_cast<thread frag_type&>(B_mat.thread_elements()) = B;
-    reinterpret_cast<thread frag_type&>(C_mat.thread_elements()) = C;
+    reinterpret_cast<thread dtype_frag_t<Atype>&>(A_mat.thread_elements()) = A;
+    reinterpret_cast<thread dtype_frag_t<Btype>&>(B_mat.thread_elements()) = B;
+    reinterpret_cast<thread dtype_frag_t<Ctype>&>(C_mat.thread_elements()) = C;
 
     mma(D_mat, A_mat, B_mat, C_mat);
 
     D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
   }
 
+  template <typename Atype, typename Btype, typename Ctype>
   METAL_FUNC static constexpr void mma(
       thread mat_type& D,
-      thread mat_type& A,
-      thread mat_type& B,
-      thread mat_type& C) {
+      thread dtype_mat_t<Atype>& A,
+      thread dtype_mat_t<Btype>& B,
+      thread dtype_mat_t<Ctype>& C) {
     simdgroup_multiply_accumulate(D, A, B, C);
   }
 
@@ -242,7 +250,7 @@ struct MMATile {
   typedef typename MMAFrag_t::mat_type mat_type;
   typedef typename MMAFrag_t::frag_type frag_type;
 
-  frag_type val_frags[kNumFrags] = {frag_type(0)};
+  frag_type val_frags[kNumFrags]; // = {frag_type(0)};
 
   METAL_FUNC MMATile() thread {}
 
@@ -409,24 +417,37 @@ struct MMATile {
   }
 };
 
-template <typename T, typename U, int M, int N, int K>
+template <
+    typename Dtype,
+    typename Atype,
+    typename Btype,
+    typename Ctype,
+    int M,
+    int N,
+    int K,
+    class MMAFragD,
+    class MMAFragA,
+    class MMAFragB,
+    class MMAFragC>
 METAL_FUNC void tile_matmad(
-    thread MMATile<T, M, N>& D,
-    thread MMATile<U, M, K>& A,
-    thread MMATile<U, K, N>& B,
-    thread MMATile<T, M, N>& C) {
+    thread MMATile<Dtype, M, N, MMAFragD>& D,
+    thread MMATile<Atype, M, K, MMAFragA>& A,
+    thread MMATile<Btype, K, N, MMAFragB>& B,
+    thread MMATile<Ctype, M, N, MMAFragC>& C) {
   STEEL_PRAGMA_UNROLL
-  for (short k = 0; k < K; ++k) {
+  for (short m = 0; m < M; ++m) {
     STEEL_PRAGMA_UNROLL
-    for (short m = 0; m < M; ++m) {
+    for (short n = 0; n < N; ++n) {
+      short m_serp = m; //(n % 2) ? (M - 1 - m) : m;
+      short n_serp = (m % 2) ? (N - 1 - n) : n;
+
       STEEL_PRAGMA_UNROLL
-      for (short n = 0; n < N; ++n) {
-        short n_serp = (m % 2) ? (N - 1 - n) : n;
-        MMATile<T, M, N>::MMAFrag_t::mma(
-            D.frag_at(m, n_serp),
-            A.frag_at(m, k),
+      for (short k = 0; k < K; ++k) {
+        MMAFragD::mma(
+            D.frag_at(m_serp, n_serp),
+            A.frag_at(m_serp, k),
             B.frag_at(k, n_serp),
-            C.frag_at(m, n_serp));
+            C.frag_at(m_serp, n_serp));
       }
     }
   }

mlx/backend/metal/metal.cpp

@@ -2,7 +2,6 @@
 #include <memory>
 
 #include "mlx/backend/metal/device.h"
-#include "mlx/backend/metal/event.h"
 #include "mlx/backend/metal/utils.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"
@@ -23,88 +22,78 @@ inline void check_error(MTL::CommandBuffer* cbuf) {
   }
 }
 
-std::function<void()> make_task(array arr, bool signal) {
-  auto task = [arr = std::move(arr), signal]() mutable {
-    auto pool = new_scoped_memory_pool();
-    auto s = arr.primitive().stream();
-    auto& d = metal::device(s.device);
-    auto command_buffer = d.get_command_buffer(s.index);
+void eval(array& arr) {
+  auto pool = new_scoped_memory_pool();
+  auto s = arr.primitive().stream();
+  auto& d = metal::device(s.device);
+  auto command_buffer = d.get_command_buffer(s.index);
 
-    for (auto& input : arr.inputs()) {
-      if (input.event().valid() &&
-          input.event().stream() != arr.primitive().stream()) {
-        input.event().wait();
-      }
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
     }
 
-    auto outputs = arr.outputs();
-    {
-      // If the array is a tracer hold a reference
-      // to its inputs so they don't get donated
-      std::vector<array> inputs;
-      if (arr.is_tracer()) {
-        inputs = arr.inputs();
-      }
+    debug_set_primitive_buffer_label(command_buffer, arr.primitive());
+    arr.primitive().eval_gpu(arr.inputs(), outputs);
+  }
+  std::unordered_set<std::shared_ptr<array::Data>> buffers;
+  for (auto& in : arr.inputs()) {
+    buffers.insert(in.data_shared_ptr());
+  }
+  for (auto& s : arr.siblings()) {
+    buffers.insert(s.data_shared_ptr());
+  }
+  // Remove the output if it was donated to by an input
+  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
+    buffers.erase(it);
+  }
 
-      debug_set_primitive_buffer_label(command_buffer, arr.primitive());
-      try {
-        arr.primitive().eval_gpu(arr.inputs(), outputs);
-      } catch (const std::exception& error) {
-        abort_with_exception(error);
-      }
-    }
-    std::vector<std::shared_ptr<array::Data>> buffers;
-    for (auto& in : arr.inputs()) {
-      buffers.push_back(in.data_shared_ptr());
-    }
-    for (auto& s : arr.siblings()) {
-      buffers.push_back(s.data_shared_ptr());
-    }
-    if (!arr.is_tracer()) {
-      arr.detach();
-    }
-    for (auto& out : outputs) {
-      out.set_status(array::Status::evaluated);
-    }
-
-    if (signal || d.command_buffer_needs_commit(s.index)) {
-      if (signal) {
-        encode_signal(arr.event());
-      }
-      d.end_encoding(s.index);
-      scheduler::notify_new_task(s);
-      command_buffer->addCompletedHandler(
-          [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
-            scheduler::notify_task_completion(s);
-            check_error(cbuf);
-          });
-      d.commit_command_buffer(s.index);
-      d.get_command_buffer(s.index);
-    } else {
-      command_buffer->addCompletedHandler(
-          [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
-            check_error(cbuf);
-          });
-    }
-  };
-  return task;
+  if (d.command_buffer_needs_commit(s.index)) {
+    d.end_encoding(s.index);
+    scheduler::notify_new_task(s);
+    command_buffer->addCompletedHandler(
+        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          scheduler::notify_task_completion(s);
+          check_error(cbuf);
+        });
+    d.commit_command_buffer(s.index);
+    d.get_command_buffer(s.index);
+  } else {
+    command_buffer->addCompletedHandler(
+        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          check_error(cbuf);
+        });
+  }
 }
 
-std::function<void()> make_synchronize_task(
-    Stream s,
-    std::shared_ptr<std::promise<void>> p) {
-  return [s, p = std::move(p)]() {
-    auto pool = new_scoped_memory_pool();
-    auto& d = metal::device(s.device);
-    auto cb = d.get_command_buffer(s.index);
-    cb->retain();
-    d.end_encoding(s.index);
-    d.commit_command_buffer(s.index);
-    cb->waitUntilCompleted();
-    check_error(cb);
-    cb->release();
-    p->set_value();
-  };
+void finalize(Stream s) {
+  auto pool = new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  d.end_encoding(s.index);
+  scheduler::notify_new_task(s);
+  cb->addCompletedHandler([s](MTL::CommandBuffer* cbuf) {
+    scheduler::notify_task_completion(s);
+    check_error(cbuf);
+  });
+  d.commit_command_buffer(s.index);
+  d.get_command_buffer(s.index);
+}
+
+void synchronize(Stream s) {
+  auto pool = new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  cb->retain();
+  d.end_encoding(s.index);
+  d.commit_command_buffer(s.index);
+  cb->waitUntilCompleted();
+  check_error(cb);
+  cb->release();
 }
 
 void start_capture(std::string path, id object) {

mlx/backend/metal/metal.h

@@ -82,7 +82,7 @@ void start_capture(std::string path = "");
 void stop_capture();
 
 /** Get information about the GPU and system settings. */
-std::unordered_map<std::string, std::variant<std::string, size_t>>
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
 device_info();
 
 } // namespace mlx::core::metal

mlx/backend/metal/metal_impl.h

@@ -14,10 +14,8 @@ void new_stream(Stream stream);
 
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
-std::function<void()> make_task(array arr, bool signal);
-
-std::function<void()> make_synchronize_task(
-    Stream s,
-    std::shared_ptr<std::promise<void>> p);
+void eval(array& arr);
+void finalize(Stream s);
+void synchronize(Stream s);
 
 } // namespace mlx::core::metal

mlx/backend/metal/normalization.cpp

@@ -18,33 +18,33 @@ void RMSNorm::eval_gpu(
   auto& out = outputs[0];
 
   // Make sure that the last dimension is contiguous
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto set_output = [&s, &out](const array& x) {
     bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
     if (no_copy && x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
       no_copy &= (s == 0 || s == x.shape().back());
     }
     if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
       return x;
     } else {
-      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-      copy_gpu(x, copies.back(), CopyType::General, s);
-      return copies.back();
+      auto x_copy = array(x.shape(), x.dtype(), nullptr, {});
+      copy_gpu(x, x_copy, CopyType::General, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
     }
   };
-  const array& x = check_input(inputs[0]);
-  const array& w = inputs[1];
 
-  if (x.is_donatable()) {
-    out.move_shared_buffer(x);
-  } else {
-    out.set_data(
-        allocator::malloc_or_wait(x.data_size() * x.itemsize()),
-        x.data_size(),
-        x.strides(),
-        x.flags());
-  }
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
 
   auto axis_size = static_cast<uint32_t>(x.shape().back());
   int n_rows = x.data_size() / axis_size;
@@ -77,7 +77,7 @@ void RMSNorm::eval_gpu(
       group_dims = MTL::Size(threadgroup_size, 1, 1);
     }
 
-    uint32_t w_stride = w.strides()[0];
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
     compute_encoder.set_compute_pipeline_state(kernel);
     compute_encoder.set_input_array(
         x.data_shared_ptr() == nullptr ? out : x, 0);
@@ -88,8 +88,6 @@ void RMSNorm::eval_gpu(
     compute_encoder.set_bytes(w_stride, 5);
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 void RMSNormVJP::eval_gpu(
@@ -101,54 +99,57 @@ void RMSNormVJP::eval_gpu(
   // Ensure row contiguity. We could relax this step by checking that the array
   // is contiguous (no broadcasts or holes) and that the input strides are the
   // same as the cotangent strides but for now this is simpler.
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto check_input = [&d, &s](const array& x) -> std::pair<array, bool> {
     if (x.flags().row_contiguous) {
-      return x;
+      return {x, false};
     }
-    // Make sure we 'll only ever allocate once. The point of that goes beyond
-    // the minor optimization. We need to ensure that there will be no
-    // reallocation such that the references won't change when we
-    // push_back(...). So tl;dr 3 possible copies x, g and gw_temp.
-    copies.reserve(3);
 
-    copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-    copy_gpu(x, copies.back(), CopyType::General, s);
-    return copies.back();
+    array x_copy(x.shape(), x.dtype(), nullptr, {});
+    copy_gpu(x, x_copy, CopyType::General, s);
+    return {x_copy, true};
   };
-  const array& x = check_input(inputs[0]);
+  bool donate_x = inputs[0].is_donatable();
+  bool donate_g = inputs[2].is_donatable();
+  auto [x, copied] = check_input(inputs[0]);
+  donate_x |= copied;
   const array& w = inputs[1];
-  const array& g = check_input(inputs[2]);
+  auto [g, g_copied] = check_input(inputs[2]);
+  donate_g |= g_copied;
   array& gx = outputs[0];
   array& gw = outputs[1];
 
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
   // Allocate space for the outputs
-  bool x_in_gx = false;
   bool g_in_gx = false;
   if (x.is_donatable()) {
-    gx.move_shared_buffer(x);
-    x_in_gx = true;
+    gx.copy_shared_buffer(x);
   } else if (g.is_donatable()) {
-    gx.move_shared_buffer(g);
+    gx.copy_shared_buffer(g);
     g_in_gx = true;
   } else {
     gx.set_data(allocator::malloc_or_wait(gx.nbytes()));
   }
+  if (g_copied && !g_in_gx) {
+    d.add_temporary(g, s.index);
+  }
 
   auto axis_size = static_cast<uint32_t>(x.shape().back());
   int n_rows = x.data_size() / axis_size;
 
   // Allocate the gradient accumulator gw and a temporary to store the
   // gradients before they are accumulated.
-  array gw_temp({n_rows, x.shape().back()}, gw.dtype(), nullptr, {});
-  bool g_in_gw = false;
-  if (!g_in_gx && g.is_donatable()) {
-    gw_temp.move_shared_buffer(g);
-    g_in_gw = true;
-  } else {
-    gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+      d.add_temporary(gw_temp, s.index);
+    }
   }
-  copies.push_back(gw_temp);
   gw.set_data(allocator::malloc_or_wait(gw.nbytes()));
 
   const int simd_size = 32;
@@ -159,9 +160,15 @@ void RMSNormVJP::eval_gpu(
     op_name += "_looped";
   }
   op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
   auto& compute_encoder = d.get_command_encoder(s.index);
   {
-    auto kernel = d.get_kernel(op_name);
+    auto kernel = d.get_kernel(op_name, "mlx", hash_name, func_consts);
 
     MTL::Size grid_dims, group_dims;
     if (axis_size <= looped_limit) {
@@ -179,11 +186,11 @@ void RMSNormVJP::eval_gpu(
       group_dims = MTL::Size(threadgroup_size, 1, 1);
     }
 
-    uint32_t w_stride = w.strides()[0];
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
     compute_encoder.set_compute_pipeline_state(kernel);
-    compute_encoder.set_input_array(x_in_gx ? gx : x, 0);
+    compute_encoder.set_input_array(x, 0);
     compute_encoder.set_input_array(w, 1);
-    compute_encoder.set_input_array(g_in_gx ? gx : (g_in_gw ? gw_temp : g), 2);
+    compute_encoder.set_input_array(g, 2);
     compute_encoder.set_output_array(gx, 3);
     compute_encoder.set_output_array(gw_temp, 4);
     compute_encoder.set_bytes(eps_, 5);
@@ -192,12 +199,12 @@ void RMSNormVJP::eval_gpu(
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
 
-  ReductionPlan plan(
-      ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
-  strided_reduce_general_dispatch(
-      gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
-
-  d.add_temporaries(std::move(copies), s.index);
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    strided_reduce_general_dispatch(
+        gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
+  }
 }
 
 void LayerNorm::eval_gpu(
@@ -208,35 +215,35 @@ void LayerNorm::eval_gpu(
   auto& out = outputs[0];
 
   // Make sure that the last dimension is contiguous
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto set_output = [&s, &out](const array& x) {
     bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
     if (no_copy && x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
       no_copy &= (s == 0 || s == x.shape().back());
     }
     if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
       return x;
     } else {
-      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-      copy_gpu(x, copies.back(), CopyType::General, s);
-      return copies.back();
+      auto x_copy = array(x.shape(), x.dtype(), nullptr, {});
+      copy_gpu(x, x_copy, CopyType::General, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
     }
   };
-  const array& x = check_input(inputs[0]);
+
+  const array x = set_output(inputs[0]);
   const array& w = inputs[1];
   const array& b = inputs[2];
 
-  if (x.is_donatable()) {
-    out.move_shared_buffer(x);
-  } else {
-    out.set_data(
-        allocator::malloc_or_wait(x.data_size() * x.itemsize()),
-        x.data_size(),
-        x.strides(),
-        x.flags());
-  }
-
   auto axis_size = static_cast<uint32_t>(x.shape().back());
   int n_rows = x.data_size() / axis_size;
 
@@ -282,8 +289,6 @@ void LayerNorm::eval_gpu(
     compute_encoder.set_bytes(b_stride, 7);
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 void LayerNormVJP::eval_gpu(
@@ -295,56 +300,58 @@ void LayerNormVJP::eval_gpu(
   // Ensure row contiguity. We could relax this step by checking that the array
   // is contiguous (no broadcasts or holes) and that the input strides are the
   // same as the cotangent strides but for now this is simpler.
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto check_input = [&s](const array& x) -> std::pair<array, bool> {
     if (x.flags().row_contiguous) {
-      return x;
+      return {x, false};
     }
-    // Make sure we 'll only ever allocate once. The point of that goes beyond
-    // the minor optimization. We need to ensure that there will be no
-    // reallocation such that the references won't change when we
-    // push_back(...). So tl;dr 3 possible copies x, g and gw_temp.
-    copies.reserve(3);
-
-    copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-    copy_gpu(x, copies.back(), CopyType::General, s);
-    return copies.back();
+    array x_copy(x.shape(), x.dtype(), nullptr, {});
+    copy_gpu(x, x_copy, CopyType::General, s);
+    return {x_copy, true};
   };
-  const array& x = check_input(inputs[0]);
+  bool donate_x = inputs[0].is_donatable();
+  bool donate_g = inputs[3].is_donatable();
+  auto [x, copied] = check_input(inputs[0]);
+  donate_x |= copied;
   const array& w = inputs[1];
   const array& b = inputs[2];
-  const array& g = check_input(inputs[3]);
+  auto [g, g_copied] = check_input(inputs[3]);
+  donate_g |= g_copied;
   array& gx = outputs[0];
   array& gw = outputs[1];
   array& gb = outputs[2];
 
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
   // Allocate space for the outputs
-  bool x_in_gx = false;
   bool g_in_gx = false;
-  if (x.is_donatable()) {
-    gx.move_shared_buffer(x);
-    x_in_gx = true;
-  } else if (g.is_donatable()) {
-    gx.move_shared_buffer(g);
+  if (donate_x) {
+    gx.copy_shared_buffer(x);
+  } else if (donate_g) {
+    gx.copy_shared_buffer(g);
     g_in_gx = true;
   } else {
     gx.set_data(allocator::malloc_or_wait(gx.nbytes()));
   }
+  if (g_copied && !g_in_gx) {
+    d.add_temporary(g, s.index);
+  }
 
   auto axis_size = static_cast<uint32_t>(x.shape().back());
   int n_rows = x.data_size() / axis_size;
 
   // Allocate a temporary to store the gradients for w and allocate the output
   // gradient accumulators.
-  array gw_temp({n_rows, x.shape().back()}, gw.dtype(), nullptr, {});
-  bool g_in_gw = false;
-  if (!g_in_gx && g.is_donatable()) {
-    gw_temp.move_shared_buffer(g);
-    g_in_gw = true;
-  } else {
-    gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+      d.add_temporary(gw_temp, s.index);
+    }
   }
-  copies.push_back(gw_temp);
   gw.set_data(allocator::malloc_or_wait(gw.nbytes()));
   gb.set_data(allocator::malloc_or_wait(gb.nbytes()));
 
@@ -354,14 +361,7 @@ void LayerNormVJP::eval_gpu(
     ReductionPlan plan(
         ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
     strided_reduce_general_dispatch(
-        g_in_gx ? gx : (g_in_gw ? gw_temp : g),
-        gb,
-        "sum",
-        plan,
-        {0},
-        compute_encoder,
-        d,
-        s);
+        g, gb, "sum", plan, {0}, compute_encoder, d, s);
   }
 
   const int simd_size = 32;
@@ -372,8 +372,14 @@ void LayerNormVJP::eval_gpu(
     op_name += "_looped";
   }
   op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
   {
-    auto kernel = d.get_kernel(op_name);
+    auto kernel = d.get_kernel(op_name, "mlx", hash_name, func_consts);
 
     MTL::Size grid_dims, group_dims;
     if (axis_size <= looped_limit) {
@@ -393,9 +399,9 @@ void LayerNormVJP::eval_gpu(
 
     uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
     compute_encoder.set_compute_pipeline_state(kernel);
-    compute_encoder.set_input_array(x_in_gx ? gx : x, 0);
+    compute_encoder.set_input_array(x, 0);
     compute_encoder.set_input_array(w, 1);
-    compute_encoder.set_input_array(g_in_gx ? gx : (g_in_gw ? gw_temp : g), 2);
+    compute_encoder.set_input_array(g, 2);
     compute_encoder.set_output_array(gx, 3);
     compute_encoder.set_output_array(gw_temp, 4);
     compute_encoder.set_bytes(eps_, 5);
@@ -404,14 +410,12 @@ void LayerNormVJP::eval_gpu(
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
 
-  if (gw.ndim() == 1 && gw.size() == axis_size) {
+  if (has_w) {
     ReductionPlan plan(
         ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
     strided_reduce_general_dispatch(
         gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 } // namespace mlx::core::fast

mlx/backend/metal/primitives.cpp

@@ -5,12 +5,10 @@
 #include <sstream>
 
 #include "mlx/backend/common/compiled.h"
-#include "mlx/backend/common/load.h"
 #include "mlx/backend/common/slicing.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/device.h"
-#include "mlx/backend/metal/event.h"
 #include "mlx/backend/metal/kernels.h"
 #include "mlx/backend/metal/slicing.h"
 #include "mlx/backend/metal/utils.h"
@@ -45,7 +43,8 @@ void reshape(const array& in, array& out, Stream s) {
     shared_buffer_reshape(in, out_strides, out);
   }
 }
-array compute_dynamic_offset(
+
+static array compute_dynamic_offset(
     const array& indices,
     const Strides& strides,
     const std::vector<int>& axes,
@@ -57,7 +56,7 @@ array compute_dynamic_offset(
   bool donate = indices.is_donatable() &&
       (indices.data_size() * indices.itemsize()) >= offset.itemsize();
   if (donate) {
-    offset.move_shared_buffer(indices);
+    offset.copy_shared_buffer(indices);
   } else {
     offset.set_data(allocator::malloc_or_wait(offset.itemsize()));
   }
@@ -88,7 +87,7 @@ array compute_dynamic_offset(
 
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
-  compute_encoder.set_input_array(donate ? offset : indices, 0);
+  compute_encoder.set_input_array(indices, 0);
   compute_encoder.set_output_array(offset, 1);
   compute_encoder.set_vector_bytes(strides, 2);
   compute_encoder.set_vector_bytes(axes, 3);
@@ -254,7 +253,7 @@ void Contiguous::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
   if (in.flags().row_contiguous ||
       (allow_col_major_ && in.flags().col_contiguous)) {
-    move_or_copy(in, out);
+    out.copy_shared_buffer(in);
   } else {
     copy_gpu(in, out, CopyType::General);
   }
@@ -302,31 +301,7 @@ void Unflatten::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Load::eval_gpu(const std::vector<array>& inputs, array& out) {
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  auto read_task = [out = unsafe_weak_copy(out),
-                    offset = offset_,
-                    reader = reader_,
-                    swap_endianness = swap_endianness_]() mutable {
-    load(out, offset, reader, swap_endianness);
-  };
-
-  // Limit the size that the command buffer will wait on to avoid timing out
-  // on the event (<4 seconds).
-  if (out.nbytes() > (1 << 28)) {
-    read_task();
-    return;
-  }
-
-  auto fut = io::thread_pool().enqueue(std::move(read_task)).share();
-
-  auto e = Event(stream());
-  e.set_value(1);
-  encode_wait(e);
-  auto signal_task = [e = std::move(e), fut = std::move(fut)]() mutable {
-    fut.wait();
-    e.signal();
-  };
-  scheduler::enqueue(io_stream(), std::move(signal_task));
+  throw std::runtime_error("[Load::eval_gpu] Not implemented.");
 }
 
 void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -452,7 +427,7 @@ void DynamicSliceUpdate::eval_gpu(
   auto& start_indices = inputs[2];
 
   if (upd.size() == 0) {
-    move_or_copy(in, out);
+    out.copy_shared_buffer(in);
     return;
   }
 
@@ -491,7 +466,7 @@ void SliceUpdate::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& upd = inputs[1];
 
   if (upd.size() == 0) {
-    move_or_copy(in, out);
+    out.copy_shared_buffer(in);
     return;
   }
 
@@ -575,8 +550,8 @@ void View::eval_gpu(const std::vector<array>& inputs, array& out) {
       strides[i] *= ibytes;
       strides[i] /= obytes;
     }
-    move_or_copy(
-        in, out, strides, in.flags(), in.data_size() * ibytes / obytes);
+    out.copy_shared_buffer(
+        in, strides, in.flags(), in.data_size() * ibytes / obytes);
   } else {
     auto tmp = array(in.shape(), in.dtype(), nullptr, {});
     tmp.set_data(allocator::malloc_or_wait(tmp.nbytes()));
@@ -587,7 +562,7 @@ void View::eval_gpu(const std::vector<array>& inputs, array& out) {
     flags.row_contiguous = true;
     auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
     flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
-    out.move_shared_buffer(tmp, out.strides(), flags, out.size());
+    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
   }
 }
 

mlx/backend/metal/rope.cpp

@@ -25,6 +25,10 @@ void RoPE::eval_gpu(
   size_t out_strides[3];
   bool donated = false;
   int ndim = in.ndim();
+  int dispatch_ndim = in.ndim();
+  while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) {
+    dispatch_ndim--;
+  }
   size_t mat_size = in.shape(-2) * in.shape(-1);
   if (dims_ < in.shape(-1)) {
     donated = true;
@@ -37,19 +41,19 @@ void RoPE::eval_gpu(
   } else if (in.flags().row_contiguous) {
     if (in.is_donatable()) {
       donated = true;
-      out.move_shared_buffer(in);
+      out.copy_shared_buffer(in);
     } else {
       out.set_data(allocator::malloc_or_wait(out.nbytes()));
     }
     strides[0] = mat_size;
     strides[1] = in.strides()[ndim - 2];
     strides[2] = in.strides()[ndim - 1];
-  } else if (ndim == 3) {
+  } else if (dispatch_ndim == 3) {
     // Handle non-contiguous 3D inputs
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    strides[0] = in.strides()[0];
-    strides[1] = in.strides()[1];
-    strides[2] = in.strides()[2];
+    strides[0] = in.strides()[ndim - 3];
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
   } else {
     // Copy non-contiguous > 3D inputs into the output and treat
     // input as donated

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -134,14 +134,17 @@ void sdpa_vector(
   size_t k_stride = k.strides()[1];
   size_t v_stride = v.strides()[1];
   MTL::Size group_dims(1024, 1, 1);
-  MTL::Size grid_dims(1, B, 1);
+  MTL::Size grid_dims(B, q.shape(2), 1);
 
   bool has_mask = mask.has_value();
+  bool query_transposed = !q.flags().row_contiguous;
   metal::MTLFCList func_consts = {
       {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
   };
   std::string hash_name = kname;
   hash_name += has_mask ? "_mask" : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
 
   // Get the kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
@@ -149,7 +152,7 @@ void sdpa_vector(
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Set its arguments
-  compute_encoder.set_input_array(q.data_shared_ptr() == nullptr ? out : q, 0);
+  compute_encoder.set_input_array(q, 0);
   compute_encoder.set_input_array(k, 1);
   compute_encoder.set_input_array(v, 2);
   compute_encoder.set_output_array(out, 3);
@@ -161,10 +164,14 @@ void sdpa_vector(
   if (has_mask) {
     auto& m = *mask;
     compute_encoder.set_input_array(m, 9);
-    int32_t seq_stride = m.ndim() >= 1 ? m.strides().back() : 0;
-    int32_t head_stride = m.ndim() >= 3 ? *(m.strides().end() - 3) : 0;
-    compute_encoder.set_bytes(seq_stride, 10);
-    compute_encoder.set_bytes(head_stride, 11);
+    auto nd = m.ndim();
+    int32_t kv_seq_stride =
+        nd >= 1 && m.shape(-1) > 1 ? m.strides()[nd - 1] : 0;
+    int32_t q_seq_stride = nd >= 2 && m.shape(-2) > 1 ? m.strides()[nd - 2] : 0;
+    int32_t head_stride = nd >= 3 && m.shape(-3) > 1 ? m.strides()[nd - 3] : 0;
+    compute_encoder.set_bytes(kv_seq_stride, 10);
+    compute_encoder.set_bytes(q_seq_stride, 11);
+    compute_encoder.set_bytes(head_stride, 12);
   }
 
   // Launch
@@ -198,7 +205,7 @@ void sdpa_vector_2pass(
   auto k_stride = k.strides()[1];
   auto v_stride = v.strides()[1];
   MTL::Size group_dims(8 * 32, 1, 1);
-  MTL::Size grid_dims(1, B, blocks);
+  MTL::Size grid_dims(B, q.shape(2), blocks);
 
   // Allocate the intermediates
   Shape intermediate_shape;
@@ -219,11 +226,14 @@ void sdpa_vector_2pass(
   d.add_temporary(maxs, s.index);
 
   bool has_mask = mask.has_value();
+  bool query_transposed = !q.flags().row_contiguous;
   metal::MTLFCList func_consts = {
       {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
   };
   std::string hash_name = kname;
   hash_name += has_mask ? "_mask" : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
 
   // Get the kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
@@ -232,7 +242,7 @@ void sdpa_vector_2pass(
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Set its arguments
-  compute_encoder.set_input_array(q.data_shared_ptr() == nullptr ? out : q, 0);
+  compute_encoder.set_input_array(q, 0);
   compute_encoder.set_input_array(k, 1);
   compute_encoder.set_input_array(v, 2);
   compute_encoder.set_output_array(intermediate, 3);
@@ -246,10 +256,14 @@ void sdpa_vector_2pass(
   if (has_mask) {
     auto& m = *mask;
     compute_encoder.set_input_array(m, 11);
-    int32_t seq_stride = m.ndim() >= 1 ? m.strides().back() : 0;
-    int32_t head_stride = m.ndim() >= 3 ? *(m.strides().end() - 3) : 0;
-    compute_encoder.set_bytes(seq_stride, 12);
-    compute_encoder.set_bytes(head_stride, 13);
+    auto nd = m.ndim();
+    int32_t kv_seq_stride =
+        nd >= 1 && m.shape(-1) > 1 ? m.strides()[nd - 1] : 0;
+    int32_t q_seq_stride = nd >= 2 && m.shape(-2) > 1 ? m.strides()[nd - 2] : 0;
+    int32_t head_stride = nd >= 3 && m.shape(-3) > 1 ? m.strides()[nd - 3] : 0;
+    compute_encoder.set_bytes(kv_seq_stride, 12);
+    compute_encoder.set_bytes(q_seq_stride, 13);
+    compute_encoder.set_bytes(head_stride, 14);
   }
 
   // Launch
@@ -274,7 +288,7 @@ void sdpa_vector_2pass(
 
   // Launch
   group_dims = MTL::Size(1024, 1, 1);
-  grid_dims = MTL::Size(1, B, 1);
+  grid_dims = MTL::Size(B, q.shape(2), 1);
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
@@ -301,16 +315,23 @@ void ScaledDotProductAttention::eval_gpu(
     if (!predicate(arr)) {
       array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
       copy_gpu(arr, arr_copy, CopyType::General, s);
-      copies.push_back(arr_copy);
+      copies.push_back(std::move(arr_copy));
       return copies.back();
     } else {
       return arr;
     }
   };
 
-  // Checks if arr is fully row contiguous
-  auto is_contiguous = [](const array& arr) {
-    return arr.flags().row_contiguous;
+  // Checks if arr is row contiguous or the sequence and head dimension are
+  // transposed
+  auto is_contiguous_or_head_seq_transposed = [](const array& arr) {
+    if (arr.flags().row_contiguous) {
+      return true;
+    }
+    auto& strides = arr.strides();
+    auto& shape = arr.shape();
+    return (strides[3] == 1) && (strides[2] == shape[3] * shape[1]) &&
+        (strides[1] == shape[3]) && (strides[0] == strides[2] * shape[2]);
   };
 
   // Returns true if the array is row contiguous except the sequence length
@@ -328,18 +349,30 @@ void ScaledDotProductAttention::eval_gpu(
   };
 
   // We are in vector mode ie single query
-  if (q_pre.shape(2) == 1) {
-    const auto& q = copy_unless(is_contiguous, q_pre);
-    // 1, heads, seq_len, head_dim
-    // mask [1, query_heads, 1, seq_len]
+  if (q_pre.shape(2) <= 8) {
+    const auto& q = copy_unless(is_contiguous_or_head_seq_transposed, q_pre);
     const auto& k = copy_unless(is_contiguous_except_seq_len, k_pre);
     const auto& v = copy_unless(is_contiguous_except_seq_len, v_pre);
 
     // Donate the query if possible
-    if (q.is_donatable() && q.size() == o.size()) {
-      o.move_shared_buffer(q);
+    if (q.is_donatable() && (q.shape(2) == 1 || !q.flags().row_contiguous) &&
+        q.size() == o.size()) {
+      o.copy_shared_buffer(q);
     } else {
-      o.set_data(allocator::malloc_or_wait(o.nbytes()));
+      if (o.shape(2) == 1) {
+        o.set_data(allocator::malloc_or_wait(o.nbytes()));
+      } else {
+        auto strides = o.strides();
+        strides[2] = o.shape(1) * o.shape(3);
+        strides[1] = o.shape(3);
+        auto flags = q.flags();
+        flags.row_contiguous = q.shape(1) == 1;
+        o.set_data(
+            allocator::malloc_or_wait(o.nbytes()),
+            o.size(),
+            std::move(strides),
+            flags);
+      }
     }
 
     auto mask =

mlx/backend/metal/scan.cpp

@@ -17,11 +17,11 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  std::vector<array> copies;
+  bool donate = inputs[0].is_donatable();
   auto in = inputs[0];
   if (in.flags().contiguous && in.strides()[axis_] != 0) {
-    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
-      out.move_shared_buffer(in);
+    if (donate && in.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(in);
     } else {
       out.set_data(
           allocator::malloc_or_wait(in.data_size() * out.itemsize()),
@@ -32,9 +32,8 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   } else {
     array arr_copy(in.shape(), in.dtype(), nullptr, {});
     copy_gpu(in, arr_copy, CopyType::General, s);
-    copies.push_back(arr_copy);
-    in = arr_copy;
-    out.move_shared_buffer(in);
+    in = std::move(arr_copy);
+    out.copy_shared_buffer(in);
   }
 
   bool contiguous = in.strides()[axis_] == 1;
@@ -69,8 +68,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   if (contiguous) {
     auto& compute_encoder = d.get_command_encoder(s.index);
     compute_encoder.set_compute_pipeline_state(kernel);
-    compute_encoder.set_input_array(
-        in.data_shared_ptr() == nullptr ? out : in, 0);
+    compute_encoder.set_input_array(in, 0);
     compute_encoder.set_output_array(out, 1);
     size_t size = in.shape(axis_);
     compute_encoder.set_bytes(size, 2);
@@ -127,8 +125,6 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
     MTL::Size group_dims(thread_group_size, 1, 1);
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 } // namespace mlx::core

mlx/backend/metal/softmax.cpp

@@ -22,31 +22,32 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& d = metal::device(s.device);
 
   // Make sure that the last dimension is contiguous
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto set_output = [&s, &out](const array& x) {
     bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
     if (no_copy && x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
       no_copy &= (s == 0 || s == x.shape().back());
     }
     if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
       return x;
     } else {
-      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-      copy_gpu(x, copies.back(), CopyType::General, s);
-      return copies.back();
+      auto x_copy = array(x.shape(), x.dtype(), nullptr, {});
+      copy_gpu(x, x_copy, CopyType::General, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
     }
   };
-  const array& in = check_input(inputs[0]);
-  if (in.is_donatable()) {
-    out.move_shared_buffer(in);
-  } else {
-    out.set_data(
-        allocator::malloc_or_wait(in.data_size() * in.itemsize()),
-        in.data_size(),
-        in.strides(),
-        in.flags());
-  }
+
+  const array in = set_output(inputs[0]);
 
   int axis_size = in.shape().back();
   int n_rows = in.data_size() / axis_size;
@@ -82,14 +83,11 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
     }
 
     compute_encoder.set_compute_pipeline_state(kernel);
-    compute_encoder.set_input_array(
-        in.data_shared_ptr() == nullptr ? out : in, 0);
+    compute_encoder.set_input_array(in, 0);
     compute_encoder.set_output_array(out, 1);
     compute_encoder.set_bytes(axis_size, 2);
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 } // namespace mlx::core

mlx/backend/metal/ternary.cpp

@@ -71,12 +71,9 @@ void ternary_op_gpu_inplace(
 
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
-  bool donate_a = a.data_shared_ptr() == nullptr;
-  bool donate_b = b.data_shared_ptr() == nullptr;
-  bool donate_c = c.data_shared_ptr() == nullptr;
-  compute_encoder.set_input_array(donate_a ? out : a, 0);
-  compute_encoder.set_input_array(donate_b ? out : b, 1);
-  compute_encoder.set_input_array(donate_c ? out : c, 2);
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(c, 2);
   compute_encoder.set_output_array(out, 3);
 
   auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
@@ -129,7 +126,7 @@ void ternary_op_gpu(
   auto& b = inputs[1];
   auto& c = inputs[2];
   TernaryOpType topt = get_ternary_op_type(a, b, c);
-  set_ternary_op_output_data(a, b, c, out, topt, true /* donate_with_move */);
+  set_ternary_op_output_data(a, b, c, out, topt);
   ternary_op_gpu_inplace(inputs, out, op, s);
 }
 

mlx/backend/metal/unary.cpp

@@ -57,8 +57,7 @@ void unary_op_gpu_inplace(
   auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
-  compute_encoder.set_input_array(
-      in.data_shared_ptr() == nullptr ? out : in, 0);
+  compute_encoder.set_input_array(in, 0);
   compute_encoder.set_output_array(out, 1);
   if (!contig) {
     // Launch up to 3D grid of threads
@@ -95,7 +94,7 @@ void unary_op_gpu(
   bool contig = in.flags().contiguous;
   if (contig) {
     if (in.is_donatable() && in.itemsize() == out.itemsize()) {
-      out.move_shared_buffer(in);
+      out.copy_shared_buffer(in);
     } else {
       out.set_data(
           allocator::malloc_or_wait(in.data_size() * out.itemsize()),
@@ -169,7 +168,7 @@ void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
     unary_op_gpu(inputs, out, get_primitive_string(this));
   } else {
     // No-op integer types
-    move_or_copy(in, out);
+    out.copy_shared_buffer(in);
   }
 }