CMakeLists.txt update

* optionally load metallib from framework

* pre-commit

* adjust logic

.circleci/config.yml

@@ -26,9 +26,9 @@ jobs:
           name: Install
           command: |
             xcodebuild -downloadComponent MetalToolchain
-            brew install python@3.9
+            brew install python@3.10
             brew install doxygen
-            python3.9 -m venv env
+            python3.10 -m venv env
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
@@ -140,7 +140,7 @@ jobs:
       - run:
           name: Install Python package
           command: |
-            uv venv --python 3.9
+            uv venv --python 3.10
             uv pip install \
               nanobind==2.4.0 \
               cmake \
@@ -273,7 +273,7 @@ jobs:
     parameters:
       python_version:
         type: string
-        default: "3.9"
+        default: "3.10"
       xcode_version:
         type: string
         default: "26.0.0"
@@ -328,7 +328,7 @@ jobs:
             << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
       - when:
           condition:
-            equal: ["3.9", << parameters.python_version >>]
+            equal: ["3.10", << parameters.python_version >>]
           steps:
             - run:
                 name: Build common package
@@ -351,7 +351,7 @@ jobs:
     parameters:
       python_version:
         type: string
-        default: "3.9"
+        default: "3.10"
       build_env:
         type: string
         default: ""
@@ -387,7 +387,7 @@ jobs:
             bash python/scripts/repair_linux.sh
       - when:
           condition:
-            equal: ["3.9", << parameters.python_version >>]
+            equal: ["3.10", << parameters.python_version >>]
           steps:
             - run:
                 name: Build common package
@@ -484,7 +484,7 @@ workflows:
               ignore: /.*/
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
               macosx_deployment_target: ["13.5", "14.0", "15.0"]
               build_env: ["PYPI_RELEASE=1"]
               xcode_version: ["26.0.0"]
@@ -503,7 +503,7 @@ workflows:
               ignore: /.*/
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
               build_env: ["PYPI_RELEASE=1"]
       - build_cuda_release:
           filters:
@@ -546,13 +546,13 @@ workflows:
       - build_release:
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
               macosx_deployment_target: ["13.5", "14.0", "15.0"]
               xcode_version: ["26.0.0"]
       - build_linux_release:
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
       - build_cuda_release
 
   build_dev_release:
@@ -564,14 +564,14 @@ workflows:
       - build_release:
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
               macosx_deployment_target: ["13.5", "14.0", "15.0"]
               build_env: ["DEV_RELEASE=1"]
               xcode_version: ["26.0.0"]
       - build_linux_release:
           matrix:
             parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
               build_env: ["DEV_RELEASE=1"]
       - build_cuda_release:
           matrix:

CMakeLists.txt

@@ -20,12 +20,17 @@ project(
   LANGUAGES C CXX
   VERSION ${MLX_PROJECT_VERSION})
 
+if(CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang")
+  add_compile_options(-Wall -Wextra)
+endif()
+
 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # ----------------------------- Configuration -----------------------------
 option(MLX_BUILD_TESTS "Build tests for mlx" ON)
@@ -173,7 +178,7 @@ if(MLX_BUILD_CPU)
     message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
     set(MLX_BUILD_ACCELERATE ON)
   else()
-    message(STATUS "Accelerate or arm neon not found, using default backend.")
+    message(STATUS "Accelerate not found, using default backend.")
     set(MLX_BUILD_ACCELERATE OFF)
   endif()
 

README.md

@@ -2,7 +2,7 @@
 
 [**Quickstart**](#quickstart) | [**Installation**](#installation) |
 [**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
-[**Examples**](#examples) 
+[**Examples**](#examples)
 
 [![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
 
@@ -11,37 +11,37 @@ brought to you by Apple machine learning research.
 
 Some key features of MLX include:
 
- - **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
    also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
    [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
    the Python API. MLX has higher-level packages like `mlx.nn` and
    `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
    more complex models.
 
- - **Composable function transformations**: MLX supports composable function
-   transformations for automatic differentiation, automatic vectorization,
-   and computation graph optimization.
+- **Composable function transformations**: MLX supports composable function
+  transformations for automatic differentiation, automatic vectorization,
+  and computation graph optimization.
 
- - **Lazy computation**: Computations in MLX are lazy. Arrays are only
-   materialized when needed.
+- **Lazy computation**: Computations in MLX are lazy. Arrays are only
+  materialized when needed.
 
- - **Dynamic graph construction**: Computation graphs in MLX are constructed
-   dynamically. Changing the shapes of function arguments does not trigger
-   slow compilations, and debugging is simple and intuitive.
+- **Dynamic graph construction**: Computation graphs in MLX are constructed
+  dynamically. Changing the shapes of function arguments does not trigger
+  slow compilations, and debugging is simple and intuitive.
 
- - **Multi-device**: Operations can run on any of the supported devices
-   (currently the CPU and the GPU).
+- **Multi-device**: Operations can run on any of the supported devices
+  (currently the CPU and the GPU).
 
- - **Unified memory**: A notable difference from MLX and other frameworks
-   is the *unified memory model*. Arrays in MLX live in shared memory.
-   Operations on MLX arrays can be performed on any of the supported
-   device types without transferring data.
+- **Unified memory**: A notable difference from MLX and other frameworks
+  is the *unified memory model*. Arrays in MLX live in shared memory.
+  Operations on MLX arrays can be performed on any of the supported
+  device types without transferring data.
 
 MLX is designed by machine learning researchers for machine learning
 researchers. The framework is intended to be user-friendly, but still efficient
 to train and deploy models. The design of the framework itself is also
 conceptually simple. We intend to make it easy for researchers to extend and
-improve MLX with the goal of quickly exploring new ideas. 
+improve MLX with the goal of quickly exploring new ideas.
 
 The design of MLX is inspired by frameworks like
 [NumPy](https://numpy.org/doc/stable/index.html),
@@ -91,7 +91,7 @@ Checkout the
 [documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
 for more information on building the C++ and Python APIs from source.
 
-## Contributing 
+## Contributing
 
 Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
 on contributing to MLX. See the
@@ -110,7 +110,7 @@ Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
 MLX useful in your research and wish to cite it, please use the following
 BibTex entry:
 
-```
+```text
 @software{mlx2023,
   author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
   title = {{MLX}: Efficient and flexible machine learning on Apple silicon},

benchmarks/python/blas/bench_gemm.py

@@ -142,9 +142,7 @@ def bench_shape(B, M, N, K, np_dtype, transpose="nn"):
     t_b = (0, 1, 2) if transpose[1] == "n" else (0, 2, 1)
 
     c_mlx = a_mx.transpose(t_a) @ b_mx.transpose(t_b)
-    c_npy = a_np.transpose(t_a).astype(np.float32) @ b_np.transpose(t_b).astype(
-        np.float32
-    )
+    c_npy = a_np.transpose(t_a).astype(np_dtype) @ b_np.transpose(t_b).astype(np_dtype)
 
     atol = 1e-5 if np_dtype == np.float32 else 1e-4
 
@@ -163,7 +161,7 @@ def get_gflop_count(B, M, N, K):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run gemm benchmarks")
 
-    dtypes = ("float32", "float16")
+    dtypes = ("float32", "float16", "complex64")
     transposes = ("nn", "nt", "tn")
     shapes = (
         (16, 234, 768, 3072),
@@ -187,7 +185,7 @@ if __name__ == "__main__":
                 diff = gflops_mx / gflops_pt - 1.0
 
                 print(
-                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100. * diff:+5.2f}%"
+                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%"
                 )
                 if gflops_pt >= 2.0 * gflops_mx:
                     print("ATTENTION ^^^^^^^")

benchmarks/python/blas/bench_gemv.py

@@ -196,7 +196,7 @@ def bench_with_out_len(ax, out_vec_len, in_vector_lens, dtype, transpose):
 
 
 for transpose in (False, True):
-    for dtype in ("float32", "float16"):
+    for dtype in ("float32", "float16", "complex64"):
         fig, axs = plt.subplots(
             len(in_vec_sizes), 2, figsize=(8.5, 11), layout="constrained"
         )
@@ -215,7 +215,7 @@ for transpose in (False, True):
         fig.suptitle(f"{device_name}: {dtype} {op_name}")
         fig.savefig(
             os.path.join(
-                results_dir, f'{device_name.replace(" ", "_")}_{dtype}_{op_name}.pdf'
+                results_dir, f"{device_name.replace(' ', '_')}_{dtype}_{op_name}.pdf"
             )
         )
         plt.close(fig)

docs/src/install.rst

@@ -16,7 +16,7 @@ silicon computer is
 To install from PyPI your system must meet the following requirements:
 
 - Using an M series chip (Apple silicon)
-- Using a native Python >= 3.9
+- Using a native Python >= 3.10
 - macOS >= 13.5
 
 .. note::
@@ -39,7 +39,7 @@ requirements:
 - Nvidia driver >= 550.54.14
 - CUDA toolkit >= 12.0
 - Linux distribution with glibc >= 2.35
-- Python >= 3.9
+- Python >= 3.10
 
 
 CPU-only (Linux)
@@ -55,7 +55,7 @@ To install the CPU-only package from PyPi your system must meet the following
 requirements:
 
 - Linux distribution with glibc >= 2.35
-- Python >= 3.9
+- Python >= 3.10
 
 
 Troubleshooting

docs/src/python/ops.rst

@@ -112,6 +112,7 @@ Operations
    max
    maximum
    mean
+   median
    meshgrid
    min
    minimum

examples/cpp/tutorial.cpp

@@ -14,14 +14,17 @@ void array_basics() {
   // Get the value out of it:
   auto s = x.item<float>();
   assert(s == 1.0);
+  (void)s;
 
   // Scalars have a size of 1:
-  size_t size = x.size();
+  int64_t size = x.size();
   assert(size == 1);
+  (void)size;
 
   // Scalars have 0 dimensions:
   int ndim = x.ndim();
   assert(ndim == 0);
+  (void)ndim;
 
   // The shape should be an empty vector:
   auto shape = x.shape();
@@ -30,6 +33,7 @@ void array_basics() {
   // The datatype should be float32:
   auto dtype = x.dtype();
   assert(dtype == mx::float32);
+  (void)dtype;
 
   // Specify the dtype when constructing the array:
   x = mx::array(1, mx::int32);

mlx/array.cpp

@@ -44,11 +44,11 @@ std::vector<array> array::make_arrays(
     const std::shared_ptr<Primitive>& primitive,
     const std::vector<array>& inputs) {
   std::vector<array> outputs;
-  for (size_t i = 0; i < shapes.size(); ++i) {
+  for (int i = 0; i < std::ssize(shapes); ++i) {
     outputs.emplace_back(std::move(shapes[i]), dtypes[i], primitive, inputs);
   }
   // For each node in |outputs|, its siblings are the other nodes.
-  for (size_t i = 0; i < outputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(outputs); ++i) {
     auto siblings = outputs;
     siblings.erase(siblings.begin() + i);
     outputs[i].set_siblings(std::move(siblings), i);
@@ -145,8 +145,9 @@ void array::set_data(allocator::Buffer buffer, Deleter d) {
   array_desc_->data_size = size();
   array_desc_->flags.contiguous = true;
   array_desc_->flags.row_contiguous = true;
-  auto max_dim = std::max_element(shape().begin(), shape().end());
-  array_desc_->flags.col_contiguous = size() <= 1 || size() == *max_dim;
+  auto max_dim =
+      static_cast<int64_t>(*std::max_element(shape().begin(), shape().end()));
+  array_desc_->flags.col_contiguous = size() <= 1 || size() == max_dim;
 }
 
 void array::set_data(
@@ -192,7 +193,7 @@ array::~array() {
   }
 
   // Break circular reference for non-detached arrays with siblings
-  if (auto n = siblings().size(); n > 0) {
+  if (auto n = std::ssize(siblings()); n > 0) {
     bool do_detach = true;
     // If all siblings have siblings.size() references except
     // the one we are currently destroying (which has siblings.size() + 1)
@@ -241,8 +242,8 @@ array::ArrayDesc::ArrayDesc(
     std::vector<array> inputs)
     : shape(std::move(shape)),
       dtype(dtype),
-      status(Status::unscheduled),
       primitive(std::move(primitive)),
+      status(Status::unscheduled),
       inputs(std::move(inputs)) {
   init();
 }
@@ -274,7 +275,7 @@ array::ArrayDesc::~ArrayDesc() {
     ad.inputs.clear();
     for (auto& [_, a] : input_map) {
       bool is_deletable =
-          (a.array_desc_.use_count() <= a.siblings().size() + 1);
+          (a.array_desc_.use_count() <= std::ssize(a.siblings()) + 1);
       // An array with siblings is deletable only if all of its siblings
       // are deletable
       for (auto& s : a.siblings()) {
@@ -283,7 +284,7 @@ array::ArrayDesc::~ArrayDesc() {
         }
         int is_input = (input_map.find(s.id()) != input_map.end());
         is_deletable &=
-            s.array_desc_.use_count() <= a.siblings().size() + is_input;
+            s.array_desc_.use_count() <= std::ssize(a.siblings()) + is_input;
       }
       if (is_deletable) {
         for_deletion.push_back(std::move(a.array_desc_));

mlx/array.h

@@ -81,22 +81,22 @@ class array {
   }
 
   /** The size of the array's datatype in bytes. */
-  size_t itemsize() const {
+  int itemsize() const {
     return size_of(dtype());
   }
 
   /** The number of elements in the array. */
-  size_t size() const {
+  int64_t size() const {
     return array_desc_->size;
   }
 
   /** The number of bytes in the array. */
-  size_t nbytes() const {
+  int64_t nbytes() const {
     return size() * itemsize();
   }
 
   /** The number of dimensions of the array. */
-  size_t ndim() const {
+  int ndim() const {
     return array_desc_->shape.size();
   }
 
@@ -329,7 +329,7 @@ class array {
    * corresponding to ``arr[-1, -1, ...]``) then ``data_size = last - first``.
    * Note, ``data_size`` is in units of ``item_size`` (not bytes).
    **/
-  size_t data_size() const {
+  int64_t data_size() const {
     return array_desc_->data_size;
   }
 
@@ -340,7 +340,7 @@ class array {
     return array_desc_->data->buffer;
   }
 
-  size_t buffer_size() const {
+  int64_t buffer_size() const {
     return allocator::allocator().size(buffer());
   }
 
@@ -530,7 +530,7 @@ array::array(
     Shape shape,
     Dtype dtype /* = TypeToDtype<T>() */)
     : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
-  if (data.size() != size()) {
+  if (std::ssize(data) != size()) {
     throw std::invalid_argument(
         "Data size and provided shape mismatch in array construction.");
   }

mlx/backend/common/common.cpp

@@ -21,8 +21,8 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
 
   // Compute the flags given the shape and strides
   bool row_contiguous = true, col_contiguous = true;
-  size_t r = 1, c = 1;
-  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+  int64_t r = 1, c = 1;
+  for (int i = std::ssize(strides_) - 1, j = 0; i >= 0; i--, j++) {
     row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
     col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
     r *= shape_[i];
@@ -60,7 +60,8 @@ void CustomTransforms::eval(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
   assert(inputs.size() > outputs.size());
-  for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
+  for (int i = 0, j = std::ssize(inputs) - std::ssize(outputs);
+       i < std::ssize(outputs);
        i++, j++) {
     outputs[i].copy_shared_buffer(inputs[j]);
   }
@@ -70,7 +71,7 @@ void Depends::eval(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
   assert(inputs.size() > outputs.size());
-  for (int i = 0; i < outputs.size(); i++) {
+  for (int i = 0; i < std::ssize(outputs); i++) {
     outputs[i].copy_shared_buffer(inputs[i]);
   }
 }
@@ -206,11 +207,11 @@ void Split::eval(
 
   auto compute_new_flags = [](const auto& shape,
                               const auto& strides,
-                              size_t in_data_size,
+                              int64_t in_data_size,
                               auto flags) {
-    size_t data_size = 1;
-    size_t f_stride = 1;
-    size_t b_stride = 1;
+    int64_t data_size = 1;
+    int64_t f_stride = 1;
+    int64_t b_stride = 1;
     flags.row_contiguous = true;
     flags.col_contiguous = true;
     for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) {
@@ -240,7 +241,7 @@ void Split::eval(
 
   std::vector<int> indices(1, 0);
   indices.insert(indices.end(), indices_.begin(), indices_.end());
-  for (int i = 0; i < indices.size(); i++) {
+  for (int i = 0; i < std::ssize(indices); i++) {
     size_t offset = indices[i] * in.strides()[axis_];
     auto [new_flags, data_size] = compute_new_flags(
         outputs[i].shape(), in.strides(), in.data_size(), in.flags());
@@ -254,7 +255,7 @@ void Squeeze::eval(const std::vector<array>& inputs, array& out) {
   const auto& in = inputs[0];
   Strides strides;
   for (int i = 0, j = 0; i < in.ndim(); ++i) {
-    if (j < axes_.size() && i == axes_[j]) {
+    if (j < std::ssize(axes_) && i == axes_[j]) {
       j++;
     } else {
       strides.push_back(in.strides(i));
@@ -272,7 +273,7 @@ void Transpose::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   Strides out_strides(out.ndim());
   auto& in = inputs[0];
-  for (int ax = 0; ax < axes_.size(); ++ax) {
+  for (int ax = 0; ax < std::ssize(axes_); ++ax) {
     out_strides[ax] = in.strides()[axes_[ax]];
   }
 

mlx/backend/common/compiled.cpp

@@ -120,7 +120,7 @@ void compiled_allocate_outputs(
     Strides strides;
     size_t data_size;
     array::Flags flags;
-    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+    for (int i = 0; i < std::ssize(inputs) && o < std::ssize(outputs); ++i) {
       auto& in = inputs[i];
       // Conditions for donation
       // - Correct size
@@ -138,7 +138,7 @@ void compiled_allocate_outputs(
         data_size = in.data_size();
       }
     }
-    for (; o < outputs.size(); ++o) {
+    for (; o < std::ssize(outputs); ++o) {
       outputs[o].set_data(
           allocator::malloc(data_size * outputs[o].itemsize()),
           data_size,
@@ -147,7 +147,7 @@ void compiled_allocate_outputs(
     }
   } else {
     int o = 0;
-    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+    for (int i = 0; i < std::ssize(inputs) && o < std::ssize(outputs); ++i) {
       auto& in = inputs[i];
       // Conditions for donation
       // - Row contiguous
@@ -162,7 +162,7 @@ void compiled_allocate_outputs(
         o++;
       }
     }
-    for (; o < outputs.size(); ++o) {
+    for (; o < std::ssize(outputs); ++o) {
       outputs[o].set_data(allocator::malloc(outputs[o].nbytes()));
     }
   }
@@ -193,7 +193,7 @@ std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
 
     // Broadcast the inputs to the output shape.
     Strides xstrides;
-    size_t j = 0;
+    int j = 0;
     for (; j < shape.size() - x.ndim(); ++j) {
       if (shape[j] == 1) {
         xstrides.push_back(out.strides()[j]);
@@ -201,7 +201,7 @@ std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
         xstrides.push_back(0);
       }
     }
-    for (size_t i = 0; i < x.ndim(); ++i, ++j) {
+    for (int i = 0; i < x.ndim(); ++i, ++j) {
       if (x.shape(i) == 1) {
         if (shape[j] == 1) {
           xstrides.push_back(out.strides()[j]);
@@ -224,13 +224,13 @@ bool compiled_use_large_index(
     const std::vector<array>& outputs,
     bool contiguous) {
   if (contiguous) {
-    size_t max_size = 0;
+    int64_t max_size = 0;
     for (const auto& in : inputs) {
       max_size = std::max(max_size, in.data_size());
     }
     return max_size > UINT32_MAX;
   } else {
-    size_t max_size = 0;
+    int64_t max_size = 0;
     for (const auto& o : outputs) {
       max_size = std::max(max_size, o.size());
     }

mlx/backend/common/load.cpp

@@ -27,7 +27,7 @@ void swap_endianness(uint8_t* data_bytes, size_t N) {
 
 namespace mlx::core {
 
-void Load::eval_cpu(const std::vector<array>& inputs, array& out) {
+void Load::eval_cpu(const std::vector<array>& /* inputs */, array& out) {
   out.set_data(allocator::malloc(out.nbytes()));
   auto read_task = [out_ptr = out.data<char>(),
                     size = out.size(),

mlx/backend/common/matmul.h

@@ -13,7 +13,7 @@ inline std::tuple<Shape, Strides, Strides> collapse_batches(
     const array& a,
     const array& b) {
   if (a.ndim() == 2) {
-    return {{1}, {0}, {0}};
+    return {Shape{1}, Strides{0}, Strides{0}};
   }
 
   Shape A_bshape{a.shape().begin(), a.shape().end() - 2};
@@ -38,7 +38,7 @@ inline std::tuple<Shape, Strides, Strides> collapse_batches(
 inline std::tuple<Shape, Strides, Strides, Strides>
 collapse_batches(const array& a, const array& b, const array& c) {
   if (a.ndim() == 2) {
-    return {{1}, {0}, {0}, {0}};
+    return {Shape{1}, Strides{0}, Strides{0}, Strides{0}};
   }
 
   Shape A_bshape{a.shape().begin(), a.shape().end() - 2};

mlx/backend/common/reduce.cpp

@@ -28,7 +28,7 @@ std::pair<Shape, Strides> shapes_without_reduction_axes(
 
 ReductionPlan get_reduction_plan(const array& x, const std::vector<int>& axes) {
   // The data is all there and we are reducing over everything
-  if (x.size() == x.data_size() && axes.size() == x.ndim() &&
+  if (x.size() == x.data_size() && std::ssize(axes) == x.ndim() &&
       x.flags().contiguous) {
     return ContiguousAllReduce;
   }
@@ -38,7 +38,7 @@ ReductionPlan get_reduction_plan(const array& x, const std::vector<int>& axes) {
     // Merge consecutive axes
     Shape shape = {x.shape(axes[0])};
     Strides strides = {x.strides()[axes[0]]};
-    for (int i = 1; i < axes.size(); i++) {
+    for (int i = 1; i < std::ssize(axes); i++) {
       if (axes[i] - 1 == axes[i - 1] && x.shape(axes[i]) > 1) {
         shape.back() *= x.shape(axes[i]);
         strides.back() = x.strides()[axes[i]];

mlx/backend/common/slicing.cpp

@@ -24,8 +24,8 @@ std::tuple<int64_t, Strides> prepare_slice(
 void shared_buffer_slice(
     const array& in,
     const Strides& out_strides,
-    size_t data_offset,
-    size_t data_size,
+    int64_t data_offset,
+    int64_t data_size,
     array& out) {
   // Compute row/col contiguity
   auto [no_bsx_size, is_row_contiguous, is_col_contiguous] =
@@ -61,7 +61,7 @@ void slice(
   if (data_end < 0) {
     data_end += in.data_size();
   }
-  size_t data_size = (data_end - data_offset);
+  int64_t data_size = (data_end - data_offset);
   shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
 }
 

mlx/backend/common/ternary.h

@@ -11,6 +11,8 @@ namespace mlx::core {
 enum class TernaryOpType {
   ScalarScalarScalar,
   VectorVectorVector,
+  VectorVectorScalar,
+  VectorScalarVector,
   General,
 };
 
@@ -25,6 +27,14 @@ get_ternary_op_type(const array& a, const array& b, const array& c) {
       (a.flags().col_contiguous && b.flags().col_contiguous &&
        c.flags().col_contiguous)) {
     topt = TernaryOpType::VectorVectorVector;
+  } else if (
+      b.data_size() == 1 && a.flags().row_contiguous &&
+      c.flags().row_contiguous) {
+    topt = TernaryOpType::VectorScalarVector;
+  } else if (
+      c.data_size() == 1 && a.flags().row_contiguous &&
+      b.flags().row_contiguous) {
+    topt = TernaryOpType::VectorVectorScalar;
   } else {
     topt = TernaryOpType::General;
   }
@@ -59,6 +69,8 @@ inline void set_ternary_op_output_data(
             b.flags());
       }
       break;
+    case TernaryOpType::VectorVectorScalar:
+    case TernaryOpType::VectorScalarVector:
     case TernaryOpType::General:
       // Try to donate an input which is row_contiguous
       if (!((a.flags().row_contiguous && maybe_donate(a)) ||

mlx/backend/common/utils.cpp

@@ -28,7 +28,7 @@ std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
     if (shape[0] != 1) {
       to_collapse.push_back(0);
     }
-    size_t size = shape[0];
+    int64_t size = shape[0];
     for (int i = 1; i < shape.size(); i++) {
       bool contiguous = true;
       size *= shape[i];
@@ -64,7 +64,7 @@ std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
       current_shape *= shape[to_collapse[k]];
     }
     out_shape.push_back(current_shape);
-    for (int j = 0; j < strides.size(); j++) {
+    for (int j = 0; j < std::ssize(strides); j++) {
       const auto& st = strides[j];
       out_strides[j].push_back(st[to_collapse[k - 1]]);
     }

mlx/backend/common/utils.h

@@ -162,7 +162,7 @@ struct ContiguousIterator {
 };
 
 inline auto check_contiguity(const Shape& shape, const Strides& strides) {
-  size_t no_broadcast_data_size = 1;
+  int64_t no_broadcast_data_size = 1;
   int64_t f_stride = 1;
   int64_t b_stride = 1;
   bool is_row_contiguous = true;
@@ -183,7 +183,7 @@ inline auto check_contiguity(const Shape& shape, const Strides& strides) {
 }
 
 inline bool is_donatable(const array& in, const array& out) {
-  constexpr size_t donation_extra = 16384;
+  constexpr int64_t donation_extra = 16384;
 
   return in.is_donatable() && in.itemsize() == out.itemsize() &&
       in.buffer_size() <= out.nbytes() + donation_extra;

mlx/backend/cpu/arange.h

@@ -10,7 +10,7 @@ namespace mlx::core {
 namespace {
 
 template <typename T>
-void arange(T start, T next, array& out, size_t size, Stream stream) {
+void arange(T start, T next, array& out, int64_t size, Stream stream) {
   auto ptr = out.data<T>();
   auto step_size = next - start;
   auto& encoder = cpu::get_command_encoder(stream);

mlx/backend/cpu/arg_reduce.cpp

@@ -19,12 +19,12 @@ void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
   auto in_ptr = in.data<InT>();
   auto out_ptr = out.data<uint32_t>();
 
-  for (uint32_t i = 0; i < out.size(); ++i) {
+  for (int64_t i = 0; i < out.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) {
+    for (int64_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;

mlx/backend/cpu/binary.cpp

@@ -17,7 +17,12 @@ namespace mlx::core {
 namespace {
 
 template <typename Op>
-void binary(const array& a, const array& b, array& out, Op op, Stream stream) {
+void binary(
+    const array& a,
+    const array& b,
+    array& out,
+    Op /* op */,
+    Stream stream) {
   auto bopt = get_binary_op_type(a, b);
   set_binary_op_output_data(a, b, out, bopt);
 
@@ -81,7 +86,7 @@ void comparison_op(
     const array& a,
     const array& b,
     array& out,
-    Op op,
+    Op /* op */,
     Stream stream) {
   auto bopt = get_binary_op_type(a, b);
   set_binary_op_output_data(a, b, out, bopt);
@@ -146,7 +151,7 @@ void binary_float(
     const array& a,
     const array& b,
     array& out,
-    Op op,
+    Op /* op */,
     Stream stream) {
   auto bopt = get_binary_op_type(a, b);
   set_binary_op_output_data(a, b, out, bopt);
@@ -187,7 +192,7 @@ void binary_int(
     const array& a,
     const array& b,
     array& out,
-    Op op,
+    Op /* op */,
     Stream stream) {
   auto bopt = get_binary_op_type(a, b);
   set_binary_op_output_data(a, b, out, bopt);

mlx/backend/cpu/binary_two.h

@@ -99,7 +99,7 @@ void binary_op_dispatch_dims(
   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) {
+  for (int64_t elem = 0; elem < std::ssize(a); elem += stride) {
     binary_op_dims<T, U, Op, 2>(
         a_ptr + a_it.loc,
         b_ptr + b_it.loc,
@@ -137,21 +137,21 @@ void binary_op(
   if (bopt == BinaryOpType::ScalarScalar) {
     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.data_size(); ++i) {
+    for (int64_t i = 0; i < b.data_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.data_size(); ++i) {
+    for (int64_t i = 0; i < a.data_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) {
+    for (int64_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++;

mlx/backend/cpu/cholesky.cpp

@@ -33,8 +33,8 @@ void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) {
                     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++) {
+    int64_t num_matrices = size / (N * N);
+    for (int64_t i = 0; i < num_matrices; i++) {
       // Compute Cholesky factorization.
       int info;
       potrf<T>(

mlx/backend/cpu/compiled.cpp

@@ -49,7 +49,7 @@ static CompilerCache& cache() {
 // GPU compile is always available if the GPU is available and since we are in
 // this file CPU compile is also available.
 namespace detail {
-bool compile_available_for_device(const Device& device) {
+bool compile_available_for_device(const Device& /* device */) {
   return true;
 }
 
@@ -168,7 +168,7 @@ inline void build_kernel(
   // Add the input arguments
   int cnt = 0;
   int strides_index = 1;
-  for (size_t i = 0; i < inputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(inputs); ++i) {
     // Skip constants from the input list
     if (is_constant(i)) {
       continue;
@@ -238,7 +238,7 @@ inline void build_kernel(
     } else {
       os << x.primitive().name();
       os << "()(";
-      for (int i = 0; i < x.inputs().size() - 1; i++) {
+      for (int i = 0; i < std::ssize(x.inputs()) - 1; i++) {
         os << "tmp_" << namer.get_name(x.inputs()[i]) << ", ";
       }
       os << "tmp_" << namer.get_name(x.inputs().back()) << ");" << std::endl;

mlx/backend/cpu/conv.cpp

@@ -860,7 +860,7 @@ void explicit_gemm_conv_1D_cpu(
     const std::vector<int>& padding_lo,
     const std::vector<int>& padding_hi,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
+    const std::vector<int>& /* wt_dilation */,
     Stream stream) {
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
   const int iH = in.shape(1); // Input spatial dim
@@ -996,131 +996,6 @@ void explicit_gemm_conv_1D_cpu(
   encoder.add_temporaries(std::move(temps));
 }
 
-void explicit_gemm_conv_2D_cpu(
-    const array& in,
-    const array& wt,
-    array out,
-    const std::vector<int>& padding_lo,
-    const std::vector<int>& padding_hi,
-    const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
-    Stream stream) {
-  const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
-  const int iH = in.shape(1); // Input spatial dim
-  const int iW = in.shape(2); // Input spatial dim
-  const int oH = out.shape(1); // Output spatial dim
-  const int oW = out.shape(2); // Output spatial dim
-  const int O = wt.shape(0); // Out channels
-  const int C = wt.shape(3); // In channels
-  const int wH = wt.shape(1); // Weight spatial dim
-  const int wW = wt.shape(2); // Weight spatial dim
-
-  auto conv_dtype = out.dtype();
-  auto& encoder = cpu::get_command_encoder(stream);
-
-  // Pad input
-  Shape padded_shape = {
-      N,
-      iH + padding_lo[0] + padding_hi[0],
-      iW + padding_lo[1] + padding_hi[1],
-      C};
-  array in_padded(padded_shape, conv_dtype, nullptr, {});
-
-  // Fill with zeros
-  std::vector<array> temps;
-  temps.push_back(array(0, conv_dtype));
-  copy_cpu(temps.back(), in_padded, CopyType::Scalar, stream);
-
-  // Pick input slice from padded
-  size_t data_offset = padding_lo[0] * in_padded.strides()[1] +
-      padding_lo[1] * in_padded.strides()[2];
-  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
-  in_padded_slice.copy_shared_buffer(
-      in_padded,
-      in_padded.strides(),
-      in_padded.flags(),
-      in_padded_slice.size(),
-      data_offset);
-  temps.push_back(in_padded_slice);
-
-  // Copy input values into the slice
-  copy_cpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, stream);
-
-  // Make strided view
-  Shape strided_shape = {N, oH, oW, wH, wW, C};
-
-  Strides strided_strides = {
-      in_padded.strides()[0],
-      in_padded.strides()[1] * wt_strides[0],
-      in_padded.strides()[2] * wt_strides[1],
-      in_padded.strides()[1],
-      in_padded.strides()[2],
-      in_padded.strides()[3]};
-  auto flags = in_padded.flags();
-
-  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
-  in_strided_view.copy_shared_buffer(
-      in_padded, strided_strides, flags, in_strided_view.size(), 0);
-
-  // Materialize strided view
-  Shape strided_reshape = {N * oH * oW, wH * wW * C};
-  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
-  copy_cpu(in_strided_view, in_strided, CopyType::General, stream);
-  temps.push_back(in_strided);
-
-  // Check wt dtype and prepare
-  auto gemm_wt = wt;
-  auto gemm_out = out;
-
-  if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
-    auto ctype =
-        wt.flags().row_contiguous ? CopyType::Vector : CopyType::General;
-    gemm_wt = array(wt.shape(), float32, nullptr, {});
-    copy_cpu(wt, gemm_wt, ctype, stream);
-    temps.push_back(gemm_wt);
-  }
-
-  if (out.dtype() != float32) {
-    gemm_out = array(out.shape(), float32, nullptr, {});
-    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
-    temps.push_back(gemm_out);
-  }
-
-  encoder.set_input_array(in_strided);
-  encoder.set_input_array(gemm_wt);
-  encoder.set_output_array(gemm_out);
-
-  encoder.dispatch([in_strided_ptr = in_strided.data<float>(),
-                    gemm_wt_ptr = gemm_wt.data<float>(),
-                    gemm_out_ptr = gemm_out.data<float>(),
-                    strided_reshape = std::move(strided_reshape),
-                    O]() {
-    // Perform gemm
-    cblas_sgemm(
-        CblasRowMajor,
-        CblasNoTrans, // no trans A
-        CblasTrans, // transB
-        strided_reshape[0], // M
-        O, // N
-        strided_reshape[1], // K
-        1.0f, // alpha
-        in_strided_ptr,
-        strided_reshape[1], // lda
-        gemm_wt_ptr,
-        strided_reshape[1], // ldb
-        0.0f, // beta
-        gemm_out_ptr,
-        O // ldc
-    );
-  });
-
-  // Copy results if needed
-  if (out.dtype() != float32) {
-    copy_cpu_inplace(gemm_out, out, CopyType::Vector, stream);
-  }
-  encoder.add_temporaries(std::move(temps));
-}
-
 void explicit_gemm_conv_ND_cpu(
     const array& in,
     const array& wt,
@@ -1128,7 +1003,7 @@ void explicit_gemm_conv_ND_cpu(
     const std::vector<int>& padding_lo,
     const std::vector<int>& padding_hi,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
+    const std::vector<int>& /* wt_dilation */,
     const bool flip,
     Stream stream) {
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
@@ -1148,7 +1023,7 @@ void explicit_gemm_conv_ND_cpu(
   // Pad input
   Shape padded_shape(in.shape().size());
   padded_shape.front() = N;
-  for (size_t i = 0; i < iDim.size(); i++) {
+  for (int i = 0; i < iDim.size(); i++) {
     padded_shape[i + 1] = iDim[i] + padding_lo[i] + padding_hi[i];
   }
   padded_shape.back() = C;
@@ -1179,20 +1054,20 @@ void explicit_gemm_conv_ND_cpu(
   // Make strided view
   Shape strided_shape(oDim.size() + wDim.size() + 2);
   strided_shape.front() = N;
-  for (size_t i = 0; i < oDim.size(); i++) {
+  for (int i = 0; i < oDim.size(); i++) {
     strided_shape[i + 1] = oDim[i];
   }
-  for (size_t i = 0; i < wDim.size(); i++) {
+  for (int i = 0; i < wDim.size(); i++) {
     strided_shape[i + 1 + oDim.size()] = wDim[i];
   }
   strided_shape.back() = C;
 
   Strides strided_strides(in.shape().size() * 2 - 2);
   strided_strides[0] = in_padded.strides()[0];
-  for (size_t i = 0; i < wt_strides.size(); i++) {
+  for (int i = 0; i < std::ssize(wt_strides); i++) {
     strided_strides[i + 1] = in_padded.strides()[i + 1] * wt_strides[i];
   }
-  for (size_t i = 1; i < in_padded.strides().size(); i++) {
+  for (int i = 1; i < std::ssize(in_padded.strides()); i++) {
     strided_strides[i + wt_strides.size()] = in_padded.strides()[i];
   }
 

mlx/backend/cpu/distributed.cpp

@@ -90,6 +90,7 @@ void Recv::eval_cpu(
     std::vector<array>& outputs) {
   assert(inputs.size() == 0);
   assert(outputs.size() == 1);
+  (void)inputs;
 
   outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
   distributed::detail::recv(group(), outputs[0], src_, stream());

mlx/backend/cpu/eig.cpp

@@ -46,7 +46,6 @@ void eig_impl(
     int info;
     {
       T work;
-      int iwork;
       geev<T>(
           &jobl,
           &jobr,
@@ -71,7 +70,7 @@ void eig_impl(
     auto eig_tmp = static_cast<T*>(eig_tmp_data.buffer.raw_ptr());
     auto vec_tmp = static_cast<T*>(vec_tmp_data.buffer.raw_ptr());
     auto work_buf = array::Data{allocator::malloc(sizeof(T) * lwork)};
-    for (size_t i = 0; i < size / (N * N); ++i) {
+    for (int64_t i = 0; i < size / (N * N); ++i) {
       geev<T>(
           &jobl,
           &jobr,

mlx/backend/cpu/eigh.cpp

@@ -165,7 +165,7 @@ void eigh_impl(
     EighWork<T> work(jobz, uplo, N);
 
     // Work loop
-    for (size_t i = 0; i < size / (N * N); ++i) {
+    for (int64_t i = 0; i < size / (N * N); ++i) {
       work.run(vec_ptr, eig_ptr);
       vec_ptr += N * N;
       eig_ptr += N;

mlx/backend/cpu/encoder.h

@@ -20,8 +20,8 @@ struct CommandEncoder {
   CommandEncoder(CommandEncoder&&) = delete;
   CommandEncoder& operator=(CommandEncoder&&) = delete;
 
-  void set_input_array(const array& a) {}
-  void set_output_array(array& a) {}
+  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.

mlx/backend/cpu/gemm.h

@@ -12,12 +12,12 @@ void matmul(
     T* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
     float alpha,
     float beta,
-    size_t batch_size,
+    int64_t batch_size,
     const Shape& a_shape,
     const Strides& a_strides,
     const Shape& b_shape,

mlx/backend/cpu/gemms/bnns.cpp

@@ -1,5 +1,4 @@
 // Copyright © 2023-2024 Apple Inc.
-
 #include <Accelerate/Accelerate.h>
 
 #include "mlx/array.h"
@@ -35,7 +34,7 @@ void matmul_bnns(
     bool b_transposed,
     size_t lda,
     size_t ldb,
-    size_t ldc,
+    size_t /* ldc */,
     float alpha,
     float beta,
     size_t batch_size,
@@ -49,9 +48,15 @@ void matmul_bnns(
   size_t K = a_shape[ndim - 1];
 
   BNNSDataType bnns_dtype = to_bnns_dtype<T>();
-
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+  if (beta != 1.0 && beta != 0.0) {
+    // scale the output
+    for (size_t i = 0; i < batch_size * M * N; ++i) {
+      out[i] *= beta;
+    }
+    beta = 1.0;
+  }
   const BNNSLayerParametersBroadcastMatMul gemm_params{
       /* float alpha = */ alpha,
       /* float beta = */ beta,
@@ -122,7 +127,7 @@ void matmul_bnns(
   auto bnns_filter =
       BNNSFilterCreateLayerBroadcastMatMul(&gemm_params, nullptr);
 
-  for (int i = 0; i < batch_size; ++i) {
+  for (size_t i = 0; i < batch_size; ++i) {
     BNNSFilterApplyTwoInput(
         bnns_filter,
         reinterpret_cast<const uint8_t*>(
@@ -143,12 +148,12 @@ void matmul<float16_t>(
     float16_t* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
     float alpha,
     float beta,
-    size_t batch_size,
+    int64_t batch_size,
     const Shape& a_shape,
     const Strides& a_strides,
     const Shape& b_shape,
@@ -178,12 +183,12 @@ void matmul<bfloat16_t>(
     bfloat16_t* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
     float alpha,
     float beta,
-    size_t batch_size,
+    int64_t batch_size,
     const Shape& a_shape,
     const Strides& a_strides,
     const Shape& b_shape,

mlx/backend/cpu/gemms/cblas.cpp

@@ -13,20 +13,20 @@ void matmul<float>(
     float* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
     float alpha,
     float beta,
-    size_t batch_size,
+    int64_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];
+  int64_t M = a_shape[ndim - 2];
+  int64_t N = b_shape[ndim - 1];
+  int64_t K = a_shape[ndim - 1];
 
   for (int i = 0; i < batch_size; ++i) {
     cblas_sgemm(
@@ -54,20 +54,20 @@ void matmul<double>(
     double* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
     float alpha,
     float beta,
-    size_t batch_size,
+    int64_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];
+  int64_t M = a_shape[ndim - 2];
+  int64_t N = b_shape[ndim - 1];
+  int64_t K = a_shape[ndim - 1];
 
   for (int i = 0; i < batch_size; ++i) {
     cblas_dgemm(
@@ -88,4 +88,47 @@ void matmul<double>(
   }
 }
 
+template <>
+void matmul<complex64_t>(
+    const complex64_t* a,
+    const complex64_t* b,
+    complex64_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    float alpha,
+    float beta,
+    int64_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  int64_t M = a_shape[ndim - 2];
+  int64_t N = b_shape[ndim - 1];
+  int64_t K = a_shape[ndim - 1];
+  auto calpha = static_cast<complex64_t>(alpha);
+  auto cbeta = static_cast<complex64_t>(beta);
+
+  for (int i = 0; i < batch_size; ++i) {
+    cblas_cgemm(
+        CblasRowMajor,
+        a_transposed ? CblasTrans : CblasNoTrans, // transA
+        b_transposed ? CblasTrans : CblasNoTrans, // transB
+        M,
+        N,
+        K,
+        &calpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        lda,
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        ldb,
+        &cbeta,
+        out + M * N * i,
+        ldc);
+  }
+}
+
 } // namespace mlx::core

mlx/backend/cpu/hadamard.cpp

@@ -11,9 +11,9 @@ namespace mlx::core {
 
 // n = 2^k component
 template <typename T>
-void hadamard_n(T* out, int n, int m, float scale, size_t size) {
+void hadamard_n(T* out, int n, int /* m */, float scale, int64_t size) {
   for (int b = 0; b < size / n; b++) {
-    size_t loc = b * n;
+    int64_t loc = b * n;
     T* data_ptr = out + loc;
     int h = 1;
     int n_over_2 = n / 2;
@@ -37,7 +37,7 @@ void hadamard_n(T* out, int n, int m, float scale, size_t size) {
 
 // m component
 template <typename T>
-void hadamard_m(T* out, int n, int m, float scale, size_t size) {
+void hadamard_m(T* out, int n, int m, float scale, int64_t size) {
   auto h_matrices = hadamard_matrices();
   auto& matrix = h_matrices[m];
   auto start = 1;
@@ -45,7 +45,7 @@ void hadamard_m(T* out, int n, int m, float scale, size_t size) {
   std::vector<bool> hmat_vec;
   while (end != std::string_view::npos) {
     auto row = matrix.substr(start, end - start);
-    for (int i = 0; i < row.length(); i++) {
+    for (int i = 0; i < std::ssize(row); i++) {
       hmat_vec.push_back(row[i] == '+');
     }
     start = end + 1;
@@ -53,7 +53,7 @@ void hadamard_m(T* out, int n, int m, float scale, size_t size) {
   }
 
   for (int b = 0; b < size / m / n; b++) {
-    size_t loc = b * n * m;
+    int64_t loc = b * n * m;
     T* data_ptr = out + loc;
     for (int i = 0; i < n; i++) {
       std::vector<float> out(m);

mlx/backend/cpu/indexing.cpp

@@ -78,7 +78,7 @@ void gather(
     can_copy = true;
 
     // Ignore leading 1s
-    int i = 0;
+    int64_t i = 0;
     for (; i < slice_sizes.size() && slice_sizes[i] == 1; ++i)
       ;
 
@@ -91,7 +91,7 @@ void gather(
     can_copy = true;
 
     // Ignore trailing 1s
-    int i = slice_sizes.size() - 1;
+    int64_t i = slice_sizes.size() - 1;
     for (; i >= 0 && slice_sizes[i] == 1; --i)
       ;
 
@@ -101,11 +101,11 @@ void gather(
       can_copy = (src.shape(i) == slice_sizes[i]);
     }
   }
-  size_t slice_size = 1;
+  int64_t slice_size = 1;
   for (auto s : slice_sizes) {
     slice_size *= s;
   }
-  size_t ind_size = slice_size == 0 ? 0 : out.size() / slice_size;
+  int64_t ind_size = slice_size == 0 ? 0 : out.size() / slice_size;
   const T* src_ptr = src.data<T>();
   T* dst_ptr = out.data<T>();
 
@@ -115,10 +115,10 @@ void gather(
     src_it = ContiguousIterator(slice_sizes, src.strides(), src.ndim());
   }
 
-  size_t out_idx = 0;
-  for (int idx = 0; idx < ind_size; idx++) {
-    size_t src_idx = 0;
-    for (int ii = 0; ii < inds.size(); ++ii) {
+  int64_t out_idx = 0;
+  for (int64_t idx = 0; idx < ind_size; idx++) {
+    int64_t src_idx = 0;
+    for (int ii = 0; ii < std::ssize(inds); ++ii) {
       auto ax = axes[ii];
       auto idx_loc = its[ii].loc;
       its[ii].step();
@@ -134,7 +134,7 @@ void gather(
           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++) {
+      for (int64_t jj = 0; jj < slice_size; jj++) {
         dst_ptr[out_idx++] = src_ptr[src_idx + src_it.loc];
         src_it.step();
       }
@@ -403,11 +403,11 @@ void scatter(
     const std::vector<int>& axes) {
   int nind = inds.size();
   auto inds_ndim = updates.ndim() - out.ndim();
-  size_t n_updates = nind ? inds[0].size() : 1;
+  int64_t n_updates = nind ? inds[0].size() : 1;
 
   Shape update_shape(
       updates.shape().begin() + inds_ndim, updates.shape().end());
-  size_t update_size = 1;
+  int64_t update_size = 1;
   for (auto us : update_shape) {
     update_size *= us;
   }
@@ -418,9 +418,9 @@ void scatter(
 
   auto out_ptr = out.data<InT>();
   auto upd_ptr = updates.data<InT>();
-  for (int i = 0; i < n_updates; ++i) {
-    size_t out_offset = 0;
-    for (int j = 0; j < inds.size(); ++j) {
+  for (int64_t i = 0; i < n_updates; ++i) {
+    int64_t out_offset = 0;
+    for (int j = 0; j < std::ssize(inds); ++j) {
       auto ax = axes[j];
       auto idx_loc = its[j].loc;
       its[j].step();
@@ -429,7 +429,7 @@ void scatter(
       out_offset += (idx_val * out.strides()[ax]);
     }
     update_it.seek(i * update_size);
-    for (int j = 0; j < update_size; ++j) {
+    for (int64_t j = 0; j < update_size; ++j) {
       OpT{}(upd_ptr[update_it.loc], out_ptr + out_offset + out_it.loc);
       update_it.step();
       out_it.step();

mlx/backend/cpu/inverse.cpp

@@ -122,7 +122,7 @@ void inverse_impl(
       stream);
 
   const int N = a.shape(-1);
-  const size_t num_matrices = a.size() / (N * N);
+  const int64_t num_matrices = a.size() / (N * N);
 
   auto& encoder = cpu::get_command_encoder(stream);
   encoder.set_output_array(inv);
@@ -130,13 +130,13 @@ void inverse_impl(
   auto inv_ptr = inv.data<T>();
   if (tri) {
     encoder.dispatch([inv_ptr, N, num_matrices, upper]() {
-      for (int i = 0; i < num_matrices; i++) {
+      for (int64_t 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++) {
+      for (int64_t i = 0; i < num_matrices; i++) {
         general_inv<T>(inv_ptr + N * N * i, N);
       }
     });

mlx/backend/cpu/masked_mm.cpp

@@ -25,7 +25,7 @@ inline void mask_matrix(
     const int64_t Y_data_str,
     const int64_t X_mask_str,
     const int64_t Y_mask_str,
-    const size_t mask_offset) {
+    const int64_t mask_offset) {
   int tX = (X + block_size - 1) / block_size;
   int tY = (Y + block_size - 1) / block_size;
 
@@ -61,13 +61,13 @@ inline void segmented_mm(
     T* out,
     bool a_transposed,
     bool b_transposed,
-    size_t lda,
-    size_t ldb,
+    int64_t lda,
+    int64_t ldb,
     const Shape& a_shape,
     const Strides& a_strides,
     const Shape& b_shape,
     const Strides& b_strides,
-    size_t num_segments,
+    int64_t num_segments,
     const Shape& segments_shape,
     const Strides& segments_strides) {
   int ndim = a_shape.size();
@@ -149,9 +149,9 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto [b_transposed, ldb, b, b_copied] =
       check_transpose(b_pre, has_op_mask, inputs.back().dtype() != bool_);
 
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
+  int64_t M = a.shape(-2);
+  int64_t N = b.shape(-1);
+  int64_t K = a.shape(-1);
 
   if (M == 0 || N == 0) {
     return;
@@ -172,8 +172,8 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
                        int batch_idx,
                        int X,
                        int Y,
-                       size_t X_data_str,
-                       size_t Y_data_str,
+                       int64_t X_data_str,
+                       int64_t Y_data_str,
                        const Shape& mask_shape,
                        const Strides& mask_strides,
                        bool is_bool) {
@@ -215,18 +215,18 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   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;
+  const void* a_mask_ptr = nullptr;
+  const void* b_mask_ptr = nullptr;
+  const void* out_mask_ptr = nullptr;
   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;
+  bool a_mask_bool = false;
+  bool b_mask_bool = false;
+  bool out_mask_bool = false;
   if (has_op_mask) {
     auto& a_mask = inputs[inputs.size() - 2];
     auto& b_mask = inputs[inputs.size() - 1];
@@ -253,7 +253,7 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, 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));
+  int64_t num_matrices = out.size() / (M * int64_t(N));
   auto ldc = out.shape(-1);
 
   encoder.dispatch([a_ptr,
@@ -394,9 +394,9 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto [a_transposed, lda, a] = check_transpose(a_pre);
   auto [b_transposed, ldb, b] = check_transpose(b_pre);
 
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
+  int64_t M = a.shape(-2);
+  int64_t N = b.shape(-1);
+  int64_t K = a.shape(-1);
 
   if (M == 0 || N == 0) {
     return;
@@ -413,7 +413,7 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   // Get batch dims
   auto batch_size_out = out.size() / (M * N);
-  size_t matrix_stride_out = M * N;
+  int64_t matrix_stride_out = M * N;
 
   auto get_batch_dims = [](const auto& v) {
     return decltype(v){v.begin(), v.end() - 2};
@@ -423,7 +423,6 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& rhs_indices = inputs[3];
 
   auto batch_shape = get_batch_dims(out.shape());
-  int batch_ndim = batch_shape.size();
 
   auto batch_shape_A = get_batch_dims(a.shape());
   auto batch_strides_A = get_batch_dims(a.strides());

mlx/backend/cpu/matmul.cpp

@@ -91,7 +91,6 @@ void matmul_general(
   auto [b_transposed, ldb, b] = check_transpose(b_pre);
   size_t M = a.shape(-2);
   size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
   if (M == 0 || N == 0) {
     return;
   }
@@ -108,6 +107,9 @@ void matmul_general(
   } else if (out.dtype() == float64) {
     matmul_dispatch<double>(
         a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == complex64) {
+    matmul_dispatch<complex64_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
   } else {
     throw std::runtime_error("[Matmul::eval_cpu] Invalid type.");
   }
@@ -128,10 +130,6 @@ void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {
 }
 
 void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
-  if (out.dtype() != float32) {
-    throw std::runtime_error(
-        "[AddMM::eval_cpu] Currently only supports float32.");
-  }
   if (out.size() == 0) {
     out.set_data(allocator::malloc(out.nbytes()));
     return;

mlx/backend/cpu/primitives.cpp

@@ -48,7 +48,7 @@ static std::pair<array, bool> compute_dynamic_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) {
+        for (int i = 0; i < std::ssize(axes); ++i) {
           offset_ += indices[i] * strides[axes[i]];
         }
         offset[0] = offset_;
@@ -124,6 +124,7 @@ void Transpose::eval_cpu(const std::vector<array>& inputs, array& out) {
 
 void Arange::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 0);
+  (void)inputs;
   out.set_data(allocator::malloc(out.nbytes()));
   switch (out.dtype()) {
     case bool_:
@@ -193,9 +194,9 @@ void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
   flags.row_contiguous = false;
   flags.col_contiguous = false;
   flags.contiguous = false;
-  for (int i = 0; i < inputs.size(); i++) {
+  for (int i = 0; i < std::ssize(inputs); i++) {
     array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
-    size_t data_offset = strides[axis_] * sizes[i];
+    int64_t data_offset = strides[axis_] * sizes[i];
     out_slice.copy_shared_buffer(
         out, strides, flags, out_slice.size(), data_offset);
     copy_cpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, stream());
@@ -205,7 +206,7 @@ void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
 void Contiguous::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  constexpr size_t extra_bytes = 16384;
+  constexpr int64_t extra_bytes = 16384;
   if (in.buffer_size() <= out.nbytes() + extra_bytes &&
       (in.flags().row_contiguous ||
        (allow_col_major_ && in.flags().col_contiguous))) {
@@ -254,8 +255,8 @@ void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
   copy_cpu(val, out, CopyType::Scalar, stream());
 
   // Find offset for start of input values
-  size_t data_offset = 0;
-  for (int i = 0; i < axes_.size(); i++) {
+  int64_t data_offset = 0;
+  for (int i = 0; i < std::ssize(axes_); i++) {
     auto ax = axes_[i] < 0 ? out.ndim() + axes_[i] : axes_[i];
     data_offset += out.strides()[ax] * low_pad_size_[i];
   }
@@ -274,10 +275,10 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
   // keys has shape (N1, ..., NK, 2)
   // out has shape (N1, ..., NK, M1, M2, ...)
   auto& keys = inputs[0];
-  size_t num_keys = keys.size() / 2;
+  int64_t num_keys = keys.size() / 2;
 
-  size_t elems_per_key = out.size() / num_keys;
-  size_t bytes_per_key = out.itemsize() * elems_per_key;
+  int64_t elems_per_key = out.size() / num_keys;
+  int64_t bytes_per_key = out.itemsize() * elems_per_key;
   out.set_data(allocator::malloc(out.nbytes()));
 
   auto kptr = inputs[0].data<uint32_t>();
@@ -291,8 +292,8 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
                     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;
+    int64_t out_skip = (bytes_per_key + 4 - 1) / 4;
+    uintptr_t 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);

mlx/backend/cpu/qrf.cpp

@@ -13,7 +13,7 @@ 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);
+  int64_t num_matrices = a.size() / (M * N);
 
   // Copy A to inplace input and make it col-contiguous
   array in(a.shape(), a.dtype(), nullptr, {});
@@ -54,7 +54,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
     auto work = allocator::malloc(sizeof(T) * lwork);
 
     // Loop over matrices
-    for (int i = 0; i < num_matrices; ++i) {
+    for (int64_t i = 0; i < num_matrices; ++i) {
       // Solve
       geqrf<T>(
           &M,
@@ -68,7 +68,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
     }
     allocator::free(work);
 
-    for (int i = 0; i < num_matrices; ++i) {
+    for (int64_t 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) {
@@ -97,7 +97,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
     work = allocator::malloc(sizeof(T) * lwork);
 
     // Loop over matrices
-    for (int i = 0; i < num_matrices; ++i) {
+    for (int64_t i = 0; i < num_matrices; ++i) {
       // Compute Q
       orgqr<T>(
           &M,
@@ -111,7 +111,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
           &info);
     }
 
-    for (int i = 0; i < num_matrices; ++i) {
+    for (int64_t 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) {

mlx/backend/cpu/quantized.cpp

@@ -445,7 +445,6 @@ void mxfp4_qmm(
     int K) {
   constexpr int group_size = 32;
   constexpr int pack_factor = get_pack_factor(4, 8);
-  constexpr int bytes_per_pack = get_bytes_per_pack(4);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {
@@ -487,7 +486,6 @@ void mxfp4_qmm_t(
     int K) {
   constexpr int group_size = 32;
   constexpr int pack_factor = get_pack_factor(4, 8);
-  constexpr int bytes_per_pack = get_bytes_per_pack(4);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {

mlx/backend/cpu/simd/accelerate_simd.h

@@ -9,7 +9,7 @@
 
 #include "mlx/backend/cpu/simd/base_simd.h"
 
-// There seems to be a bug in sims/base.h
+// There seems to be a bug in simd/base_simd.h
 // __XROS_2_0 is not defined, the expression evaluates
 // to true instead of false setting the SIMD library
 // higher than it should be even on macOS < 15
@@ -253,12 +253,12 @@ Simd<T, N> pow(Simd<T, N> base, Simd<T, N> exp) {
   } else {
     Simd<T, N> res = 1;
     // Raising an integer to a negative power is undefined
-    if (any(exp < 0)) {
+    if (any(exp < static_cast<T>(0))) {
       return 0;
     }
-    while (any(exp > 0)) {
+    while (any(exp > static_cast<T>(0))) {
       res = select((exp & 1) != 0, res * base, res);
-      base = select(exp > 0, base * base, base);
+      base = select(exp > static_cast<T>(0), base * base, base);
       exp = exp >> 1;
     }
     return res;

mlx/backend/cpu/simd/math.h

@@ -79,7 +79,8 @@ Simd<T, N> sincos(Simd<T, N> in) {
 
   // Get the polynom selection mask. There is one polynom for 0 <= x <= Pi/4
   // and another one for Pi/4<x<=Pi/2. Both branches will be computed.
-  auto poly_mask = (emm2 & 2) != 0;
+  auto poly_mask =
+      (emm2 & static_cast<uint32_t>(2)) != static_cast<uint32_t>(0);
 
   // The magic pass: "Extended precision modular arithmetic"
   // x = ((x - y * DP1) - y * DP2) - y * DP3
@@ -87,8 +88,8 @@ Simd<T, N> sincos(Simd<T, N> in) {
   x = fma(y, Simd<float, N>(-2.4187564849853515625e-4f), x);
   x = fma(y, Simd<float, N>(-3.77489497744594108e-8f), x);
 
-  sign_mask_sin = sign_mask_sin ^ ((emm2 & 4) != 0);
-  auto sign_mask_cos = ((emm2 - 2) & 4) != 0;
+  sign_mask_sin = sign_mask_sin ^ ((emm2 & 4) != static_cast<uint32_t>(0));
+  auto sign_mask_cos = ((emm2 - 2) & 4) != static_cast<uint32_t>(0);
 
   // Evaluate the first polynom  (0 <= x <= Pi/4) in y1,
   // and the second polynom      (Pi/4 <= x <= 0) in y2

mlx/backend/cpu/sort.cpp

@@ -15,6 +15,18 @@ namespace mlx::core {
 
 namespace {
 
+// NaN-aware comparator that places NaNs at the end
+template <typename T>
+bool nan_aware_less(T a, T b) {
+  if constexpr (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
+    if (std::isnan(a))
+      return false;
+    if (std::isnan(b))
+      return true;
+  }
+  return a < b;
+}
+
 template <typename T>
 struct StridedIterator {
   using iterator_category = std::random_access_iterator_tag;
@@ -27,7 +39,7 @@ struct StridedIterator {
   StridedIterator() = default;
 
   explicit StridedIterator(T* ptr, int64_t stride, difference_type offset = 0)
-      : ptr_(ptr + offset * stride), stride_(stride) {}
+      : stride_(stride), ptr_(ptr + offset * stride) {}
 
   explicit StridedIterator(array& arr, int axis, difference_type offset = 0)
       : StridedIterator(arr.data<T>(), arr.strides()[axis], offset) {}
@@ -108,8 +120,8 @@ template <typename T>
 void sort(array& out, int axis) {
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + out.ndim() : axis;
-  size_t in_size = out.size();
-  size_t n_rows = in_size / out.shape(axis);
+  int64_t in_size = out.size();
+  int64_t n_rows = in_size / out.shape(axis);
 
   auto remaining_shape = out.shape();
   remaining_shape.erase(remaining_shape.begin() + axis);
@@ -124,13 +136,13 @@ void sort(array& out, int axis) {
   ContiguousIterator src_it(
       remaining_shape, remaining_strides, remaining_shape.size());
   auto out_ptr = out.data<T>();
-  for (int i = 0; i < n_rows; i++) {
+  for (int64_t 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);
 
-    std::stable_sort(st, ed);
+    std::stable_sort(st, ed, nan_aware_less<T>);
     src_it.step();
   }
 }
@@ -139,7 +151,7 @@ template <typename T, typename IdxT = uint32_t>
 void argsort(const array& in, array& out, int axis) {
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + in.ndim() : axis;
-  size_t n_rows = in.size() / in.shape(axis);
+  int64_t n_rows = in.size() / in.shape(axis);
 
   auto in_remaining_shape = in.shape();
   in_remaining_shape.erase(in_remaining_shape.begin() + axis);
@@ -164,7 +176,7 @@ void argsort(const array& in, array& out, int axis) {
       out_remaining_shape, out_remaining_strides, out_remaining_shape.size());
   auto in_ptr = in.data<T>();
   auto out_ptr = out.data<IdxT>();
-  for (int i = 0; i < n_rows; i++) {
+  for (int64_t i = 0; i < n_rows; i++) {
     const T* data_ptr = in_ptr + in_it.loc;
     IdxT* idx_ptr = out_ptr + out_it.loc;
 
@@ -184,6 +196,15 @@ void argsort(const array& in, array& out, int axis) {
     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];
+
+      // Handle NaNs (place them at the end)
+      if (std::is_floating_point<T>::value) {
+        if (std::isnan(v1))
+          return false;
+        if (std::isnan(v2))
+          return true;
+      }
+
       return v1 < v2 || (v1 == v2 && a < b);
     });
   }
@@ -193,8 +214,8 @@ template <typename T>
 void partition(array& out, int axis, int kth) {
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + out.ndim() : axis;
-  size_t in_size = out.size();
-  size_t n_rows = in_size / out.shape(axis);
+  int64_t in_size = out.size();
+  int64_t n_rows = in_size / out.shape(axis);
 
   auto remaining_shape = out.shape();
   remaining_shape.erase(remaining_shape.begin() + axis);
@@ -211,7 +232,7 @@ void partition(array& out, int axis, int kth) {
   ContiguousIterator src_it(
       remaining_shape, remaining_strides, remaining_shape.size());
   auto out_ptr = out.data<T>();
-  for (int i = 0; i < n_rows; i++) {
+  for (int64_t i = 0; i < n_rows; i++) {
     T* data_ptr = out_ptr + src_it.loc;
     src_it.step();
 
@@ -219,7 +240,7 @@ void partition(array& out, int axis, int kth) {
     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, nan_aware_less<T>);
   }
 }
 
@@ -227,7 +248,7 @@ template <typename T, typename IdxT = uint32_t>
 void argpartition(const array& in, array& out, int axis, int kth) {
   // Get axis, shape and stride info
   axis = axis < 0 ? axis + in.ndim() : axis;
-  size_t n_rows = in.size() / in.shape(axis);
+  int64_t n_rows = in.size() / in.shape(axis);
 
   auto in_remaining_shape = in.shape();
   in_remaining_shape.erase(in_remaining_shape.begin() + axis);
@@ -256,7 +277,7 @@ void argpartition(const array& in, array& out, int axis, int kth) {
   auto in_ptr = in.data<T>();
   auto out_ptr = out.data<IdxT>();
 
-  for (int i = 0; i < n_rows; i++) {
+  for (int64_t 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();
@@ -276,6 +297,15 @@ void argpartition(const array& in, array& out, int axis, int kth) {
     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];
+
+      // Handle NaNs (place them at the end)
+      if (std::is_floating_point<T>::value) {
+        if (std::isnan(v1))
+          return false;
+        if (std::isnan(v2))
+          return true;
+      }
+
       return v1 < v2 || (v1 == v2 && a < b);
     });
   }

mlx/backend/cpu/svd.cpp

@@ -27,7 +27,7 @@ void svd_impl(
   const int N = a.shape(-1);
   const int K = std::min(M, N);
 
-  size_t num_matrices = a.size() / (M * N);
+  int64_t num_matrices = a.size() / (M * N);
 
   // lapack clobbers the input, so we have to make a copy.
   array in(a.shape(), a.dtype(), nullptr, {});
@@ -83,8 +83,6 @@ void svd_impl(
 
     auto jobz = (u_ptr) ? "A" : "N";
 
-    // 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
@@ -123,7 +121,7 @@ void svd_impl(
     auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
 
     // Loop over matrices.
-    for (int i = 0; i < num_matrices; i++) {
+    for (int64_t i = 0; i < num_matrices; i++) {
       gesdd<T>(
           /* jobz = */ jobz,
           // M and N are swapped since lapack expects column-major.
@@ -155,10 +153,10 @@ void svd_impl(
 
 template <typename T>
 void compute_svd(
-    const array& a,
-    bool compute_uv,
-    std::vector<array>& outputs,
-    Stream stream) {}
+    const array& /* a */,
+    bool /* compute_uv */,
+    std::vector<array>& /* outputs */,
+    Stream /* stream */) {}
 
 void SVD::eval_cpu(
     const std::vector<array>& inputs,

mlx/backend/cpu/ternary.h

@@ -136,7 +136,7 @@ void ternary_op(
   if (topt == TernaryOpType::ScalarScalarScalar) {
     *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
   } else if (topt == TernaryOpType::VectorVectorVector) {
-    for (size_t i = 0; i < out.size(); ++i) {
+    for (int64_t i = 0; i < out.size(); ++i) {
       *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
       a_ptr++;
       b_ptr++;

mlx/backend/cpu/unary.h

@@ -10,8 +10,8 @@
 namespace mlx::core {
 
 template <typename T, typename U = T, typename Op>
-void unary_op(const T* a, U* out, size_t shape, size_t stride) {
-  for (size_t i = 0; i < shape; i += 1) {
+void unary_op(const T* a, U* out, int64_t shape, int64_t stride) {
+  for (int64_t i = 0; i < shape; i += 1) {
     out[i] = Op{}(*a);
     a += stride;
   }
@@ -38,14 +38,14 @@ void unary_op(const array& a, array& out, Op) {
       src++;
     }
   } else {
-    size_t shape = ndim > 0 ? a.shape().back() : 1;
-    size_t stride = ndim > 0 ? a.strides().back() : 1;
+    int64_t shape = ndim > 0 ? a.shape().back() : 1;
+    int64_t stride = ndim > 0 ? a.strides().back() : 1;
     if (ndim <= 1) {
       unary_op<T, U, Op>(src, dst, shape, stride);
       return;
     }
     auto it = ContiguousIterator(a.shape(), a.strides(), ndim - 1);
-    for (size_t elem = 0; elem < a.size(); elem += shape) {
+    for (int64_t elem = 0; elem < a.size(); elem += shape) {
       unary_op<T, U, Op>(src + it.loc, dst + elem, shape, stride);
       it.step();
     }

mlx/backend/cpu/unary_ops.h

@@ -77,7 +77,8 @@ struct Real {
 struct Sigmoid {
   template <int N, typename T>
   Simd<T, N> operator()(Simd<T, N> x) {
-    return 1.0f / (1.0f + simd::exp(-x));
+    auto y = 1.0f / (1.0f + simd::exp(simd::abs(x)));
+    return simd::select(x < Simd<T, N>{0}, y, Simd<T, N>{1} - y);
   }
   SINGLE()
 };

mlx/backend/cuda/CMakeLists.txt

@@ -170,6 +170,10 @@ target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 # Suppress nvcc warnings on MLX headers.
 target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
                                    --diag_suppress=997>)
+# Supress warnings: note: parameter passing for argument of type
+# ‘std::pair<float, float>’ when C++17 is enabled changed to match C++14 in GCC
+# 10.1
+target_compile_options(mlx PRIVATE -Wno-psabi)
 
 # Install CCCL headers for JIT.
 install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda

mlx/backend/cuda/allocator.cpp

@@ -30,15 +30,20 @@ SmallSizePool::SmallSizePool() {
   next_free_ = buffer_;
 
   CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));
+
+  int device_count = 0;
+  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
+  for (int i = 0; i < device_count; ++i) {
 #if CUDART_VERSION >= 13000
-  cudaMemLocation loc;
-  loc.type = cudaMemLocationTypeDevice;
-  loc.id = 0;
+    cudaMemLocation loc;
+    loc.type = cudaMemLocationTypeDevice;
+    loc.id = i;
 #else
-  int loc = 0;
+    int loc = i;
 #endif // CUDART_VERSION >= 13000
-  CHECK_CUDA_ERROR(
-      cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetReadMostly, loc));
+    CHECK_CUDA_ERROR(
+        cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+  }
 
   auto curr = next_free_;
   for (size_t i = 1; i < num_blocks; ++i) {
@@ -86,7 +91,7 @@ CudaAllocator::CudaAllocator()
   // TODO: Set memory limit for multi-device.
   size_t free, total;
   CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
-  memory_limit_ = total * 0.8;
+  memory_limit_ = total * 0.95;
   max_pool_size_ = memory_limit_;
 }
 

mlx/backend/cuda/compiled.cpp

@@ -332,9 +332,9 @@ void Compiled::eval_gpu(
     encoder.set_output_array(out);
   }
 
-  auto kernel = mod.get_kernel(kernel_name);
+  auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name);
   auto [num_blocks, block_dims] =
-      get_launch_args(outputs[0], large, work_per_thread);
+      get_launch_args(outputs[0], large, work_per_thread, max_block_dims);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
 }
 

mlx/backend/cuda/conv.cpp

@@ -382,20 +382,19 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
   }
 
   if (op_graph) {
-    // Setup inputs and outputs.
-    register_args(encoder, backend_type, in, wt, out, out_);
-
     // Find a plan for the graph and execute it.
     auto plan = find_cudnn_plan_from_op_graph(
         encoder.device().cudnn_handle(), backend_type, dtype, *op_graph);
-    if (!plan) {
-      throw std::runtime_error("[conv] Unable to find an execution plan.");
-    }
-    auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
-    if (encode_cudnn_plan(encoder, *plan, {'x', 'w', 'y'}, x, w, y)) {
-      conv_cache().emplace(
-          cache_key, std::make_pair(backend_type, std::move(*plan)));
-      return;
+    if (plan) {
+      // Setup inputs and outputs.
+      register_args(encoder, backend_type, in, wt, out, out_);
+
+      auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
+      if (encode_cudnn_plan(encoder, *plan, {'x', 'w', 'y'}, x, w, y)) {
+        conv_cache().emplace(
+            cache_key, std::make_pair(backend_type, std::move(*plan)));
+        return;
+      }
     }
   }
 

mlx/backend/cuda/cudnn_utils.cpp

@@ -210,6 +210,9 @@ std::optional<cudnn_frontend::ExecutionPlan> find_cudnn_plan_from_op_graph(
     Dtype dtype,
     cudnn_frontend::OperationGraph& op_graph) {
   auto engine_configs = get_cudnn_engine_configs(backend_type, dtype, op_graph);
+  if (engine_configs.empty()) {
+    return std::nullopt;
+  }
   return find_cudnn_plan_from_engine_configs(handle, engine_configs, op_graph);
 }
 

mlx/backend/cuda/device.cpp

@@ -14,10 +14,6 @@ namespace mlx::core::cu {
 
 namespace {
 
-// Can be tuned with MLX_MAX_OPS_PER_BUFFER
-// This should be less than 255
-constexpr int default_max_nodes_per_graph = 20;
-
 #define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
 
 void check_cudnn_error(const char* name, cudnnStatus_t err) {
@@ -68,8 +64,8 @@ Device::~Device() {
 
 void Device::make_current() {
   // We need to set/get current CUDA device very frequently, cache it to reduce
-  // actual calls of CUDA APIs. This function assumes single-thread in host.
-  static int current = 0;
+  // actual calls of CUDA APIs.
+  static thread_local int current = 0;
   if (current != device_) {
     CHECK_CUDA_ERROR(cudaSetDevice(device_));
     current = device_;
@@ -95,6 +91,7 @@ CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
 
 CommandEncoder::CaptureContext::~CaptureContext() {
   if (!use_cuda_graphs()) {
+    enc.node_count_++;
     return;
   }
 
@@ -196,6 +193,7 @@ CommandEncoder::CommandEncoder(Device& d)
     : device_(d),
       stream_(d),
       graph_(d),
+      worker_(d),
       graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {}
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
@@ -220,12 +218,6 @@ void CommandEncoder::set_output_array(const array& arr) {
   active_outputs_.push_back(id);
 }
 
-void CommandEncoder::maybe_commit() {
-  if (node_count_ >= env::max_ops_per_buffer(default_max_nodes_per_graph)) {
-    commit();
-  }
-}
-
 void CommandEncoder::add_kernel_node(
     void* func,
     dim3 grid_dim,
@@ -233,6 +225,7 @@ void CommandEncoder::add_kernel_node(
     uint32_t smem_bytes,
     void** params) {
   if (!use_cuda_graphs()) {
+    node_count_++;
     CHECK_CUDA_ERROR(cudaLaunchKernel(
         func, grid_dim, block_dim, params, smem_bytes, stream()));
     return;
@@ -253,6 +246,7 @@ void CommandEncoder::add_kernel_node(
     uint32_t smem_bytes,
     void** params) {
   if (!use_cuda_graphs()) {
+    node_count_++;
     CHECK_CUDA_ERROR(cuLaunchKernel(
         func,
         grid_dim.x,
@@ -295,6 +289,7 @@ void CommandEncoder::add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params) {
 
 void CommandEncoder::add_graph_node(cudaGraph_t child) {
   if (!use_cuda_graphs()) {
+    node_count_++;
     CudaGraphExec graph_exec;
     graph_exec.instantiate(child);
     device_.make_current();
@@ -306,12 +301,16 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
   insert_graph_dependencies(GraphNode{node, 'G'});
 }
 
+int CommandEncoder::get_num_ops() {
+  return node_count_;
+}
+
 void CommandEncoder::commit() {
   nvtx3::scoped_range r("CommandEncoder::commit");
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
   }
-  if (node_count_ > 0) {
+  if (use_cuda_graphs() && node_count_ > 0) {
     if (!from_nodes_.empty()) {
       CHECK_CUDA_ERROR(cudaGraphAddDependencies(
           graph_,
@@ -354,7 +353,6 @@ void CommandEncoder::commit() {
     CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
 
     // Reset state
-    node_count_ = 0;
     graph_node_count_ = 0;
     empty_node_count_ = 0;
     from_nodes_.clear();
@@ -366,6 +364,7 @@ void CommandEncoder::commit() {
 
   // Put completion handlers in a batch.
   worker_.commit(stream_);
+  node_count_ = 0;
 }
 
 void CommandEncoder::synchronize() {

mlx/backend/cuda/device.h

@@ -83,7 +83,7 @@ class CommandEncoder {
   }
 
   void add_completed_handler(std::function<void()> task);
-  void maybe_commit();
+  int get_num_ops();
   void commit();
 
   Device& device() {
@@ -140,7 +140,7 @@ class Device {
   Device(const Device&) = delete;
   Device& operator=(const Device&) = delete;
 
-  // Make this device the current cuda device, required by some cuda calls.
+  // Make this device the current cuda device, this method is thread-safe.
   void make_current();
 
   CommandEncoder& get_command_encoder(Stream s);

mlx/backend/cuda/device/unary_ops.cuh

@@ -257,8 +257,8 @@ struct Round {
 struct Sigmoid {
   template <typename T>
   __device__ T operator()(T x) {
-    T y = 1 / (1 + exp(-abs(x)));
-    return (x < 0) ? 1 - y : y;
+    T y = 1 / (1 + exp(abs(x)));
+    return (x < 0) ? y : 1 - y;
   }
 };
 

mlx/backend/cuda/device/utils.cuh

@@ -1,6 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
-// This file must not include any host-only code, utilies that work under both
+// This file must not include any host-only code, utilities that work under both
 // host and device can be put here.
 //
 // See more about the requirements at:
@@ -202,7 +202,7 @@ struct Limits<
   }
 };
 
-// CUDA 11 does not have host side arithmatic operators for half types.
+// CUDA 11 does not have host side arithmetic operators for half types.
 template <typename T>
 struct Limits<
     T,

mlx/backend/cuda/eval.cpp

@@ -5,19 +5,24 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/gpu/available.h"
 #include "mlx/primitives.h"
+#include "mlx/scheduler.h"
 
 #include <nvtx3/nvtx3.hpp>
 
 namespace mlx::core::gpu {
 
+// Can be tuned with MLX_MAX_OPS_PER_BUFFER
+constexpr int default_max_nodes_per_graph = 20;
+
 bool is_available() {
   return true;
 }
 
 void new_stream(Stream s) {
-  // Force initalization of CUDA by creating an event, so the CUDA runtime and
-  // our CUDA event pool get destroyed last.
-  cu::CudaEvent(cudaEventDefault);
+  // Force initalization of CUDA, so CUDA runtime get destroyed at last.
+  cudaFree(nullptr);
+  // Make sure CUDA event pool get destroyed after device and stream.
+  cu::CudaEvent::init_pool();
   // Ensure the static stream objects get created.
   cu::get_command_encoder(s);
 }
@@ -35,7 +40,8 @@ void eval(array& arr) {
     arr.primitive().eval_gpu(arr.inputs(), outputs);
   }
 
-  auto& encoder = cu::get_command_encoder(arr.primitive().stream());
+  auto& stream = arr.primitive().stream();
+  auto& encoder = cu::get_command_encoder(stream);
   // Keep used buffers alive until kernel finishes running.
   for (auto& in : arr.inputs()) {
     // Except for the donated one.
@@ -46,7 +52,14 @@ void eval(array& arr) {
   for (auto& s : arr.siblings()) {
     encoder.add_temporary(s);
   }
-  encoder.maybe_commit();
+
+  if (encoder.get_num_ops() >=
+      env::max_ops_per_buffer(default_max_nodes_per_graph)) {
+    scheduler::notify_new_task(stream);
+    encoder.add_completed_handler(
+        [stream]() { scheduler::notify_task_completion(stream); });
+    encoder.commit();
+  }
 }
 
 void finalize(Stream s) {

mlx/backend/cuda/event.cu

@@ -22,11 +22,15 @@ namespace cu {
 namespace {
 
 // Manage cached cudaEvent_t objects.
-struct CudaEventPool {
-  static CudaEventHandle create(int flags) {
-    auto& cache = cache_for(flags);
+class CudaEventPool {
+ public:
+  CudaEventHandle create(Device& d, int flags) {
+    if (!on_creation_thread()) {
+      return CudaEventHandle(d, flags);
+    }
+    auto& cache = cache_for(d, flags);
     if (cache.empty()) {
-      return CudaEventHandle(flags);
+      return CudaEventHandle(d, flags);
     } else {
       CudaEventHandle ret = std::move(cache.back());
       cache.pop_back();
@@ -34,54 +38,89 @@ struct CudaEventPool {
     }
   }
 
-  static void release(CudaEventHandle event) {
-    cache_for(event.flags).push_back(std::move(event));
+  void release(CudaEventHandle event) {
+    if (!on_creation_thread()) {
+      // Event will be destroyed directly instead of getting moved to cache.
+      return;
+    }
+    cache_for(event.device, event.flags).push_back(std::move(event));
   }
 
-  static std::vector<CudaEventHandle>& cache_for(int flags) {
-    static std::map<int, std::vector<CudaEventHandle>> cache;
-    return cache[flags];
+ private:
+  std::vector<CudaEventHandle>& cache_for(Device& d, int flags) {
+    return cache_[d.cuda_device()][flags];
   }
+
+  bool on_creation_thread() {
+    return std::this_thread::get_id() == thread_id_;
+  }
+
+  // The CudaEvent may be created and destroyed on different threads (for
+  // example when waiting on GPU work in CPU stream), we don't want to make
+  // the cache thread-safe as it adds overhead, so we just skip cache when
+  // using events in worker threads.
+  std::thread::id thread_id_{std::this_thread::get_id()};
+
+  // {device: {flags: [events]}}
+  std::map<int, std::map<int, std::vector<CudaEventHandle>>> cache_;
 };
 
+CudaEventPool& cuda_event_pool() {
+  static CudaEventPool pool;
+  return pool;
+}
+
 } // namespace
 
-CudaEventHandle::CudaEventHandle(int flags) : flags(flags) {
+CudaEventHandle::CudaEventHandle(Device& d, int flags)
+    : device(d), flags(flags) {
+  device.make_current();
   CHECK_CUDA_ERROR(cudaEventCreateWithFlags(&handle_, flags));
   assert(handle_ != nullptr);
 }
 
-CudaEvent::CudaEvent(int flags) : event_(CudaEventPool::create(flags)) {}
+CudaEvent::CudaEvent(Device& d, int flags)
+    : event_(cuda_event_pool().create(d, flags)) {}
 
 CudaEvent::~CudaEvent() {
-  CudaEventPool::release(std::move(event_));
+  cuda_event_pool().release(std::move(event_));
 }
 
 void CudaEvent::wait() {
   nvtx3::scoped_range r("cu::CudaEvent::wait");
+  event_.device.make_current();
   cudaEventSynchronize(event_);
 }
 
 void CudaEvent::wait(cudaStream_t stream) {
+  event_.device.make_current();
   cudaStreamWaitEvent(stream, event_);
 }
 
 void CudaEvent::record(cudaStream_t stream) {
+  event_.device.make_current();
   cudaEventRecord(event_, stream);
 }
 
 bool CudaEvent::completed() const {
+  // Note: cudaEventQuery can be safely called from any device.
   return cudaEventQuery(event_) == cudaSuccess;
 }
 
+// static
+void CudaEvent::init_pool() {
+  cuda_event_pool();
+}
+
 // Wraps CudaEvent with a few features:
 // 1. The class can be copied.
 // 2. Make wait/record work with CPU streams.
 // 3. Add checks for waiting on un-recorded event.
 class CopyableCudaEvent {
  public:
-  CopyableCudaEvent()
+  explicit CopyableCudaEvent(Device& d)
       : event_(std::make_shared<CudaEvent>(
+            d,
             cudaEventDisableTiming | cudaEventBlockingSync)) {}
 
   void wait() {
@@ -245,7 +284,7 @@ struct EventImpl {
       nvtx3::mark("Using slow AtomicEvent");
       atomic = std::make_unique<cu::AtomicEvent>();
     } else {
-      cuda = std::make_unique<cu::CopyableCudaEvent>();
+      cuda = std::make_unique<cu::CopyableCudaEvent>(cu::device(s.device));
     }
   }
 };

mlx/backend/cuda/event.h

@@ -13,9 +13,12 @@
 
 namespace mlx::core::cu {
 
+class Device;
+
 // RAII-managed move-only wrapper of cudaEvent_t.
 struct CudaEventHandle : public CudaHandle<cudaEvent_t, cudaEventDestroy> {
-  CudaEventHandle(int flags);
+  CudaEventHandle(Device& d, int flags);
+  Device& device;
   int flags;
 };
 
@@ -23,7 +26,7 @@ struct CudaEventHandle : public CudaHandle<cudaEvent_t, cudaEventDestroy> {
 // on GPU stream in CPU stream, but can not wait on CPU stream.
 class CudaEvent {
  public:
-  explicit CudaEvent(int flags);
+  CudaEvent(Device& d, int flags);
   ~CudaEvent();
 
   CudaEvent(CudaEvent&&) = default;
@@ -40,6 +43,9 @@ class CudaEvent {
   // returns true if record() has not been called.
   bool completed() const;
 
+  // Internal: make sure event pool is initialized.
+  static void init_pool();
+
  private:
   CudaEventHandle event_;
 };

mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -50,8 +50,10 @@ cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
       return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
                                            : CUBLAS_COMPUTE_32F;
     case float64:
-    case complex64:
       return CUBLAS_COMPUTE_64F;
+    case complex64:
+      return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
+                                           : CUBLAS_COMPUTE_32F;
     default:
       throw std::runtime_error(fmt::format(
           "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
@@ -126,12 +128,13 @@ CublasGemm::CublasGemm(
       N_(b_cols) {
   heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
 
-  auto scale_type = dtype_to_cublas_type(dtype);
+  scale_type_ = dtype_to_cublas_type(dtype);
   if (dtype == bfloat16 || dtype == float16) {
-    scale_type = CUDA_R_32F;
+    scale_type_ = CUDA_R_32F;
   }
+
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescCreate(
-      &matmul_desc_, dtype_to_compute_type(dtype), scale_type));
+      &matmul_desc_, dtype_to_compute_type(dtype), scale_type_));
   int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
       matmul_desc_,
@@ -352,6 +355,16 @@ void CublasGemm::execute(
     }
   }
 
+  const void* alpha_ptr = α
+  const void* beta_ptr = β
+  complex64_t alpha_c, beta_c;
+  if (scale_type_ == CUDA_C_32F) {
+    alpha_c = complex64_t{alpha, 0.0f};
+    beta_c = complex64_t{beta, 0.0f};
+    alpha_ptr = &alpha_c;
+    beta_ptr = &beta_c;
+  }
+
   void* workspace_ptr = nullptr;
   if (heuristic_.workspaceSize > 0) {
     // Ensure workspace is 256-byte aligned
@@ -368,12 +381,12 @@ void CublasGemm::execute(
   CHECK_CUBLAS_ERROR(cublasLtMatmul(
       handle_,
       matmul_desc_,
-      &alpha,
+      alpha_ptr,
       b, // a and b are swapped
       a_desc_,
       a,
       b_desc_,
-      &beta,
+      beta_ptr,
       c ? c : out,
       c ? c_desc_ : out_desc_,
       out,

mlx/backend/cuda/gemms/cublas_gemm.h

@@ -115,6 +115,7 @@ class CublasGemm {
 
   uint64_t M_;
   uint64_t N_;
+  cudaDataType_t scale_type_;
   cublasLtMatmulPreference_t pref_{nullptr};
   cublasLtHandle_t handle_{nullptr};
   cublasLtMatmulDesc_t matmul_desc_{nullptr};

mlx/backend/cuda/gemms/gemv.cu

@@ -13,6 +13,37 @@ namespace cg = cooperative_groups;
 
 static constexpr int rows_per_block = 8;
 
+// Accumulator type selection per input element type T.
+template <typename T>
+struct GemvAccType {
+  using type = T;
+};
+
+template <>
+struct GemvAccType<__half> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<__nv_bfloat16> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<float> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<double> {
+  using type = double;
+};
+
+template <>
+struct GemvAccType<cu::complex64_t> {
+  using type = cu::complex64_t;
+};
+
 template <typename T, int rows_per_block, int n_per_thread>
 __device__ void
 gemv_impl(const T* mat, const T* vec, T* out, int rows, int cols) {
@@ -24,7 +55,8 @@ gemv_impl(const T* mat, const T* vec, T* out, int rows, int cols) {
   int row = g_idx.x * rows_per_block + t_idx.y;
 
   if (row < rows) {
-    float sum = 0.0f;
+    using Acc = typename GemvAccType<T>::type;
+    Acc sum = Acc(0);
     for (int col = n_per_thread * warp.thread_rank(); col < cols;
          col += (WARP_SIZE * n_per_thread)) {
       auto local_mat =
@@ -32,12 +64,11 @@ gemv_impl(const T* mat, const T* vec, T* out, int rows, int cols) {
       auto local_vec = unsafe_load_vector<n_per_thread>(vec + col, 0);
 #pragma unroll
       for (int j = 0; j < n_per_thread; ++j) {
-        sum +=
-            static_cast<float>(local_mat[j]) * static_cast<float>(local_vec[j]);
+        sum += static_cast<Acc>(local_mat[j]) * static_cast<Acc>(local_vec[j]);
       }
     }
 
-    sum = cg::reduce(warp, sum, cg::plus<float>{});
+    sum = cg::reduce(warp, sum, cg::plus<Acc>{});
     if (warp.thread_rank() == 0) {
       out[row] = static_cast<T>(sum);
     }
@@ -107,7 +138,7 @@ void gemv(
   encoder.set_input_array(a);
   encoder.set_input_array(b);
   encoder.set_output_array(out);
-  dispatch_float_types(out.dtype(), "gemv", [&](auto type_tag) {
+  dispatch_inexact_types(out.dtype(), "gemv", [&](auto type_tag) {
     using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
     dim3 block_dims{WARP_SIZE, rows_per_block};
     const DataType* mat;

mlx/backend/cuda/jit_module.cpp

@@ -99,6 +99,30 @@ const std::filesystem::path& ptx_cache_dir() {
   return cache;
 }
 
+std::filesystem::path get_ptx_path(
+    const std::filesystem::path& cache_dir,
+    const std::string& module_name) {
+#ifdef _WIN32
+  constexpr int max_file_name_length = 140;
+#else
+  constexpr int max_file_name_length = 245;
+#endif
+
+  if (module_name.size() <= max_file_name_length) {
+    return cache_dir / (module_name + ".ptx");
+  }
+
+  auto ptx_path = cache_dir;
+  int offset = 0;
+  while (module_name.size() - offset > max_file_name_length) {
+    ptx_path /= module_name.substr(offset, max_file_name_length);
+    offset += max_file_name_length;
+  }
+  ptx_path /= module_name.substr(offset) + ".ptx";
+
+  return ptx_path;
+}
+
 // Try to read the cached |ptx| and |ptx_kernels| from |cache_dir|.
 bool read_cached_ptx(
     const std::filesystem::path& cache_dir,
@@ -109,7 +133,7 @@ bool read_cached_ptx(
     return false;
   }
 
-  auto ptx_path = cache_dir / (module_name + ".ptx");
+  auto ptx_path = get_ptx_path(cache_dir, module_name);
   std::error_code error;
   auto ptx_size = std::filesystem::file_size(ptx_path, error);
   if (error) {
@@ -122,7 +146,7 @@ bool read_cached_ptx(
   ptx.resize(ptx_size);
   ptx_file.read(ptx.data(), ptx_size);
 
-  std::ifstream txt_file(cache_dir / (module_name + ".txt"), std::ios::binary);
+  std::ifstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary);
   std::string line;
   while (std::getline(txt_file, line)) {
     auto tab = line.find('\t');
@@ -144,16 +168,26 @@ void write_cached_ptx(
     return;
   }
 
-  std::ofstream ptx_file(cache_dir / (module_name + ".ptx"), std::ios::binary);
+  auto ptx_path = get_ptx_path(cache_dir, module_name);
+
+  // Ensure that the directory exists
+  auto parent = ptx_path.parent_path();
+  if (parent != cache_dir) {
+    std::filesystem::create_directories(parent);
+  }
+
+  // Write the compiled code and mangled names
+  std::ofstream ptx_file(ptx_path, std::ios::binary);
   if (!ptx.empty()) {
     ptx_file.write(&ptx.front(), ptx.size());
   }
-  std::ofstream txt_file(cache_dir / (module_name + ".txt"), std::ios::binary);
+  std::ofstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary);
   for (const auto& [name, mangled] : ptx_kernels) {
     txt_file << name << "\t" << mangled << std::endl;
   }
 
-  std::ofstream source_file(cache_dir / (module_name + ".cu"));
+  // Write the generated code
+  std::ofstream source_file(ptx_path.replace_extension(".cu"));
   source_file << source_code;
 }
 
@@ -297,7 +331,8 @@ void load_module(
     const std::string& ptx,
     const std::vector<std::pair<std::string, std::string>>& ptx_kernels,
     CUmodule& module_,
-    std::unordered_map<std::string, std::pair<CUfunction, bool>>& kernels) {
+    std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>>&
+        kernels) {
   // Load module.
   char jit_log[4089] = {};
   CUjit_option options[] = {
@@ -314,7 +349,7 @@ void load_module(
   for (const auto& [name, mangled] : ptx_kernels) {
     CUfunction kernel;
     CHECK_CUDA_ERROR(cuModuleGetFunction(&kernel, module_, mangled.c_str()));
-    kernels[name] = std::make_pair(kernel, false);
+    kernels[name] = std::make_tuple(kernel, false, 0);
   }
 }
 
@@ -358,7 +393,7 @@ JitModule::~JitModule() {
   CHECK_CUDA_ERROR(cuModuleUnload(module_));
 }
 
-CUfunction JitModule::get_kernel(
+std::pair<CUfunction, uint> JitModule::get_kernel_and_dims(
     const std::string& kernel_name,
     std::function<void(CUfunction)> configure_kernel) {
   auto it = kernels_.find(kernel_name);
@@ -369,14 +404,22 @@ CUfunction JitModule::get_kernel(
 
   // If it is the first time we run this kernel then configure it. Do it only
   // once!
-  if (!it->second.second) {
+  auto kernel = std::get<0>(it->second);
+  if (!std::get<1>(it->second)) {
     if (configure_kernel) {
-      configure_kernel(it->second.first);
+      configure_kernel(kernel);
     }
-    it->second.second = true;
+    std::get<1>(it->second) = true;
+    std::get<2>(it->second) = max_occupancy_block_dim(kernel);
   }
 
-  return it->second.first;
+  return {kernel, std::get<2>(it->second)};
+}
+
+CUfunction JitModule::get_kernel(
+    const std::string& kernel_name,
+    std::function<void(CUfunction)> configure_kernel) {
+  return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first;
 }
 
 std::unordered_map<std::string, JitModule>& get_jit_module_cache() {

mlx/backend/cuda/jit_module.h

@@ -99,10 +99,13 @@ class JitModule {
   CUfunction get_kernel(
       const std::string& kernel_name,
       std::function<void(CUfunction)> configure_kernel = nullptr);
+  std::pair<CUfunction, uint> get_kernel_and_dims(
+      const std::string& kernel_name,
+      std::function<void(CUfunction)> configure_kernel = nullptr);
 
  private:
   CUmodule module_{nullptr};
-  std::unordered_map<std::string, std::pair<CUfunction, bool>> kernels_;
+  std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>> kernels_;
 };
 
 std::unordered_map<std::string, JitModule>& get_jit_module_cache();

mlx/backend/cuda/kernel_utils.cu

@@ -35,12 +35,10 @@ std::tuple<dim3, uint> get_launch_args(
     const Shape& shape,
     const Strides& strides,
     bool large,
-    int work_per_thread) {
+    int work_per_thread /* = 1 */,
+    uint max_block_dim /* = 1024 */) {
   size_t nthreads = cuda::ceil_div(size, work_per_thread);
-  uint block_dim = 1024;
-  if (block_dim > nthreads) {
-    block_dim = nthreads;
-  }
+  uint block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
   dim3 num_blocks;
   if (large) {
     num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);

mlx/backend/cuda/kernel_utils.cuh

@@ -1,8 +1,8 @@
 // Copyright © 2025 Apple Inc.
 
-// This file includes host-only utilies for writing CUDA kernels, the difference
-// from backend/cuda/device/utils.cuh is that the latter file only include
-// device-only code.
+// This file includes host-only utilities for writing CUDA kernels, the
+// difference from backend/cuda/device/utils.cuh is that the latter file only
+// include device-only code.
 
 #pragma once
 
@@ -120,19 +120,28 @@ dim3 get_2d_grid_dims(
     size_t divisor);
 std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);
 
-// Get the num_blocks and block_dims that maximize occupancy for |kernel|,
-// assuming each thread handles |work_per_thread| elements of |arr|.
+// Get the num_blocks and block_dims assuming each thread handles
+// |work_per_thread| elements of |arr|.
 std::tuple<dim3, uint> get_launch_args(
     size_t size,
     const Shape& shape,
     const Strides& strides,
     bool large,
-    int work_per_thread = 1);
+    int work_per_thread = 1,
+    uint max_block_dim = 1024);
 
-inline std::tuple<dim3, uint>
-get_launch_args(const array& arr, bool large, int work_per_thread = 1) {
+inline std::tuple<dim3, uint> get_launch_args(
+    const array& arr,
+    bool large,
+    int work_per_thread = 1,
+    uint max_block_dim = 1024) {
   return get_launch_args(
-      arr.size(), arr.shape(), arr.strides(), large, work_per_thread);
+      arr.size(),
+      arr.shape(),
+      arr.strides(),
+      large,
+      work_per_thread,
+      max_block_dim);
 }
 
 } // namespace mlx::core

mlx/backend/cuda/matmul.cpp

@@ -93,6 +93,10 @@ void gemm_and_bias(
       a_batch_strides.back(),
       b_batch_strides.back());
   if (bias) {
+    if (a.dtype() == complex64) {
+      throw std::runtime_error(
+          "[gemm_and_bias] complex64 bias epilogue isn’t supported in cublasLtMatmul.");
+    }
     gemm.set_bias(encoder, *bias);
   }
   gemm.run(
@@ -157,7 +161,8 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   /////////////////////////////////////////////////////////////////////////////
   // Dispatch to GEMM with epilogue or AddMM
 
-  if (beta_ == 1 && c.strides(-1) == 1 && c.data_size() == out.shape(-1)) {
+  if (beta_ == 1 && a.dtype() != complex64 && c.strides(-1) == 1 &&
+      c.data_size() == out.shape(-1)) {
     out.set_data(allocator::malloc(out.nbytes()));
     gemm_and_bias(
         encoder,

mlx/backend/cuda/no_cuda.cpp

@@ -35,9 +35,9 @@ std::vector<array> precompiled_cuda_kernel(
     const std::vector<ScalarArg>&,
     std::tuple<int, int, int>,
     std::tuple<int, int, int>,
-    int shared_memory,
-    std::optional<float> init_value,
-    bool ensure_row_contiguous,
+    int /* shared_memory */,
+    std::optional<float> /* init_value */,
+    bool /* ensure_row_contiguous */,
     StreamOrDevice) {
   throw std::runtime_error("[cuda_kernel] No CUDA back-end.");
 }

mlx/backend/cuda/reduce/col_reduce.cu

@@ -181,6 +181,47 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
   }
 }
 
+template <typename T, typename U, typename Op, int N_READS = 4>
+__global__ void col_reduce_small(
+    const T* in,
+    U* out,
+    const __grid_constant__ ColReduceArgs args,
+    size_t total) {
+  Op op;
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  const auto idx = grid.thread_rank() * N_READS;
+  const auto before_axis = idx / args.reduction_stride;
+  const auto after_axis = idx % args.reduction_stride;
+  const auto offset =
+      before_axis * args.reduction_stride * args.reduction_size + after_axis;
+
+  if (idx >= total) {
+    return;
+  }
+
+  in += offset;
+  out += idx;
+
+  AlignedVector<U, N_READS> accumulator;
+  for (int i = 0; i < N_READS; i++) {
+    accumulator[i] = ReduceInit<Op, T>::value();
+  }
+
+  for (int i = 0; i < args.reduction_size; i++) {
+    auto values = load_vector<N_READS>(in, 0);
+
+    for (int j = 0; j < N_READS; j++) {
+      accumulator[j] = op(accumulator[j], cast_to<U>(values[j]));
+    }
+
+    in += args.reduction_stride;
+  }
+
+  store_vector(out, 0, accumulator);
+}
+
 } // namespace cu
 
 inline auto output_grid_for_col_reduce(
@@ -206,7 +247,7 @@ void col_reduce_looped(
     Reduce::ReduceType reduce_type,
     const std::vector<int>& axes,
     const ReductionPlan& plan,
-    cu::ColReduceArgs args) {
+    const cu::ColReduceArgs& args) {
   // Allocate data for the output using in's layout to access them as
   // contiguously as possible.
   allocate_same_layout(out, in, axes);
@@ -230,12 +271,55 @@ void col_reduce_looped(
         auto kernel =
             cu::col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS>;
         encoder.add_kernel_node(
-            kernel, grid, blocks, 0, indata, out.data<U>(), args);
+            kernel,
+            grid,
+            blocks,
+            0,
+            indata,
+            out.data<U>(),
+            static_cast<cu::ColReduceArgs>(args));
       });
     });
   });
 }
 
+void col_reduce_small(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    const cu::ColReduceArgs& args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes);
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+
+      constexpr int N_READS = 16 / sizeof(T);
+      auto tmp_grid = get_2d_grid_dims(out.shape(), out.strides());
+      auto [grid, block] = get_grid_and_block(tmp_grid.x, tmp_grid.y, 1);
+      auto kernel = cu::col_reduce_small<T, U, OP, N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          grid,
+          block,
+          0,
+          in.data<T>(),
+          out.data<U>(),
+          static_cast<cu::ColReduceArgs>(args),
+          out.size());
+    });
+  });
+}
+
 void col_reduce(
     cu::CommandEncoder& encoder,
     const array& in,
@@ -258,6 +342,13 @@ void col_reduce(
   // Make the args struct to help route to the best kernel
   cu::ColReduceArgs args(in, plan, axes);
 
+  // Small col reduce with a single or contiguous reduction axis
+  if (args.non_col_reductions == 1 && args.reduction_size <= 32 &&
+      args.reduction_stride % (16 / in.itemsize()) == 0) {
+    col_reduce_small(encoder, in, out, reduce_type, axes, plan, args);
+    return;
+  }
+
   // Fallback col reduce
   col_reduce_looped(encoder, in, out, reduce_type, axes, plan, args);
 }

mlx/backend/cuda/reduce/row_reduce.cu

@@ -7,8 +7,6 @@
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
-#include <cub/block/block_load.cuh>
-#include <cub/block/block_reduce.cuh>
 
 namespace mlx::core {
 
@@ -83,7 +81,8 @@ struct RowReduceArgs {
 };
 
 template <typename T, typename U, typename ReduceOp, int N = 4, int M = 1>
-__global__ void row_reduce_simple(T* in, U* out, size_t n_rows, int size) {
+__global__ void
+row_reduce_simple(const T* in, U* out, size_t n_rows, int size) {
   auto grid = cg::this_grid();
   auto block = cg::this_thread_block();
   auto warp = cg::tiled_partition<WARP_SIZE>(block);
@@ -91,8 +90,8 @@ __global__ void row_reduce_simple(T* in, U* out, size_t n_rows, int size) {
   const U init = cu::ReduceInit<ReduceOp, T>::value();
   ReduceOp op;
 
-  T vals[M][N];
-  U accs[M];
+  AlignedVector<T, N> vals[M];
+  AlignedVector<U, M> accs;
   for (int i = 0; i < M; i++) {
     accs[i] = init;
   }
@@ -101,43 +100,31 @@ __global__ void row_reduce_simple(T* in, U* out, size_t n_rows, int size) {
       min(n_rows - M, static_cast<size_t>(grid.block_rank() * M));
   const size_t full_blocks = size / (block.size() * N);
   const size_t final_offset = full_blocks * (block.size() * N);
-  in += start_row * size;
+  in += start_row * size + block.thread_rank() * N;
   out += start_row;
 
-  if (size % N == 0) {
-    for (size_t r = 0; r < full_blocks; r++) {
-      for (int k = 0; k < M; k++) {
-        cub::LoadDirectBlockedVectorized<T, N>(
-            block.thread_rank(),
-            in + k * size + r * (block.size() * N),
-            vals[k]);
-        for (int j = 0; j < N; j++) {
-          accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
-        }
-      }
-    }
-  } else {
-    for (size_t r = 0; r < full_blocks; r++) {
-      for (int k = 0; k < M; k++) {
-        cub::LoadDirectBlocked(
-            block.thread_rank(),
-            in + k * size + r * (block.size() * N),
-            vals[k]);
-        for (int j = 0; j < N; j++) {
-          accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
-        }
+  for (size_t r = 0; r < full_blocks; r++) {
+    for (int k = 0; k < M; k++) {
+      vals[k] = load_vector<N>(in + k * size, 0);
+    }
+    for (int k = 0; k < M; k++) {
+      for (int j = 0; j < N; j++) {
+        accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
       }
     }
+
+    in += block.size() * N;
   }
 
   if (final_offset < size) {
     for (int k = 0; k < M; k++) {
-      cub::LoadDirectBlocked(
-          block.thread_rank(),
-          in + k * size + final_offset,
-          vals[k],
-          size,
-          cast_to<T>(init));
+      for (int i = 0; i < N; i++) {
+        vals[k][i] = ((final_offset + block.thread_rank() * N + i) < size)
+            ? in[k * size + i]
+            : cast_to<T>(init);
+      }
+    }
+    for (int k = 0; k < M; k++) {
       for (int j = 0; j < N; j++) {
         accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
       }
@@ -145,13 +132,11 @@ __global__ void row_reduce_simple(T* in, U* out, size_t n_rows, int size) {
   }
 
   __shared__ U shared_accumulators[32 * M];
-  block_reduce(block, warp, accs, shared_accumulators, op, init);
+  block_reduce(block, warp, accs.val, shared_accumulators, op, init);
 
   if (block.thread_rank() == 0) {
     if (grid.block_rank() * M + M <= n_rows) {
-      for (int i = 0; i < M; i++) {
-        out[i] = accs[i];
-      }
+      store_vector(out, 0, accs);
     } else {
       short offset = grid.block_rank() * M + M - n_rows;
       for (int i = offset; i < M; i++) {
@@ -161,17 +146,10 @@ __global__ void row_reduce_simple(T* in, U* out, size_t n_rows, int size) {
   }
 }
 
-template <
-    typename T,
-    typename U,
-    typename Op,
-    int NDIM,
-    int BLOCK_DIM,
-    int N_READS = 4>
+template <typename T, typename U, typename Op, int NDIM, int N_READS = 4>
 __global__ void row_reduce_looped(
-    T* in,
+    const T* in,
     U* out,
-    size_t out_size,
     const __grid_constant__ RowReduceArgs args) {
   auto grid = cg::this_grid();
   auto block = cg::this_thread_block();
@@ -185,36 +163,60 @@ __global__ void row_reduce_looped(
   U init = ReduceInit<Op, T>::value();
   total[0] = init;
   LoopedElemToLoc<NDIM, (NDIM > 2)> loop(args.reduce_ndim);
-  size_t full_blocks = args.row_size / (BLOCK_DIM * N_READS);
-  size_t final_offset = full_blocks * BLOCK_DIM * N_READS;
+  const size_t full_blocks = args.row_size / (block.size() * N_READS);
+  const size_t final_offset = full_blocks * (block.size() * N_READS);
 
   in += elem_to_loc(out_idx, args.shape.data(), args.strides.data(), args.ndim);
+  in += block.thread_rank() * N_READS;
 
-  for (size_t n = 0; n < args.non_row_reductions; n++) {
-    for (size_t r = 0; r < full_blocks; r++) {
-      T vals[N_READS];
-      cub::LoadDirectBlockedVectorized<T, N_READS>(
-          block.thread_rank(),
-          in + loop.location() + r * BLOCK_DIM * N_READS,
-          vals);
-      for (int i = 0; i < N_READS; i++) {
-        total[0] = op(total[0], cast_to<U>(vals[i]));
-      }
+  // Unaligned reduce
+  if (final_offset < args.row_size) {
+    bool mask[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      mask[i] =
+          (final_offset + block.thread_rank() * N_READS + i) < args.row_size;
     }
-    if (final_offset < args.row_size) {
-      T vals[N_READS];
-      cub::LoadDirectBlocked(
-          block.thread_rank(),
-          in + loop.location() + final_offset,
-          vals,
-          args.row_size - final_offset,
-          cast_to<T>(init));
-      for (int i = 0; i < N_READS; i++) {
-        total[0] = op(total[0], cast_to<U>(vals[i]));
+
+    for (size_t n = 0; n < args.non_row_reductions; n++) {
+      const T* inlocal = in + loop.location();
+
+      for (size_t r = 0; r < full_blocks; r++) {
+        auto vals = load_vector<N_READS>(inlocal, 0);
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+        inlocal += block.size() * N_READS;
       }
+
+      {
+        T vals[N_READS];
+        for (int i = 0; i < N_READS; i++) {
+          vals[i] = mask[i] ? inlocal[i] : cast_to<T>(init);
+        }
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+      }
+
+      loop.next(args.reduce_shape.data(), args.reduce_strides.data());
+    }
+  }
+
+  // Aligned case
+  else {
+    for (size_t n = 0; n < args.non_row_reductions; n++) {
+      const T* inlocal = in + loop.location();
+
+      for (size_t r = 0; r < full_blocks; r++) {
+        auto vals = load_vector<N_READS>(inlocal, 0);
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+        inlocal += block.size() * N_READS;
+      }
+
+      loop.next(args.reduce_shape.data(), args.reduce_strides.data());
     }
-    // TODO: Maybe block.sync() here?
-    loop.next(args.reduce_shape.data(), args.reduce_strides.data());
   }
 
   __shared__ U shared_accumulators[32];
@@ -234,8 +236,6 @@ void row_reduce_simple(
     Reduce::ReduceType reduce_type,
     const std::vector<int>& axes,
     const ReductionPlan& plan) {
-  constexpr int N_READS = 8;
-
   // Allocate data for the output using in's layout to avoid elem_to_loc in the
   // kernel.
   allocate_same_layout(out, in, axes);
@@ -250,14 +250,15 @@ void row_reduce_simple(
       using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       using U = typename cu::ReduceResult<OP, T>::type;
 
-      // Cub doesn't like const pointers for vectorized loads. (sigh)
-      T* indata = const_cast<T*>(in.data<T>());
+      constexpr int N_READS = 16 / sizeof(T);
 
       // Calculate the grid and block dims
       size_t reductions = (plan.shape.back() + N_READS - 1) / N_READS;
       dim3 grid = get_2d_grid_dims(out.shape(), out.strides());
-      int threads = std::min(1024UL, reductions);
-      threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+      int warps = (reductions + WARP_SIZE - 1) / WARP_SIZE;
+      warps /= 4;
+      warps = std::max(std::min(warps, 32), 1);
+      int threads = warps * WARP_SIZE;
       dim3 block(threads, 1, 1);
 
       // Pick the kernel
@@ -267,6 +268,7 @@ void row_reduce_simple(
         kernel = cu::row_reduce_simple<T, U, OP, N_READS, 2>;
       }
 
+      T* indata = const_cast<T*>(in.data<T>());
       int size = plan.shape.back();
       encoder.add_kernel_node(
           kernel, grid, block, 0, indata, out.data<U>(), out.size(), size);
@@ -282,8 +284,6 @@ void row_reduce_looped(
     const std::vector<int>& axes,
     const ReductionPlan& plan,
     cu::RowReduceArgs args) {
-  constexpr int N_READS = 8;
-
   // Allocate data for the output using in's layout to access them as
   // contiguously as possible.
   allocate_same_layout(out, in, axes);
@@ -295,34 +295,27 @@ void row_reduce_looped(
       using OP = MLX_GET_TYPE(reduce_type_tag);
       using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       using U = typename cu::ReduceResult<OP, T>::type;
-      // Cub doesn't like const pointers for vectorized loads. (sigh)
-      T* indata = const_cast<T*>(in.data<T>());
+
+      constexpr int N_READS = 16 / sizeof(T);
 
       // Calculate the grid and block dims
       args.sort_access_pattern(in, axes);
       dim3 grid = get_2d_grid_dims(out.shape(), out.strides());
       size_t reductions = (args.row_size + N_READS - 1) / N_READS;
-      int threads = std::min(1024UL, reductions);
-      threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+      int warps = (reductions + WARP_SIZE - 1) / WARP_SIZE;
+      warps /= 4;
+      warps = std::max(std::min(warps, 32), 1);
+      int threads = warps * WARP_SIZE;
       dim3 block(threads, 1, 1);
 
       // Pick the kernel
-      auto kernel = cu::row_reduce_looped<T, U, OP, 1, 32, N_READS>;
+      auto kernel = cu::row_reduce_looped<T, U, OP, 1, N_READS>;
       dispatch_reduce_ndim(args.reduce_ndim, [&](auto reduce_ndim) {
-        dispatch_block_dim(threads, [&](auto threads_constant) {
-          kernel = cu::row_reduce_looped<
-              T,
-              U,
-              OP,
-              reduce_ndim.value,
-              threads_constant.value,
-              N_READS>;
-          block.x = threads_constant.value;
-        });
+        kernel = cu::row_reduce_looped<T, U, OP, reduce_ndim.value, N_READS>;
       });
 
       encoder.add_kernel_node(
-          kernel, grid, block, 0, indata, out.data<U>(), out.size(), args);
+          kernel, grid, block, 0, in.data<T>(), out.data<U>(), args);
     });
   });
 }

mlx/backend/cuda/ternary.cu

@@ -156,7 +156,25 @@ void ternary_op_gpu_inplace(
     using DType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
 
     auto topt = get_ternary_op_type(a, b, c);
-    if (topt == TernaryOpType::General) {
+    if (topt == TernaryOpType::VectorVectorVector ||
+        topt == TernaryOpType::ScalarScalarScalar) {
+      dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
+        using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+        constexpr int N_READS = 16 / sizeof(DType);
+        auto [num_blocks, block_dims] = get_launch_args(
+            out.data_size(), out.shape(), out.strides(), large(), N_READS);
+        encoder.add_kernel_node(
+            cu::ternary_v<Op, DType, IdxT, N_READS>,
+            num_blocks,
+            block_dims,
+            0,
+            a.data<bool>(),
+            b.data<DType>(),
+            c.data<DType>(),
+            out.data<DType>(),
+            out.data_size());
+      });
+    } else {
       dispatch_bool(
           a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
               c.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
@@ -225,23 +243,6 @@ void ternary_op_gpu_inplace(
                   ndim);
             }
           });
-    } else {
-      dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
-        using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-        constexpr int N_READS = 16 / sizeof(DType);
-        auto [num_blocks, block_dims] = get_launch_args(
-            out.data_size(), out.shape(), out.strides(), large(), N_READS);
-        encoder.add_kernel_node(
-            cu::ternary_v<Op, DType, IdxT, N_READS>,
-            num_blocks,
-            block_dims,
-            0,
-            a.data<bool>(),
-            b.data<DType>(),
-            c.data<DType>(),
-            out.data<DType>(),
-            out.data_size());
-      });
     }
   });
 }

mlx/backend/cuda/unary.cu

@@ -1,284 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/common/unary.h"
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/device/unary_ops.cuh"
-#include "mlx/backend/cuda/kernel_utils.cuh"
-#include "mlx/dtype_utils.h"
-#include "mlx/primitives.h"
-
-#include <cooperative_groups.h>
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace cu {
-
-namespace cg = cooperative_groups;
-
-template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void unary_v(const In* in, Out* out, IdxT size) {
-  IdxT index = cg::this_grid().thread_rank();
-
-  if ((index + 1) * N_READS > size) {
-    for (IdxT i = index * N_READS; i < size; ++i) {
-      out[i] = Op{}(in[i]);
-    }
-  } else {
-    auto in_vec = load_vector<N_READS>(in, index);
-
-    AlignedVector<Out, N_READS> out_vec;
-#pragma unroll
-    for (int i = 0; i < N_READS; ++i) {
-      out_vec[i] = Op{}(in_vec[i]);
-    }
-
-    store_vector<N_READS>(out, index, out_vec);
-  }
-}
-
-template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void unary_g(
-    const In* in,
-    Out* out,
-    IdxT size_rest,
-    const __grid_constant__ Shape shape,
-    const __grid_constant__ Strides strides,
-    int ndim) {
-  auto block = cg::this_thread_block();
-  auto grid = cg::this_grid();
-  IdxT index_rest =
-      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
-  if (index_rest >= size_rest) {
-    return;
-  }
-
-  auto shape_x = shape[ndim - 1];
-  auto stride_x = strides[ndim - 1];
-  IdxT index_x =
-      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
-  auto idx =
-      elem_to_loc(index_rest * shape_x, shape.data(), strides.data(), ndim);
-  auto in_vec =
-      load_vector<N_READS>(in + idx, index_x, shape_x, stride_x, In(0));
-  AlignedVector<Out, N_READS> out_vec;
-#pragma unroll
-  for (int i = 0; i < N_READS; ++i) {
-    out_vec[i] = Op{}(in_vec[i]);
-  }
-  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
-}
-
-template <typename Op, typename In, typename Out>
-constexpr bool supports_unary_op() {
-  if (std::is_same_v<Op, Abs> || std::is_same_v<Op, Negative> ||
-      std::is_same_v<Op, Sign> || std::is_same_v<Op, Square>) {
-    return std::is_same_v<In, Out>;
-  }
-  if (std::is_same_v<Op, ArcCosh> || std::is_same_v<Op, ArcSinh> ||
-      std::is_same_v<Op, ArcTanh> || std::is_same_v<Op, Erf> ||
-      std::is_same_v<Op, ErfInv> || std::is_same_v<Op, Expm1> ||
-      std::is_same_v<Op, Sigmoid>) {
-    return std::is_same_v<In, Out> && is_floating_v<In>;
-  }
-  if (std::is_same_v<Op, BitwiseInvert>) {
-    return std::is_same_v<In, Out> && std::is_integral_v<In> &&
-        !std::is_same_v<In, bool>;
-  }
-  if (std::is_same_v<Op, Ceil> || std::is_same_v<Op, Floor>) {
-    return std::is_same_v<In, Out> && !mlx::core::is_complex_v<In>;
-  }
-  if (std::is_same_v<Op, Conjugate>) {
-    return std::is_same_v<In, Out> && mlx::core::is_complex_v<In>;
-  }
-  if (std::is_same_v<Op, ArcCos> || std::is_same_v<Op, ArcSin> ||
-      std::is_same_v<Op, ArcTan> || std::is_same_v<Op, Cos> ||
-      std::is_same_v<Op, Cosh> || std::is_same_v<Op, Exp> ||
-      std::is_same_v<Op, Log> || std::is_same_v<Op, Log2> ||
-      std::is_same_v<Op, Log10> || std::is_same_v<Op, Log1p> ||
-      std::is_same_v<Op, Round> || std::is_same_v<Op, Rsqrt> ||
-      std::is_same_v<Op, Sqrt> || std::is_same_v<Op, Sin> ||
-      std::is_same_v<Op, Sinh> || std::is_same_v<Op, Tan> ||
-      std::is_same_v<Op, Tanh>) {
-    return std::is_same_v<In, Out> && is_inexact_v<In>;
-  }
-  if (std::is_same_v<Op, Imag> || std::is_same_v<Op, Real>) {
-    return mlx::core::is_complex_v<In> && std::is_same_v<Out, float>;
-  }
-  if (std::is_same_v<Op, LogicalNot>) {
-    return std::is_same_v<In, Out> && std::is_same_v<In, bool>;
-  }
-  return false;
-}
-
-} // namespace cu
-
-template <typename Op>
-void unary_op_gpu_inplace(
-    const std::vector<array>& inputs,
-    array& out,
-    const char* op,
-    const Stream& s) {
-  auto& in = inputs[0];
-  if (in.size() == 0) {
-    return;
-  }
-  bool contig = in.flags().contiguous;
-  bool large;
-  if (!contig) {
-    large = in.data_size() > INT32_MAX || out.size() > INT32_MAX;
-  } else {
-    large = in.data_size() > UINT32_MAX;
-  }
-
-  auto& encoder = cu::get_command_encoder(s);
-  encoder.set_input_array(in);
-  encoder.set_output_array(out);
-  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
-    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
-      using CTYPE_IN = MLX_GET_TYPE(in_type_tag);
-      using CTYPE_OUT = MLX_GET_TYPE(out_type_tag);
-      if constexpr (cu::supports_unary_op<Op, CTYPE_IN, CTYPE_OUT>()) {
-        dispatch_bool(large, [&](auto large) {
-          using InType = cuda_type_t<CTYPE_IN>;
-          using OutType = cuda_type_t<CTYPE_OUT>;
-          if (contig) {
-            using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-            constexpr int N_READS = 16 / sizeof(OutType);
-            auto [num_blocks, block_dims] = get_launch_args(
-                out.data_size(), out.shape(), out.strides(), large, N_READS);
-            encoder.add_kernel_node(
-                cu::unary_v<Op, InType, OutType, IdxT, N_READS>,
-                num_blocks,
-                block_dims,
-                0,
-                in.data<InType>(),
-                out.data<OutType>(),
-                out.data_size());
-          } else {
-            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
-            auto [shape, strides] = collapse_contiguous_dims(in);
-            auto ndim = shape.size();
-            int work_per_thread = 1;
-            auto kernel = cu::unary_g<Op, InType, OutType, IdxT, 1>;
-            auto dim0 = ndim > 0 ? shape.back() : 1;
-            auto rest = out.size() / dim0;
-            if (dim0 >= 4) {
-              kernel = cu::unary_g<Op, InType, OutType, IdxT, 4>;
-              work_per_thread = 4;
-            }
-            dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
-            auto block_dims = get_block_dims(dim0, rest, 1);
-            uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
-            uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
-            encoder.add_kernel_node(
-                kernel,
-                {num_blocks_x, num_blocks_y},
-                block_dims,
-                0,
-                in.data<InType>(),
-                out.data<OutType>(),
-                rest,
-                const_param(shape),
-                const_param(strides),
-                ndim);
-          }
-        });
-      } else {
-        throw std::runtime_error(fmt::format(
-            "Can not do unary op {} on input of {} with output of {}.",
-            op,
-            dtype_to_string(in.dtype()),
-            dtype_to_string(out.dtype())));
-      }
-    });
-  });
-}
-
-template <typename Op>
-void unary_op_gpu(
-    const std::vector<array>& inputs,
-    array& out,
-    const char* op,
-    const Stream& s) {
-  set_unary_output_data(inputs[0], out);
-  unary_op_gpu_inplace<Op>(inputs, out, op, s);
-}
-
-#define UNARY_GPU(func)                                               \
-  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
-    nvtx3::scoped_range r(#func "::eval_gpu");                        \
-    auto& s = out.primitive().stream();                               \
-    unary_op_gpu<cu::func>(inputs, out, name(), s);                   \
-  }
-
-UNARY_GPU(Abs)
-UNARY_GPU(ArcCos)
-UNARY_GPU(ArcCosh)
-UNARY_GPU(ArcSin)
-UNARY_GPU(ArcSinh)
-UNARY_GPU(ArcTan)
-UNARY_GPU(ArcTanh)
-UNARY_GPU(BitwiseInvert)
-UNARY_GPU(Ceil)
-UNARY_GPU(Conjugate)
-UNARY_GPU(Cos)
-UNARY_GPU(Cosh)
-UNARY_GPU(Erf)
-UNARY_GPU(ErfInv)
-UNARY_GPU(Exp)
-UNARY_GPU(Expm1)
-UNARY_GPU(Floor)
-UNARY_GPU(Imag)
-UNARY_GPU(Log1p)
-UNARY_GPU(LogicalNot)
-UNARY_GPU(Negative)
-UNARY_GPU(Real)
-UNARY_GPU(Sigmoid)
-UNARY_GPU(Sign)
-UNARY_GPU(Sin)
-UNARY_GPU(Sinh)
-UNARY_GPU(Square)
-UNARY_GPU(Tan)
-UNARY_GPU(Tanh)
-
-void Log::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Log::eval_gpu");
-  auto& s = out.primitive().stream();
-  switch (base_) {
-    case Base::e:
-      unary_op_gpu<cu::Log>(inputs, out, name(), s);
-      break;
-    case Base::two:
-      unary_op_gpu<cu::Log2>(inputs, out, name(), s);
-      break;
-    case Base::ten:
-      unary_op_gpu<cu::Log10>(inputs, out, name(), s);
-      break;
-  }
-}
-
-void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Round::eval_gpu");
-  assert(inputs.size() == 1);
-  const auto& in = inputs[0];
-  auto& s = out.primitive().stream();
-  if (issubdtype(in.dtype(), inexact)) {
-    unary_op_gpu<cu::Round>(inputs, out, name(), s);
-  } else {
-    // No-op integer types
-    out.copy_shared_buffer(in);
-  }
-}
-
-void Sqrt::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Sort::eval_gpu");
-  auto& s = out.primitive().stream();
-  if (recip_) {
-    unary_op_gpu<cu::Rsqrt>(inputs, out, "Rsqrt", s);
-  } else {
-    unary_op_gpu<cu::Sqrt>(inputs, out, "Sqrt", s);
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/utils.h

@@ -1,6 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
-// This file include utilies that are used by C++ code (i.e. .cpp files).
+// This file include utilities that are used by C++ code (i.e. .cpp files).
 
 #pragma once
 
@@ -12,6 +12,7 @@ namespace mlx::core {
 
 namespace cu {
 class Device;
+
 }
 
 struct Dtype;
@@ -86,4 +87,17 @@ class CudaStream : public CudaHandle<cudaStream_t, cudaStreamDestroy> {
   explicit CudaStream(cu::Device& device);
 };
 
+template <typename T>
+inline uint max_occupancy_block_dim(T kernel) {
+  int _, block_dim;
+  if constexpr (std::is_same_v<T, CUfunction>) {
+    CHECK_CUDA_ERROR(
+        cuOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel, 0, 0, 0));
+  } else {
+    CHECK_CUDA_ERROR(
+        cudaOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel));
+  }
+  return block_dim;
+}
+
 } // namespace mlx::core

mlx/backend/cuda/worker.cpp

@@ -5,9 +5,9 @@
 
 namespace mlx::core::cu {
 
-Worker::Worker()
-    : signal_stream_(device(mlx::core::Device::gpu)),
-      signal_event_(cudaEventDisableTiming | cudaEventBlockingSync),
+Worker::Worker(Device& d)
+    : signal_stream_(d),
+      signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync),
       worker_(&Worker::thread_fn, this) {}
 
 Worker::~Worker() {

mlx/backend/cuda/worker.h

@@ -15,7 +15,7 @@ namespace mlx::core::cu {
 // Run tasks in worker thread, synchronized with cuda stream.
 class Worker {
  public:
-  Worker();
+  explicit Worker(Device& d);
   ~Worker();
 
   Worker(const Worker&) = delete;

mlx/backend/gpu/primitives.cpp

@@ -51,7 +51,7 @@ void Contiguous::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Contiguous::eval_gpu");
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  constexpr size_t extra_bytes = 16384;
+  constexpr int64_t extra_bytes = 16384;
   if (in.buffer_size() <= out.nbytes() + extra_bytes &&
       (in.flags().row_contiguous ||
        (allow_col_major_ && in.flags().col_contiguous))) {

mlx/backend/gpu/slicing.cpp

@@ -11,7 +11,7 @@ void slice_gpu(
     array& out,
     const Shape& start_indices,
     const Shape& strides,
-    const Stream& s) {
+    const Stream& /* s */) {
   slice(in, out, start_indices, strides);
 }
 
@@ -27,7 +27,7 @@ void pad_gpu(
 
   // Find offset for start of input values
   size_t data_offset = 0;
-  for (int i = 0; i < axes.size(); i++) {
+  for (int i = 0; i < std::ssize(axes); i++) {
     auto ax = axes[i] < 0 ? out.ndim() + axes[i] : axes[i];
     data_offset += out.strides()[ax] * low_pad_size[i];
   }

mlx/backend/metal/allocator.cpp

@@ -32,7 +32,6 @@ namespace metal {
 
 MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
-      residency_set_(device_),
       buffer_cache_(
           vm_page_size,
           [](MTL::Buffer* buf) { return buf->length(); },
@@ -41,7 +40,8 @@ MetalAllocator::MetalAllocator()
               residency_set_.erase(buf);
             }
             buf->release();
-          }) {
+          }),
+      residency_set_(device_) {
   auto pool = metal::new_scoped_memory_pool();
   auto memsize = std::get<size_t>(device_info().at("memory_size"));
   auto max_rec_size =

mlx/backend/metal/allocator.h

@@ -65,7 +65,6 @@ class MetalAllocator : public allocator::Allocator {
   size_t peak_memory_{0};
   size_t max_pool_size_;
   size_t wired_limit_{0};
-  bool relaxed_{true};
   size_t num_resources_{0};
   size_t resource_limit_{0};
 

mlx/backend/metal/compiled.cpp

@@ -109,7 +109,7 @@ inline void build_kernel(
 
   // Read constant / contiguous inputs in tmps
   std::vector<array> nc_inputs;
-  for (int i = 0; i < inputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(inputs); ++i) {
     auto& x = inputs[i];
     auto& xname = namer.get_name(x);
 
@@ -134,7 +134,7 @@ inline void build_kernel(
   }
 
   // Initialize the indices for non-contiguous inputs
-  for (int i = 0; i < nc_inputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(nc_inputs); ++i) {
     auto& xname = namer.get_name(nc_inputs[i]);
     os += fmt::format("  {0} index_{1} = ", idx_type, xname);
     if (ndim == 1) {
@@ -174,7 +174,7 @@ inline void build_kernel(
       os += fmt::format("  for (int d = {0}; d >= 0; --d) {{\n", ndim - 3);
     }
     os += "    uint l = zpos % output_shape[d];\n";
-    for (int i = 0; i < nc_inputs.size(); ++i) {
+    for (int i = 0; i < std::ssize(nc_inputs); ++i) {
       auto& xname = namer.get_name(nc_inputs[i]);
       os += fmt::format("    index_{0} += ", xname);
       if (dynamic_dims) {
@@ -195,7 +195,7 @@ inline void build_kernel(
   }
 
   // Read non-contiguous inputs into tmps
-  for (int i = 0; i < nc_inputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(nc_inputs); ++i) {
     auto& x = nc_inputs[i];
     auto& xname = namer.get_name(x);
     os += fmt::format(
@@ -214,7 +214,7 @@ inline void build_kernel(
     } else {
       os += x.primitive().name();
       os += "()(";
-      for (int i = 0; i < x.inputs().size() - 1; i++) {
+      for (int i = 0; i < std::ssize(x.inputs()) - 1; i++) {
         os += fmt::format("tmp_{0}, ", namer.get_name(x.inputs()[i]));
       }
       os += fmt::format("tmp_{0});\n", namer.get_name(x.inputs().back()));
@@ -227,7 +227,7 @@ inline void build_kernel(
   }
   // Increment indices and close per thread loop
   if (work_per_thread > 1) {
-    for (int i = 0; i < nc_inputs.size(); ++i) {
+    for (int i = 0; i < std::ssize(nc_inputs); ++i) {
       auto& x = nc_inputs[i];
       auto& xname = namer.get_name(x);
       if (!dynamic_dims) {
@@ -396,7 +396,7 @@ void Compiled::eval_gpu(
   int cnt = 0;
   int stride_idx = 1; // idx 0 is the output strides
   Strides in_strides;
-  for (int i = 0; i < inputs.size(); i++) {
+  for (int i = 0; i < std::ssize(inputs); i++) {
     if (is_constant_(i)) {
       continue;
     }

mlx/backend/metal/conv.cpp

@@ -990,7 +990,7 @@ void conv_3D_gpu(
     const std::vector<int>& wt_dilation,
     const std::vector<int>& in_dilation,
     bool flip,
-    std::vector<array>& copies) {
+    std::vector<array>& /* copies */) {
   // Make conv params
   MLXConvParams<3> conv_params{
       /* const int  N = */ static_cast<int>(in.shape(0)),

mlx/backend/metal/custom_kernel.cpp

@@ -68,7 +68,7 @@ std::string write_signature(
   int index = 0;
   constexpr int max_constant_array_size = 8;
   // Add inputs
-  for (int i = 0; i < inputs.size(); ++i) {
+  for (int i = 0; i < std::ssize(inputs); ++i) {
     const auto& name = input_names[i];
     const auto& arr = inputs[i];
     auto dtype = get_type_string(arr.dtype());
@@ -109,7 +109,7 @@ std::string write_signature(
     }
   }
   // Add outputs
-  for (int i = 0; i < output_names.size(); ++i) {
+  for (int i = 0; i < std::ssize(output_names); ++i) {
     const auto& name = output_names[i];
     const auto& dtype = output_dtypes[i];
     kernel_source += "  device ";
@@ -126,8 +126,8 @@ std::string write_signature(
     kernel_source += " [[buffer(";
     kernel_source += std::to_string(index);
     kernel_source += ")]]";
-    if (index < inputs.size() + output_names.size() - 1 ||
-        attributes.size() > 0) {
+    if (index < std::ssize(inputs) + std::ssize(output_names) - 1 ||
+        std::ssize(attributes) > 0) {
       kernel_source += ",\n";
     } else {
       kernel_source += ") {\n";
@@ -138,7 +138,7 @@ std::string write_signature(
   index = 0;
   for (const auto& attr : attributes) {
     kernel_source += attr;
-    if (index < attributes.size() - 1) {
+    if (index < std::ssize(attributes) - 1) {
       kernel_source += ",\n";
     } else {
       kernel_source += ") {\n";
@@ -327,6 +327,10 @@ CustomKernelFunction metal_kernel(
 void CustomKernel::eval_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
+  // silence some warnings
+  (void)is_precompiled_;
+  (void)shared_memory_;
+
   auto& s = stream();
 
   std::vector<array> copies;
@@ -377,7 +381,7 @@ void CustomKernel::eval_gpu(
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
   int index = 0;
-  for (int i = 0; i < checked_inputs.size(); i++) {
+  for (int i = 0; i < std::ssize(checked_inputs); i++) {
     const array& in = checked_inputs[i];
     auto& shape_info = shape_infos_[i];
     compute_encoder.set_input_array(in, index);
@@ -404,7 +408,7 @@ void CustomKernel::eval_gpu(
   }
 
   const auto [tx, ty, tz] = threadgroup_;
-  auto tg_size = tx * ty * tz;
+  unsigned long tg_size = tx * ty * tz;
   auto max_tg_size = kernel->maxTotalThreadsPerThreadgroup();
   if (tg_size > max_tg_size) {
     std::ostringstream msg;

mlx/backend/metal/device.cpp

@@ -72,6 +72,19 @@ MTL::Library* try_load_bundle(
   }
   return nullptr;
 }
+
+MTL::Library* try_load_framework(
+    MTL::Device* device,
+    NS::URL* url,
+    const std::string& lib_name) {
+  std::string resource_path = std::string(url->fileSystemRepresentation()) +
+      "/" + lib_name + ".metallib";
+  auto [lib, error] = load_library_from_path(device, resource_path.c_str());
+  if (lib) {
+    return lib;
+  }
+  return nullptr;
+}
 #endif
 
 // Firstly, search for the metallib in the same path as this binary
@@ -103,6 +116,20 @@ std::pair<MTL::Library*, NS::Error*> load_swiftpm_library(
       return {library, nullptr};
     }
   }
+  // if SWIFTPM_BUNDLE is a framework identifier, try loading from that
+  auto frameworks = NS::Bundle::allFrameworks();
+  for (int i = 0, c = (int)frameworks->count(); i < c; i++) {
+    auto bundle = reinterpret_cast<NS::Bundle*>(frameworks->object(i));
+    if (!strcmp(bundle->bundleIdentifier()->utf8String(), SWIFTPM_BUNDLE)) {
+      library = try_load_framework(device, bundle->resourceURL(), lib_name);
+      if (library != nullptr) {
+        return {library, nullptr};
+      }
+    }
+  }
+#else
+  (void)device;
+  (void)lib_name;
 #endif
   return {nullptr, nullptr};
 }
@@ -471,6 +498,10 @@ void Device::end_encoding(int index) {
 CommandEncoder& Device::get_command_encoder(int index) {
   auto& stream = get_stream_(index);
   if (stream.encoder == nullptr) {
+    // Ensure there is an active command buffer
+    if (stream.buffer == nullptr) {
+      get_command_buffer(index);
+    }
     stream.encoder = std::make_unique<CommandEncoder>(stream);
     stream.fence = std::make_shared<Fence>(device_->newFence());
   }
@@ -685,7 +716,7 @@ MTL::LinkedFunctions* Device::get_linked_functions_(
   auto lfuncs = MTL::LinkedFunctions::linkedFunctions();
 
   std::vector<NS::Object*> objs(funcs.size());
-  for (int i = 0; i < funcs.size(); i++) {
+  for (int i = 0; i < std::ssize(funcs); i++) {
     objs[i] = funcs[i];
   }
 
@@ -724,7 +755,7 @@ MTL::ComputePipelineState* Device::get_kernel_(
   mtl_linked_funcs->release();
 
   // Add kernel to cache
-  auto inserted = kernel_map_.insert({hash_name, kernel});
+  kernel_map_.insert({hash_name, kernel});
 
   return kernel;
 }

mlx/backend/metal/device.h

@@ -137,7 +137,7 @@ struct DeviceStream {
   // Data updated between command buffers
   MTL::CommandBuffer* buffer{nullptr};
   int buffer_ops{0};
-  size_t buffer_sizes{0};
+  int64_t buffer_sizes{0};
 
   // The command encoder, fence, and temporaries are updated between command
   // encoders

mlx/backend/metal/eval.cpp

@@ -71,7 +71,7 @@ void eval(array& arr) {
     d.get_command_buffer(s.index);
   } else {
     command_buffer->addCompletedHandler(
-        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+        [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
           check_error(cbuf);
         });
   }
@@ -82,7 +82,7 @@ void finalize(Stream s) {
   auto& d = metal::device(s.device);
   auto cb = d.get_command_buffer(s.index);
   d.end_encoding(s.index);
-  cb->addCompletedHandler([s](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
+  cb->addCompletedHandler([](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
   d.commit_command_buffer(s.index);
   d.get_command_buffer(s.index);
 }

mlx/backend/metal/fence.cpp

@@ -76,7 +76,7 @@ void Fence::wait(Stream stream, const array& x) {
     auto command_buffer = d.get_command_buffer(idx);
     command_buffer->encodeWait(static_cast<MTL::Event*>(f.fence), f.count);
     command_buffer->addCompletedHandler(
-        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+        [fence_ = fence_](MTL::CommandBuffer* /* cbuf */) {});
     return;
   }
 
@@ -96,7 +96,7 @@ void Fence::wait(Stream stream, const array& x) {
   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(Stream stream, const array& x) {
@@ -124,7 +124,7 @@ void Fence::update(Stream stream, const array& x) {
     command_buffer->encodeSignalEvent(
         static_cast<MTL::Event*>(f.fence), f.count);
     command_buffer->addCompletedHandler(
-        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+        [fence_ = fence_](MTL::CommandBuffer* /* cbuf */) {});
     return;
   }
 
@@ -154,7 +154,7 @@ void Fence::update(Stream stream, const array& x) {
   compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
 
   d.get_command_buffer(idx)->addCompletedHandler(
-      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+      [fence_ = fence_](MTL::CommandBuffer* /* cbuf */) {});
 }
 
 } // namespace mlx::core

mlx/backend/metal/fft.cpp

@@ -60,7 +60,7 @@ struct FourStepParams {
 void fft_op(
     const array& in,
     array& out,
-    size_t axis,
+    int64_t axis,
     bool inverse,
     bool real,
     const FourStepParams four_step_params,
@@ -93,7 +93,7 @@ std::vector<int> plan_stockham_fft(int n) {
   if (n == 1) {
     return plan;
   }
-  for (int i = 0; i < radices.size(); i++) {
+  for (int i = 0; i < std::ssize(radices); i++) {
     int radix = radices[i];
     // Manually tuned radices for powers of 2
     if (is_power_of_2(orig_n) && orig_n < 512 && radix > 4) {
@@ -150,7 +150,6 @@ FFTPlan plan_fft(int n) {
       }
       // See if we can use Rader's algorithm to Stockham decompose n - 1
       auto rader_factors = prime_factors(factor - 1);
-      int last_factor = -1;
       for (int rf : rader_factors) {
         // We don't nest Rader's algorithm so if `factor - 1`
         // isn't Stockham decomposable we give up and do Bluestein's.
@@ -182,7 +181,7 @@ int compute_elems_per_thread(FFTPlan plan) {
   steps.insert(steps.end(), plan.stockham.begin(), plan.stockham.end());
   steps.insert(steps.end(), plan.rader.begin(), plan.rader.end());
   std::set<int> used_radices;
-  for (int i = 0; i < steps.size(); i++) {
+  for (int i = 0; i < std::ssize(steps); i++) {
     int radix = radices[i % radices.size()];
     if (steps[i] > 0) {
       used_radices.insert(radix);
@@ -261,7 +260,7 @@ int primitive_root(int n) {
 
 std::tuple<array, array, array> compute_raders_constants(
     int rader_n,
-    const Stream& s) {
+    const Stream& /* s */) {
   int proot = primitive_root(rader_n);
   // Fermat's little theorem
   int inv = mod_exp(proot, rader_n - 2, rader_n);
@@ -313,8 +312,6 @@ std::pair<array, array> compute_bluestein_constants(int n, int bluestein_n) {
   // w_k = np.exp(-1j * np.pi / N * (np.arange(-N + 1, N) ** 2))
   // w_q = np.fft.fft(1/w_k)
   // return w_k, w_q
-  int length = 2 * n - 1;
-
   std::vector<std::complex<float>> w_k_vec(n);
   std::vector<std::complex<float>> w_q_vec(bluestein_n, 0);
 
@@ -484,8 +481,6 @@ void four_step_fft(
     std::vector<array>& copies,
     const Stream& s,
     bool in_place) {
-  auto& d = metal::device(s.device);
-
   if (plan.bluestein_n == -1) {
     // Fast no transpose implementation for powers of 2.
     FourStepParams four_step_params = {
@@ -513,7 +508,7 @@ void four_step_fft(
 void fft_op(
     const array& in,
     array& out,
-    size_t axis,
+    int64_t axis,
     bool inverse,
     bool real,
     const FourStepParams four_step_params,
@@ -617,11 +612,11 @@ void fft_op(
 
   // Start of radix/rader step constants
   int index = 4;
-  for (int i = 0; i < plan.stockham.size(); i++) {
+  for (int i = 0; i < std::ssize(plan.stockham); i++) {
     func_consts.push_back(make_int(&plan.stockham[i], index));
     index += 1;
   }
-  for (int i = 0; i < plan.rader.size(); i++) {
+  for (int i = 0; i < std::ssize(plan.rader); i++) {
     func_consts.push_back(make_int(&plan.rader[i], index));
     index += 1;
   }
@@ -776,8 +771,8 @@ void nd_fft_op(
   array temp1(temp_shape, complex64, nullptr, {});
   array temp2(temp_shape, complex64, nullptr, {});
   std::vector<array> temp_arrs = {temp1, temp2};
-  for (int i = axes.size() - 1; i >= 0; i--) {
-    int reverse_index = axes.size() - i - 1;
+  for (int i = std::ssize(axes) - 1; i >= 0; i--) {
+    int reverse_index = std::ssize(axes) - i - 1;
     // For 5D and above, we don't want to reallocate our two temporary arrays
     bool inplace = reverse_index >= 3 && i != 0;
     // Opposite order for fft vs ifft
@@ -785,9 +780,8 @@ void nd_fft_op(
     size_t axis = axes[index];
     // Mirror np.fft.(i)rfftn and perform a real transform
     // only on the final axis.
-    bool step_real = (real && index == axes.size() - 1);
-    auto step_shape = inverse ? out.shape(axis) : in.shape(axis);
-    const array& in_arr = i == axes.size() - 1 ? in : temp_arrs[1 - i % 2];
+    bool step_real = (real && index == std::ssize(axes) - 1);
+    const array& in_arr = i == std::ssize(axes) - 1 ? in : temp_arrs[1 - i % 2];
     array& out_arr = i == 0 ? out : temp_arrs[i % 2];
     fft_op(in_arr, out_arr, axis, inverse, step_real, inplace, s);
   }

mlx/backend/metal/hadamard.cpp

@@ -43,7 +43,7 @@ std::string gen_hadamard_codelet(int m) {
   while (end != std::string_view::npos) {
     source << "  tmp[" << index << "] = ";
     auto row = matrix.substr(start, end - start);
-    for (int i = 0; i < row.length(); i++) {
+    for (int i = 0; i < std::ssize(row); i++) {
       source << " " << row[i] << " x[" << i << "]";
     }
     source << ";" << std::endl;

mlx/backend/metal/indexing.cpp

@@ -52,7 +52,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  size_t slice_size = 1;
+  int64_t slice_size = 1;
   for (auto s : slice_sizes_) {
     slice_size *= s;
   }
@@ -94,8 +94,8 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
     auto kernel = d.get_kernel(kernel_name, lib);
     compute_encoder.set_compute_pipeline_state(kernel);
 
-    size_t dim_x = (slice_size + work_per_thread - 1) / work_per_thread;
-    size_t dim_y = indices.size();
+    int64_t dim_x = (slice_size + work_per_thread - 1) / work_per_thread;
+    int64_t dim_y = indices.size();
     auto group_dims = get_block_dims(dim_x, dim_y, 1);
     MTL::Size grid_dims = MTL::Size(dim_x, dim_y, 1);
 
@@ -110,7 +110,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
 
   int idx_ndim = nidx ? inputs[1].ndim() : 0;
-  size_t ndim = src.ndim();
+  int64_t ndim = src.ndim();
 
   std::string kernel_name = fmt::format(
       "gather{0}{1}_{2}_{3}_{4}",
@@ -149,8 +149,8 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   // Launch 3D grid of threads
   // First two dimensions for the indices, the last one for the slice
-  size_t dim0 = 1;
-  size_t dim1 = 1;
+  int64_t dim0 = 1;
+  int64_t dim1 = 1;
   if (nidx) {
     if (inputs[1].ndim() >= 1) {
       dim0 = inputs[1].shape(0);
@@ -159,13 +159,13 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
       dim1 = inputs[1].size() / dim0;
     }
   }
-  size_t dim2 = slice_size;
+  int64_t dim2 = slice_size;
   auto group_dims = get_block_dims(dim0, dim1, dim2);
   MTL::Size grid_dims = MTL::Size(dim0, dim1, dim2);
 
   // Collect all idx shapes and strides into one place
   std::vector<int> idx_shapes;
-  std::vector<size_t> idx_strides;
+  std::vector<int64_t> idx_strides;
   std::vector<char> idx_contigs;
   for (int i = 0; i < nidx; ++i) {
     idx_shapes.insert(
@@ -246,7 +246,7 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& d = metal::device(s.device);
 
   int idx_ndim = nidx ? inputs[1].ndim() : 0;
-  size_t idx_size = nidx ? inputs[1].size() : 1;
+  int64_t idx_size = nidx ? inputs[1].size() : 1;
 
   auto idx_to_out = idx_size / out.size();
   int nwork;
@@ -345,7 +345,7 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& compute_encoder = d.get_command_encoder(s.index);
   auto kernel = d.get_kernel(kernel_name, lib);
 
-  size_t nthreads = upd.size();
+  int64_t nthreads = upd.size();
 
   compute_encoder.set_compute_pipeline_state(kernel);
 
@@ -354,8 +354,8 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   compute_encoder.set_output_array(out, 2);
 
   // Set update info
-  size_t upd_ndim = upd.ndim();
-  size_t upd_size = 1;
+  int64_t upd_ndim = upd.ndim();
+  int64_t upd_size = 1;
   for (int i = idx_ndim; i < upd.ndim(); ++i) {
     upd_size *= upd.shape(i);
   }
@@ -378,7 +378,7 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
 
   if (upd_ndim == 0) {
-    // Need placeholders so Metal doesn't compalain
+    // Need placeholders so Metal doesn't complain
     int shape_ = 0;
     int64_t stride_ = 0;
     compute_encoder.set_bytes(shape_, 3);
@@ -391,9 +391,9 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   compute_encoder.set_bytes(upd_size, 6);
 
   // Set output info
-  size_t out_ndim = out.ndim();
+  int64_t out_ndim = out.ndim();
   if (out_ndim == 0) {
-    // Need placeholders so Metal doesn't compalain
+    // Need placeholders so Metal doesn't complain
     int shape_ = 0;
     int64_t stride_ = 0;
     compute_encoder.set_bytes(shape_, 7);
@@ -448,7 +448,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  size_t ndim = src.ndim();
+  int64_t ndim = src.ndim();
 
   bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
 
@@ -486,8 +486,8 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Grid [size post, index size, size pre]
-  size_t size_pre = 1;
-  size_t size_post = 1;
+  int64_t size_pre = 1;
+  int64_t size_post = 1;
   for (int i = 0; i < axis_; ++i) {
     size_pre *= idx.shape(i);
   }
@@ -541,7 +541,7 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  size_t ndim = src.ndim();
+  int64_t ndim = src.ndim();
 
   bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
 
@@ -602,8 +602,8 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Grid [size post, index size, size pre]
-  size_t size_pre = 1;
-  size_t size_post = 1;
+  int64_t size_pre = 1;
+  int64_t size_post = 1;
   for (int i = 0; i < axis_; ++i) {
     size_pre *= idx.shape(i);
   }

mlx/backend/metal/jit_kernels.cpp

@@ -144,8 +144,7 @@ MTL::ComputePipelineState* get_ternary_kernel(
     auto t_str = get_type_string(type);
     std::string kernel_source = metal::utils();
     concatenate(kernel_source, metal::ternary_ops(), metal::ternary());
-    const std::array<std::pair<std::string, std::string>, 4> kernel_types = {{
-        {"v2", "ternary_v2"},
+    const std::array<std::pair<std::string, std::string>, 3> kernel_types = {{
         {"g1large", "ternary_g_nd1"},
         {"g2large", "ternary_g_nd2"},
         {"g3large", "ternary_g_nd3"},
@@ -154,13 +153,29 @@ MTL::ComputePipelineState* get_ternary_kernel(
       kernel_source +=
           get_template_definition(name + "_" + lib_name, func, t_str, op);
     }
+
+    kernel_source += get_template_definition(
+        "v2_" + lib_name, "ternary_v2", t_str, op, false, false);
+    kernel_source += get_template_definition(
+        "sv2_" + lib_name, "ternary_v2", t_str, op, true, false);
+    kernel_source += get_template_definition(
+        "vs2_" + lib_name, "ternary_v2", t_str, op, false, true);
+
     if (get_work_per_thread(type) > 1) {
-      kernel_source +=
-          get_template_definition("vn_" + lib_name, "ternary_v", t_str, op);
+      kernel_source += get_template_definition(
+          "vn_" + lib_name, "ternary_v", t_str, op, false, false);
+      kernel_source += get_template_definition(
+          "svn_" + lib_name, "ternary_v", t_str, op, true, false);
+      kernel_source += get_template_definition(
+          "vsn_" + lib_name, "ternary_v", t_str, op, false, true);
     }
 
-    kernel_source +=
-        get_template_definition("v_" + lib_name, "ternary_v", t_str, op, 1);
+    kernel_source += get_template_definition(
+        "v_" + lib_name, "ternary_v", t_str, op, false, false, 1);
+    kernel_source += get_template_definition(
+        "sv_" + lib_name, "ternary_v", t_str, op, true, false, 1);
+    kernel_source += get_template_definition(
+        "vs_" + lib_name, "ternary_v", t_str, op, false, true, 1);
     kernel_source += get_template_definition(
         "g1_" + lib_name, "ternary_g_nd1", t_str, op, "int");
     kernel_source += get_template_definition(

mlx/backend/metal/kernels/complex.h

@@ -104,6 +104,27 @@ constexpr bool operator==(complex64_t a, complex64_t b) {
 constexpr complex64_t operator+(complex64_t a, complex64_t b) {
   return {a.real + b.real, a.imag + b.imag};
 }
+
+constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
+constexpr threadgroup complex64_t& operator+=(
+    threadgroup complex64_t& a,
+    complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
+constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
 constexpr complex64_t operator+(float a, complex64_t b) {
   return {a + b.real, b.imag};
 }

mlx/backend/metal/kernels/gemv.metal

@@ -15,6 +15,15 @@ using namespace metal;
 
 #define MLX_MTL_CONST static constant constexpr const
 
+template <typename U>
+struct DefaultAccT {
+  using type = float;
+};
+template <>
+struct DefaultAccT<complex64_t> {
+  using type = complex64_t;
+};
+
 template <
     typename T,
     const int BM, /* Threadgroup rows (in simdgroups) */
@@ -24,8 +33,10 @@ template <
     const int TM, /* Thread rows (in elements) */
     const int TN, /* Thread cols (in elements) */
     const bool kDoAxpby, /* Do out = alpha * out + beta * bias */
-    typename AccT = float>
+    typename AccT = typename DefaultAccT<T>::type>
 struct GEMVKernel {
+  using acc_type = AccT;
+
   MLX_MTL_CONST int threadsM = BM * SM;
   MLX_MTL_CONST int threadsN = BN * SN;
 
@@ -246,8 +257,10 @@ template <
     const int TM, /* Thread rows (in elements) */
     const int TN, /* Thread cols (in elements) */
     const bool kDoAxpby, /* Do out = alpha * out + beta * bias */
-    typename AccT = float>
+    typename AccT = typename DefaultAccT<T>::type>
 struct GEMVTKernel {
+  using acc_type = AccT;
+
   MLX_MTL_CONST int threadsM = BM * SM;
   MLX_MTL_CONST int threadsN = BN * SN;
 
@@ -453,7 +466,7 @@ template <
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
-  threadgroup float tgp_memory
+  threadgroup typename gemv_kernel::acc_type tgp_memory
       [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   // Update batch offsets
@@ -530,6 +543,7 @@ template <
 instantiate_gemv_blocks(float32, float);
 instantiate_gemv_blocks(float16, half);
 instantiate_gemv_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_blocks(complex64, complex64_t);
 
 template <
     typename T,
@@ -565,7 +579,7 @@ template <
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, false>;
-  threadgroup float tgp_memory
+  threadgroup typename gemv_kernel::acc_type tgp_memory
       [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   uint32_t indx_vec;
@@ -636,6 +650,7 @@ template <
 instantiate_gemv_bs_blocks(float32, float);
 instantiate_gemv_bs_blocks(float16, half);
 instantiate_gemv_bs_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_bs_blocks(complex64, complex64_t);
 
 ///////////////////////////////////////////////////////////////////////////////
 /// Vector matrix multiplication
@@ -672,7 +687,7 @@ template <
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
-  threadgroup float tgp_memory
+  threadgroup typename gemv_kernel::acc_type tgp_memory
       [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   // Update batch offsets
@@ -738,7 +753,8 @@ template <
 // clang-format off
 instantiate_gemv_t_blocks(float32, float);
 instantiate_gemv_t_blocks(float16, half);
-instantiate_gemv_t_blocks(bfloat16, bfloat16_t); // clang-format on
+instantiate_gemv_t_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_t_blocks(complex64, complex64_t); // clang-format on
 
 template <
     typename T,
@@ -773,8 +789,8 @@ template <
     uint3 lid [[thread_position_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
-  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, false, float>;
-  threadgroup float tgp_memory
+  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, false>;
+  threadgroup typename gemv_kernel::acc_type tgp_memory
       [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   uint32_t indx_vec;
@@ -848,4 +864,5 @@ template <
 // clang-format off
 instantiate_gemv_t_bs_blocks(float32, float);
 instantiate_gemv_t_bs_blocks(float16, half);
-instantiate_gemv_t_bs_blocks(bfloat16, bfloat16_t); // clang-format on
+instantiate_gemv_t_bs_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_t_bs_blocks(complex64, complex64_t); // clang-format on

mlx/backend/metal/kernels/sort.h

@@ -19,11 +19,28 @@ METAL_FUNC void thread_swap(thread T& a, thread T& b) {
   b = w;
 }
 
+template <typename T, typename = void>
+struct Init {
+  static constexpr constant T v = Limits<T>::max;
+};
+
+template <typename T>
+struct Init<T, metal::enable_if_t<metal::is_floating_point_v<T>>> {
+  static constexpr constant T v = metal::numeric_limits<T>::quiet_NaN();
+};
+
 template <typename T>
 struct LessThan {
-  static constexpr constant T init = Limits<T>::max;
-
-  METAL_FUNC bool operator()(T a, T b) {
+  static constexpr constant T init = Init<T>::v;
+  METAL_FUNC bool operator()(T a, T b) const {
+    if constexpr (
+        metal::is_floating_point_v<T> || metal::is_same_v<T, complex64_t>) {
+      bool an = isnan(a);
+      bool bn = isnan(b);
+      if (an | bn) {
+        return (!an) & bn;
+      }
+    }
     return a < b;
   }
 };

mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal

@@ -3,9 +3,8 @@
 #include <metal_stdlib>
 
 // clang-format off
-#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
-
 #include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
 #include "mlx/backend/metal/kernels/steel/conv/conv.h"
 #include "mlx/backend/metal/kernels/steel/conv/params.h"
 #include "mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h"

mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal

@@ -3,9 +3,8 @@
 #include <metal_stdlib>
 
 // clang-format off
-#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
-
 #include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
 #include "mlx/backend/metal/kernels/steel/conv/conv.h"
 #include "mlx/backend/metal/kernels/steel/conv/params.h"
 #include "mlx/backend/metal/kernels/steel/utils.h"

mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.metal

@@ -23,10 +23,12 @@
     instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 1, 2) \
     instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 32, 32, 2, 2) \
     instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 64, 16, 1, 2) \
-    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2)
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 32,  8, 4, 1)
 
 instantiate_gemm_shapes_helper(float16, half, float16, half);
 instantiate_gemm_shapes_helper(bfloat16, bfloat16_t, bfloat16, bfloat16_t);
 
 instantiate_gemm_shapes_helper(float32, float, float32, float);
+instantiate_gemm_shapes_helper(complex64, complex64_t, complex64, complex64_t);
 // clang-format on

mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.metal

@@ -1,8 +1,8 @@
 // Copyright © 2024 Apple Inc.
 
 // clang-format off
-#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
 #include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
 #include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h"
 
 #define instantiate_gemm(                                              \

mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal

@@ -60,6 +60,7 @@
 instantiate_gemm_shapes_helper(float16, half, float32, float);
 instantiate_gemm_shapes_helper(bfloat16, bfloat16_t, float32, float);
 instantiate_gemm_shapes_helper(float32, float, float32, float);
+instantiate_gemm_shapes_helper(complex64, complex64_t, complex64, complex64_t);
 
 #define instantiate_accum(oname, otype, aname, atype)      \
   instantiate_kernel(                                      \
@@ -71,4 +72,5 @@ instantiate_gemm_shapes_helper(float32, float, float32, float);
 
 instantiate_accum(bfloat16, bfloat16_t, float32, float);
 instantiate_accum(float16, half, float32, float);
-instantiate_accum(float32, float, float32, float); // clang-format on
+instantiate_accum(float32, float, float32, float);
+instantiate_accum(complex64, complex64_t, complex64, complex64_t); // clang-format on

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

@@ -421,6 +421,16 @@ METAL_FUNC void tile_matmad(
   }
 }
 
+template <typename InT>
+struct TransformNone<complex64_t, InT> {
+  static METAL_FUNC complex64_t apply(complex64_t x) {
+    return x;
+  }
+  static METAL_FUNC complex64_t apply(complex64_t x, complex64_t) {
+    return x;
+  }
+};
+
 template <
     typename T,
     typename U,
@@ -731,5 +741,406 @@ struct BlockMMA {
   }
 };
 
+template <
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    short lda_tgp,
+    short ldb_tgp,
+    typename AccumType,
+    typename Epilogue>
+struct BlockMMA<
+    complex64_t,
+    U,
+    BM,
+    BN,
+    BK,
+    WM,
+    WN,
+    transpose_a,
+    transpose_b,
+    lda_tgp,
+    ldb_tgp,
+    AccumType,
+    Epilogue> {
+  static_assert(
+      metal::is_same_v<AccumType, float>,
+      "BlockMMA<complex64_t,...> expects float accumulators");
+  static_assert(
+      metal::is_same_v<U, complex64_t>,
+      "For complex BlockMMA, U must be complex64_t; use a different epilogue for projections");
+  // MMAFrag size
+  STEEL_CONST short kFragSize = 8;
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TM_stride = kFragSize * WM;
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TN_stride = kFragSize * WN;
+
+  // Warp tile size along M
+  STEEL_CONST short TM = BM / (kFragSize * WM);
+  // Warp tile size along N
+  STEEL_CONST short TN = BN / (kFragSize * WN);
+
+  // Threadgroup A strides
+  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
+  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
+
+  // Threadgroup B strides
+  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
+  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
+
+  // Threadgroup strides along K
+  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
+  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
+
+  // When indexing complex as float[2]
+  STEEL_CONST short A_str_m_f = A_str_m * 2;
+  STEEL_CONST short A_str_k_f = A_str_k * 2;
+  STEEL_CONST short B_str_k_f = B_str_k * 2;
+  STEEL_CONST short B_str_n_f = B_str_n * 2;
+  STEEL_CONST short tile_stride_a_f = tile_stride_a * 2;
+  STEEL_CONST short tile_stride_b_f = tile_stride_b * 2;
+
+  // Accumulators (real/imag)
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile_r;
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile_i;
+
+  // Offsets within threadgroup
+  short sm, sn;
+  short As_offset, Bs_offset;
+
+  /* Constructor */
+  METAL_FUNC BlockMMA(
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    // Determine thread position in simdgroup matrix
+    short tm = kFragSize * (simd_group_id / WN);
+    short tn = kFragSize * (simd_group_id % WN);
+
+    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+    sm = simd_coord.y;
+    sn = simd_coord.x;
+
+    // Determine thread and simdgroup offset
+    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // (M,K)
+    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // (K,N)
+
+    sm += tm;
+    sn += tn;
+  }
+
+  /* Karatsuba MMA: 3 real MMAs per K-chunk */
+  METAL_FUNC void mma(
+      const threadgroup complex64_t* As,
+      const threadgroup complex64_t* Bs) {
+    // Adjust for simdgroup and thread location
+    As += As_offset;
+    Bs += Bs_offset;
+    threadgroup const float* As_f =
+        reinterpret_cast<threadgroup const float*>(As);
+    threadgroup const float* Bs_f =
+        reinterpret_cast<threadgroup const float*>(Bs);
+
+    // Iterate over BK in blocks of kFragSize
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < BK; kk += kFragSize) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      MMATile<AccumType, TM, 1, MMAFrag_acc_t> Ar, Ai;
+      Ar.template load<float, WM, 1, A_str_m_f, A_str_k_f>(As_f + 0);
+      Ai.template load<float, WM, 1, A_str_m_f, A_str_k_f>(As_f + 1);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      MMATile<AccumType, 1, TN, MMAFrag_acc_t> Br, Bi;
+      Br.template load<float, 1, WN, B_str_k_f, B_str_n_f>(Bs_f + 0);
+      Bi.template load<float, 1, WN, B_str_k_f, B_str_n_f>(Bs_f + 1);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      // P = Ar*Br ; Q = Ai*Bi ; R = (Ar+Ai)*(Br+Bi)
+      MMATile<AccumType, TM, TN, MMAFrag_acc_t> P, Q, R;
+
+      tile_matmad(P, Ar, Br, P);
+      tile_matmad(Q, Ai, Bi, Q);
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Ar)::kElemsPerTile; ++i)
+        Ar.elems()[i] += Ai.elems()[i];
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Br)::kElemsPerTile; ++i)
+        Br.elems()[i] += Bi.elems()[i];
+
+      tile_matmad(R, Ar, Br, R);
+
+      // C_r += P - Q ; C_i -= Q
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; ++i) {
+        const auto p = P.elems()[i];
+        const auto q = Q.elems()[i];
+        const auto r = R.elems()[i];
+        Ctile_r.elems()[i] += (p - q);
+        Ctile_i.elems()[i] += (r - p - q);
+      }
+
+      // Progress to next simdgroup tile
+      As_f += tile_stride_a_f;
+      Bs_f += tile_stride_b_f;
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(device U* D, const int ldd) {
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        thread const auto& r = Ctile_r.frag_at(i, j);
+        thread const auto& im = Ctile_i.frag_at(i, j);
+        int off = (i * TM_stride) * ldd + (j * TN_stride);
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+          D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void
+  store_result_slice(device U* D, const int ldd, short2 start, short2 stop) {
+    D += sm * ldd + sn;
+    start -= short2(sn, sm);
+    stop -= short2(sn, sm);
+
+    if (stop.y <= 0 || stop.x <= 0)
+      return;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; ++i) {
+      const int row = i * TM_stride;
+      if (row >= start.y && row < stop.y) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < TN; ++j) {
+          const int off = row * ldd + (j * TN_stride);
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; ++k) {
+            const int col = j * TN_stride + k;
+            if (col >= start.x && col < stop.x) {
+              D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+            }
+          }
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void
+  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
+    D += sm * ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < TN; j++) {
+          int off = (i * TM_stride) * ldd + (j * TN_stride);
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename UnaryEpilogue>
+  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) {
+      complex64_t out = epilogue_op.apply(
+          complex64_t(Ctile_r.elems()[i], Ctile_i.elems()[i]));
+      Ctile_r.elems()[i] = out.real;
+      Ctile_i.elems()[i] = out.imag;
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread auto& r = Ctile_r.frag_at(i, j);
+        thread auto& im = Ctile_i.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+          complex64_t out = epilogue_op.apply(
+              complex64_t(r[k], im[k]), C[offset_c + k * fdc]);
+          r[k] = out.real;
+          im[k] = out.imag;
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue_safe(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread auto& r = Ctile_r.frag_at(i, j);
+        thread auto& im = Ctile_i.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+        complex64_t tmp[kelems];
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          if ((j * TN_stride + k) < dst_tile_dims.x &&
+              (i * TM_stride) < dst_tile_dims.y) {
+            tmp[k] = C[offset_c + k * fdc];
+          } else {
+            tmp[k] = complex64_t(0.0f, 0.0f);
+          }
+        }
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          complex64_t out = epilogue_op.apply(complex64_t(r[k], im[k]), tmp[k]);
+          r[k] = out.real;
+          im[k] = out.imag;
+        }
+      }
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+
+    constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread const auto& r = Ctile_r.frag_at(i, j);
+        thread const auto& im = Ctile_i.frag_at(i, j);
+        int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+        int off_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          D[off_d + k] =
+              epilogue_op.apply(complex64_t(r[k], im[k]), C[off_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void store_result_safe(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < TN; j++) {
+          // Get accumulated result and associated offset in Cr, Ci
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+          int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+          int off_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+          // Apply epilogue
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[off_d + k] = epilogue_op.apply(
+                  complex64_t(r[k], im[k]), C[off_c + k * fdc]);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
 } // namespace steel
 } // namespace mlx