only interrupt during an eval

.circleci/config.yml

@@ -24,8 +24,8 @@ jobs:
         type: boolean
         default: false
     macos:
-      xcode: "16.2.0"
-    resource_class: m2pro.medium
+      xcode: "15.2.0"
+    resource_class: macos.m1.medium.gen1
     steps:
       - checkout
       - run:
@@ -89,14 +89,15 @@ jobs:
             pip install numpy
             sudo apt-get update
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
       - run:
           name: Install Python package
           command: |
-            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF
+              CMAKE_COMPILE_WARNING_AS_ERROR=ON" \
               CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
               python3 setup.py build_ext --inplace
-            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF \
+              CMAKE_COMPILE_WARNING_AS_ERROR=ON" \
               CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
               python3 setup.py develop
       - run:
@@ -109,8 +110,6 @@ jobs:
           name: Run Python tests
           command: |
             python3 -m unittest discover python/tests -v
-            mpirun --bind-to none -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py
       - run:
           name: Build CPP only
           command: |
@@ -125,15 +124,10 @@ jobs:
     parameters:
       xcode_version:
         type: string
-        default: "16.2.0"
-      macosx_deployment_target:
-        type: string
-        default: ""
+        default: "15.2.0"
     macos:
       xcode: << parameters.xcode_version >>
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    resource_class: m2pro.medium
+    resource_class: macos.m1.medium.gen1
     steps:
       - checkout
       - run:
@@ -155,7 +149,7 @@ jobs:
           command: |
             source env/bin/activate
             DEBUG=1 CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
-            CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
+            CMAKE_ARGS="CMAKE_COMPILE_WARNING_AS_ERROR=ON" \
               pip install -e . -v
       - run:
           name: Generate package stubs
@@ -219,18 +213,13 @@ jobs:
         default: "3.9"
       xcode_version:
         type: string
-        default: "16.2.0"
+        default: "15.2.0"
       build_env:
         type: string
         default: ""
-      macosx_deployment_target:
-        type: string
-        default: ""
     macos:
       xcode: << parameters.xcode_version >>
-    resource_class: m2pro.medium
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
+    resource_class: macos.m1.medium.gen1
     steps:
       - checkout
       - run:
@@ -251,7 +240,7 @@ jobs:
           name: Install Python package
           command: |
             source env/bin/activate
-            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
+            DEV_RELEASE=1 \
               CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
               pip install . -v
       - run:
@@ -346,7 +335,7 @@ workflows:
       - mac_build_and_test:
           matrix:
             parameters:
-              macosx_deployment_target: ["13.5", "14.0"]
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
       - linux_build_and_test
       - build_documentation 
 
@@ -366,70 +355,8 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
+              xcode_version: ["15.0.0", "15.2.0"]
               build_env: ["PYPI_RELEASE=1"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
       - build_documentation:
           filters:
             tags:
@@ -452,7 +379,7 @@ workflows:
           requires: [ hold ]
           matrix:
             parameters:
-              macosx_deployment_target: ["13.5", "14.0"]
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
       - linux_build_and_test:
           requires: [ hold ]
   nightly_build:
@@ -465,54 +392,7 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
+              xcode_version: ["15.0.0", "15.2.0"]
   weekly_build:
     when:
       and:
@@ -523,70 +403,8 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
               build_env: ["DEV_RELEASE=1"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
   linux_test_release:
     when:
       and:

.gitignore

@@ -36,7 +36,6 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
-uv.lock
 
 # vim
 *.swp

CMakeLists.txt

@@ -34,7 +34,6 @@ option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
 option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
 option(MLX_BUILD_METAL "Build metal backend" ON)
 option(MLX_BUILD_CPU "Build cpu backend" ON)
-option(MLX_BUILD_CUDA "Build cuda backend" OFF)
 option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
@@ -84,10 +83,6 @@ if(MLX_BUILD_METAL)
   set(QUARTZ_LIB "-framework QuartzCore")
 endif()
 
-if(MLX_BUILD_CUDA)
-  enable_language(CUDA)
-endif()
-
 if(MLX_BUILD_METAL AND NOT METAL_LIB)
   message(STATUS "Metal not found. Unable to build GPU")
   set(MLX_BUILD_METAL OFF)
@@ -217,6 +212,24 @@ else()
   set(MLX_BUILD_ACCELERATE OFF)
 endif()
 
+find_package(MPI)
+if(MPI_FOUND)
+  execute_process(
+    COMMAND zsh "-c" "mpirun --version"
+    OUTPUT_VARIABLE MPI_VERSION
+    ERROR_QUIET)
+  if(${MPI_VERSION} MATCHES ".*Open MPI.*")
+    target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
+  elseif(MPI_VERSION STREQUAL "")
+    set(MPI_FOUND FALSE)
+    message(
+      WARNING "MPI found but mpirun is not available. Building without MPI.")
+  else()
+    set(MPI_FOUND FALSE)
+    message(WARNING "MPI which is not OpenMPI found. Building without MPI.")
+  endif()
+endif()
+
 message(STATUS "Downloading json")
 FetchContent_Declare(
   json

CONTRIBUTING.md

@@ -5,26 +5,26 @@ possible.
 
 ## Pull Requests
 
-1. Fork and submit pull requests to the repo.
+1. Fork and submit pull requests to the repo. 
 2. If you've added code that should be tested, add tests.
 3. If a change is likely to impact efficiency, run some of the benchmarks before
    and after the change. Examples of benchmarks can be found in `benchmarks/python/`.
 4. If you've changed APIs, update the documentation.
-5. Every PR should have passing tests and at least one review.
+5. Every PR should have passing tests and at least one review. 
 6. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`.
    This should install hooks for running `black` and `clang-format` to ensure
    consistent style for C++ and python code.
-
+ 
    You can also run the formatters manually as follows:
-
-   ```shell
-   clang-format -i file.cpp
-   ```
-
-   ```shell
-   black file.py
-   ```
-
+ 
+     ```
+     clang-format -i file.cpp
+     ```
+ 
+     ```
+     black file.py
+     ```
+ 
    or run `pre-commit run --all-files` to check all files in the repo.
 
 ## Issues

MANIFEST.in

@@ -1,6 +1,4 @@
 include CMakeLists.txt
-include mlx.pc.in
 recursive-include mlx/ *
-include cmake/*
 include python/src/*
 include python/mlx/py.typed # support type hinting as in PEP-561

benchmarks/python/gather_mm_bench.py

@@ -1,74 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate([x[i] @ w[j].T for i, j in enumerate(idx.tolist())], axis=0)
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_mm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_mm(x, w1.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_mm(x, w2.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = x @ w1.T
-        x = x @ w2.T
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_mm()

benchmarks/python/gather_qmm_bench.py

@@ -1,84 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate(
-            [
-                mx.quantized_matmul(x[i], w[0][j], w[1][j], w[2][j], transpose=True)
-                for i, j in enumerate(idx.tolist())
-            ],
-            axis=0,
-        )
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_qmm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = mx.quantized_matmul(x, *w1, transpose=True)
-        x = mx.quantized_matmul(x, *w2, transpose=True)
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_qmm()

benchmarks/python/sdpa_bench.py

@@ -28,34 +28,11 @@ def bench(f, *args):
     return (e - s) * 1e-9
 
 
-def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype):
-    np_dtype = getattr(np, dtype)
-
-    shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D)
-    shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D)
-
-    scale = 1.0 / math.sqrt(D)
-
-    q_np = np.random.normal(0.0, 1.0, shape_q).astype(np_dtype)
-    k_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-    v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-
-    q_mx = mx.array(q_np)
-    k_mx = mx.array(k_np)
-    v_mx = mx.array(v_np)
-
-    if mask is not None:
-        if mask == "additive":
-            mask_np = np.random.normal(0.0, 1.0, (B, qH, qL, kL)).astype(np_dtype)
-            mask = mx.array(mask_np)
-        elif mask == "bool":
-            mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
-            mask = mx.array(mask_np)
-
-    return q_mx, k_mx, v_mx, scale, mask
+def mlx_sdpa_fused_inner(q, k, v, scale):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None)
 
 
-def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
+def mlx_sdpa_unfused_inner(q, k, v, scale, f32softmax=False):
     q_dtype = q.dtype
     q = q * mx.array(scale, q_dtype)
     n_q_heads = q.shape[-3]
@@ -64,7 +41,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
 
     B = q.shape[0]
     L = q.shape[2]
-    kL = k.shape[2]
 
     if n_repeats > 1:
         q = mx.reshape(q, [B, n_kv_heads, n_repeats, L, -1])
@@ -72,27 +48,10 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
         v = mx.expand_dims(v, 2)
 
     scores = q @ mx.swapaxes(k, -1, -2)
-
-    if mask is not None:
-
-        if mask == "causal":
-            q_offset = max(0, kL - L)
-            q_indices = mx.arange(q_offset, q_offset + L)
-            k_indices = mx.arange(kL)
-            mask = q_indices[:, None] >= k_indices[None]
-
-        if n_repeats > 1 and mask.ndim >= 3:
-            if mask.shape[-3] == 1:
-                mask = mx.expand_dims(mask, -3)
-            else:
-                mask = mx.unflatten(mask, -3, (n_kv_heads, n_repeats))
-
-        if mask.dtype == mx.bool_:
-            scores = mx.where(mask, scores, -np.float32(np.inf))
-        else:
-            scores += mask
-
-    scores = mx.softmax(scores, axis=-1, precise=True)
+    if f32softmax:
+        scores = mx.softmax(scores.astype(mx.float32), axis=-1).astype(q_dtype)
+    else:
+        scores = mx.softmax(scores, axis=-1)
 
     out = scores @ v
     if n_repeats > 1:
@@ -101,55 +60,74 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
     return out
 
 
-def mlx_fused_attn(q, k, v, scale, mask):
-    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention(f, q, k, v, scale, mask=None, transpose=False):
-    if transpose:
-        q_t = mx.transpose(q, (0, 2, 1, 3))
-        k_t = mx.transpose(k, (0, 2, 1, 3))
-        v_t = mx.transpose(v, (0, 2, 1, 3))
-        o_t = f(q_t, k_t, v_t, scale=scale, mask=mask)
-        return mx.transpose(o_t, (0, 2, 1, 3))
-    else:
-        return f(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention_bench(f, q, k, v, scale, mask=None, transpose=False):
+def mlx_spda_unfused(q, k, v, scale, transpose):
     q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
 
     for i in range(N_iter_func):
-        q_out = do_attention(f, q_out, k, v, scale, mask=mask, transpose=transpose)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_unfused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
 
     mx.eval(q_out)
     return q_out
 
 
-def bench_shape(
-    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None
-):
-    q_mx, k_mx, v_mx, scale, mask = prepare_inputs(
-        B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, mask_in, transpose, dtype
+def mlx_spda_fused(q, k, v, scale, transpose):
+    q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
+
+    for i in range(N_iter_func):
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_fused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+
+    mx.eval(q_out)
+    return q_out
+
+
+def bench_shape(B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose=True):
+    shape_q = (
+        (B, qsl, n_q_heads, head_dim) if transpose else (B, n_q_heads, qsl, head_dim)
+    )
+    shape_kv = (
+        (B, ksl, n_kv_heads, head_dim) if transpose else (B, n_kv_heads, ksl, head_dim)
     )
 
-    time_mlx_unfused = bench(
-        do_attention_bench, mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
-    time_mlx_fused = bench(
-        do_attention_bench, mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    q_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_q).astype(np_dtype)
+    k_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
+    v_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
 
-    o_mlx_fused = do_attention(mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose)
-    o_mlx_unfused = do_attention(
-        mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    scale = math.sqrt(1.0 / head_dim)
 
-    atol = 1e-5 if dtype == "float32" else 2e-4
+    q_mx = mx.array(q_np)
+    k_mx = mx.array(k_np)
+    v_mx = mx.array(v_np)
 
-    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol, rtol=atol):
+    time_mlx_unfused = bench(mlx_spda_unfused, q_mx, k_mx, v_mx, scale, transpose)
+    time_mlx_fused = bench(mlx_spda_fused, q_mx, k_mx, v_mx, scale, transpose)
+
+    if transpose:
+        q_mx = mx.transpose(q_mx, (0, 2, 1, 3))
+        k_mx = mx.transpose(k_mx, (0, 2, 1, 3))
+        v_mx = mx.transpose(v_mx, (0, 2, 1, 3))
+
+    o_mlx_fused = mlx_sdpa_fused_inner(q_mx, k_mx, v_mx, scale)
+    o_mlx_unfused = mlx_sdpa_unfused_inner(q_mx, k_mx, v_mx, scale, f32softmax=True)
+
+    atol = 1e-5 if np_dtype == np.float32 else 1e-4
+
+    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol):
         print(
-            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
+            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
         )
 
     return time_mlx_fused, time_mlx_unfused
@@ -173,51 +151,39 @@ if __name__ == "__main__":
           (  1,   128,   128,       64,   32,    32),
           (  1,   256,   256,       64,   32,    32),
           (  1,   512,   512,       64,   32,    32),
-          (  1,  1024,  1024,       64,   32,     8),
-          (  1,  2048,  2048,       64,   32,     8),
-          (  1,  4096,  4096,       64,   32,     8),
+          (  1,  1024,  1024,       64,   32,    32),
+          (  1,  2048,  2048,       64,   32,    32),
+          (  1,  4096,  4096,       64,   32,    32),
     )
 
     shapes_80 = (
         # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
-          (  1,  1024,  1024,       80,   32,     8),
-          (  1,  2048,  2048,       80,   32,     8),
-          (  1,  4096,  4096,       80,   32,     8),
+          (  1,  1024,  1024,       80,   32,    32),
+          (  1,  2048,  2048,       80,   32,    32),
+          (  1,  4096,  4096,       80,   32,    32),
     )
 
     shapes_128 = (
         # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
-          (  1,  1024,  1024,      128,   32,     8),
-          (  1,  2048,  2048,      128,   32,     8),
-          (  1,  4096,  4096,      128,   32,     8),
+          (  1,  1024,  1024,      128,   32,    32),
+          (  1,  2048,  2048,      128,   32,    32),
+          (  1,  4096,  4096,      128,   32,    32),
     )
     # fmt: on
 
     shapes = shapes_64 + shapes_80 + shapes_128
 
-    masks = [None, "bool", "causal"]
-
-    print(
-        "  B,   qsl,   ksl, hdim, n_qh, n_kvh, t,   dtype,     mask, t_unfs, t_fuse, diff%"
-    )
+    print("  B,   qsl,   ksl, hdim, n_qh, n_kvh, tpose,   dtype, t_unfs, t_fuse, diff%")
 
     for dtype in dtypes:
         for transpose in transposes:
             for B, qsl, ksl, head_dim, n_q_heads, n_kv_heads in shapes:
-                for mask_in in masks:
-                    time_mlx_fused, time_mlx_unfused = bench_shape(
-                        B,
-                        qsl,
-                        ksl,
-                        head_dim,
-                        n_q_heads,
-                        n_kv_heads,
-                        dtype,
-                        transpose,
-                        mask_in,
-                    )
-                    diff = time_mlx_unfused / time_mlx_fused - 1.0
-                    t_str = 1 if transpose else 0
-                    print(
-                        f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:1d}, {dtype}, {str(mask_in):>8}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
-                    )
+                np_dtype = getattr(np, dtype)
+                time_mlx_fused, time_mlx_unfused = bench_shape(
+                    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose
+                )
+                diff = time_mlx_unfused / time_mlx_fused - 1.0
+                t_str = 1 if transpose else 0
+                print(
+                    f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:5d}, {dtype}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
+                )

docs/Doxyfile

@@ -13,7 +13,7 @@ EXCLUDE_PATTERNS       = */private/*
 CREATE_SUBDIRS         = NO
 FULL_PATH_NAMES        = YES
 RECURSIVE              = YES
-GENERATE_HTML          = NO
+GENERATE_HTML          = YES
 GENERATE_LATEX         = NO
 GENERATE_XML           = YES
 XML_PROGRAMLISTING     = YES

docs/src/dev/extensions.rst

@@ -93,9 +93,9 @@ Primitives
 ^^^^^^^^^^^
 
 A :class:`Primitive` is part of the computation graph of an :class:`array`. It
-defines how to create output arrays given input arrays. Further, a
+defines how to create outputs arrays given a input arrays. Further, a
 :class:`Primitive` has methods to run on the CPU or GPU and for function
-transformations such as ``vjp`` and ``jvp``.  Let's go back to our example to be
+transformations such as ``vjp`` and ``jvp``.  Lets go back to our example to be
 more concrete:
 
 .. code-block:: C++
@@ -128,7 +128,7 @@ more concrete:
         /** The vector-Jacobian product. */
         std::vector<array> vjp(
             const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
+            const array& cotan,
             const std::vector<int>& argnums,
             const std::vector<array>& outputs) override;
 
@@ -247,7 +247,9 @@ point-wise. This is captured in the templated function :meth:`axpby_impl`.
       float alpha_,
       float beta_,
       mx::Stream stream) {
-    out.set_data(mx::allocator::malloc(out.nbytes()));
+    // Allocate the output with `malloc_or_wait` which synchronously allocates
+    // memory, potentially waiting if the system is under memory pressure
+    out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
 
     // Get the CPU command encoder and register input and output arrays
     auto& encoder = mx::cpu::get_command_encoder(stream);
@@ -391,7 +393,7 @@ below.
         auto& d = metal::device(s.device);
 
         // Allocate output memory
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
         // Resolve name of kernel
         std::ostringstream kname;
@@ -469,7 +471,7 @@ one we just defined:
             const std::vector<array>& tangents,
             const std::vector<int>& argnums) {
         // Forward mode diff that pushes along the tangents
-        // The jvp transform on the primitive can be built with ops
+        // The jvp transform on the primitive can built with ops
         // that are scheduled on the same stream as the primitive
 
         // If argnums = {0}, we only push along x in which case the
@@ -481,7 +483,7 @@ one we just defined:
             auto scale_arr = array(scale, tangents[0].dtype());
             return {multiply(scale_arr, tangents[0], stream())};
         }
-        // If argnums = {0, 1}, we take contributions from both
+        // If, argnums = {0, 1}, we take contributions from both
         // which gives us jvp = tangent_x * alpha + tangent_y * beta
         else {
             return {axpby(tangents[0], tangents[1], alpha_, beta_, stream())};
@@ -735,7 +737,7 @@ Let's look at a simple script and its results:
 
     print(f"c shape: {c.shape}")
     print(f"c dtype: {c.dtype}")
-    print(f"c is correct: {mx.all(c == 6.0).item()}")
+    print(f"c correct: {mx.all(c == 6.0).item()}")
 
 Output:
 
@@ -743,7 +745,7 @@ Output:
 
     c shape: [3, 4]
     c dtype: float32
-    c is correct: True
+    c correctness: True
 
 Results
 ^^^^^^^

docs/src/index.rst

@@ -70,7 +70,6 @@ are the CPU and GPU.
    python/fft
    python/linalg
    python/metal
-   python/memory_management
    python/nn
    python/optimizers
    python/distributed

docs/src/python/array.rst

@@ -38,7 +38,6 @@ Array
     array.log10
     array.log1p
     array.log2
-    array.logcumsumexp
     array.logsumexp
     array.max
     array.mean

docs/src/python/fft.rst

@@ -20,5 +20,3 @@ FFT
   irfft2
   rfftn
   irfftn
-  fftshift
-  ifftshift

docs/src/python/linalg.rst

@@ -20,6 +20,5 @@ Linear Algebra
     eigh
     lu
     lu_factor
-    pinv
     solve
     solve_triangular

docs/src/python/memory_management.rst

@@ -1,16 +0,0 @@
-Memory Management
-=================
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-  get_active_memory
-  get_peak_memory
-  reset_peak_memory
-  get_cache_memory
-  set_memory_limit
-  set_cache_limit
-  set_wired_limit
-  clear_cache

docs/src/python/metal.rst

@@ -8,5 +8,13 @@ Metal
 
   is_available
   device_info
+  get_active_memory
+  get_peak_memory
+  reset_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit
+  set_wired_limit
+  clear_cache
   start_capture
   stop_capture

docs/src/python/ops.rst

@@ -36,12 +36,10 @@ Operations
    bitwise_or
    bitwise_xor
    block_masked_mm
-   broadcast_arrays
    broadcast_to
    ceil
    clip
    concatenate
-   contiguous
    conj
    conjugate
    convolve
@@ -103,7 +101,6 @@ Operations
    log10
    log1p
    logaddexp
-   logcumsumexp
    logical_not
    logical_and
    logical_or

docs/src/python/optimizers/common_optimizers.rst

@@ -18,4 +18,3 @@ Common Optimizers
    AdamW
    Adamax
    Lion
-   MultiOptimizer

docs/src/python/transforms.rst

@@ -9,7 +9,6 @@ Transforms
   :toctree: _autosummary
 
    eval
-   async_eval
    compile
    custom_function
    disable_compile

examples/extensions/axpby/axpby.cpp

@@ -72,7 +72,9 @@ void axpby_impl(
     float alpha_,
     float beta_,
     mx::Stream stream) {
-  out.set_data(mx::allocator::malloc(out.nbytes()));
+  // Allocate the output with `malloc_or_wait` which synchronously allocates
+  // memory, potentially waiting if the system is under memory pressure
+  out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
 
   // Get the CPU command encoder and register input and output arrays
   auto& encoder = mx::cpu::get_command_encoder(stream);
@@ -158,12 +160,12 @@ void Axpby::eval_gpu(
   // Allocate output memory with strides based on specialization
   if (contiguous_kernel) {
     out.set_data(
-        mx::allocator::malloc(x.data_size() * out.itemsize()),
+        mx::allocator::malloc_or_wait(x.data_size() * out.itemsize()),
         x.data_size(),
         x.strides(),
         x.flags());
   } else {
-    out.set_data(mx::allocator::malloc(out.nbytes()));
+    out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
   }
 
   // Resolve name of kernel (corresponds to axpby.metal)

mlx/CMakeLists.txt

@@ -5,7 +5,6 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/dtype_utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/export.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
@@ -49,16 +48,5 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
 if(MLX_BUILD_METAL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
 else()
-  target_sources(mlx
-                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/no_metal.cpp)
-endif()
-
-if(MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda)
-endif()
-
-if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
-else()
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_gpu)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_metal)
 endif()

mlx/allocator.cpp

@@ -4,11 +4,12 @@
 #include <sstream>
 
 #include "mlx/allocator.h"
+#include "mlx/scheduler.h"
 
 namespace mlx::core::allocator {
 
 Buffer malloc(size_t size) {
-  auto buffer = allocator().malloc(size);
+  auto buffer = allocator().malloc(size, /* allow_swap */ true);
   if (size && !buffer.ptr()) {
     std::ostringstream msg;
     msg << "[malloc] Unable to allocate " << size << " bytes.";
@@ -21,4 +22,45 @@ void free(Buffer buffer) {
   allocator().free(buffer);
 }
 
+Buffer CommonAllocator::malloc(size_t size, bool) {
+  void* ptr = std::malloc(size + sizeof(size_t));
+  if (ptr != nullptr) {
+    *static_cast<size_t*>(ptr) = size;
+  }
+  return Buffer{ptr};
+}
+
+void CommonAllocator::free(Buffer buffer) {
+  std::free(buffer.ptr());
+}
+
+size_t CommonAllocator::size(Buffer buffer) const {
+  if (buffer.ptr() == nullptr) {
+    return 0;
+  }
+  return *static_cast<size_t*>(buffer.ptr());
+}
+
+Buffer malloc_or_wait(size_t size) {
+  auto buffer = allocator().malloc(size);
+
+  while (size && !buffer.ptr() && scheduler::n_active_tasks() > 0) {
+    scheduler::wait_for_one();
+    buffer = allocator().malloc(size);
+  }
+
+  // Try swapping if needed
+  if (size && !buffer.ptr()) {
+    buffer = allocator().malloc(size, /* allow_swap = */ true);
+  }
+
+  if (size && !buffer.ptr()) {
+    std::ostringstream msg;
+    msg << "[malloc_or_wait] Unable to allocate " << size << " bytes.";
+    throw std::runtime_error(msg.str());
+  }
+
+  return buffer;
+}
+
 } // namespace mlx::core::allocator

mlx/allocator.h

@@ -32,10 +32,14 @@ Buffer malloc(size_t size);
 
 void free(Buffer buffer);
 
+// Wait for running tasks to finish and free up memory
+// if allocation fails
+Buffer malloc_or_wait(size_t size);
+
 class Allocator {
   /** Abstract base class for a memory allocator. */
  public:
-  virtual Buffer malloc(size_t size) = 0;
+  virtual Buffer malloc(size_t size, bool allow_swap = false) = 0;
   virtual void free(Buffer buffer) = 0;
   virtual size_t size(Buffer buffer) const = 0;
 
@@ -49,4 +53,16 @@ class Allocator {
 
 Allocator& allocator();
 
+class CommonAllocator : public Allocator {
+  /** A general CPU allocator. */
+ public:
+  virtual Buffer malloc(size_t size, bool allow_swap = false) override;
+  virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
+
+ private:
+  CommonAllocator() = default;
+  friend Allocator& allocator();
+};
+
 } // namespace mlx::core::allocator

mlx/array.h

@@ -339,11 +339,11 @@ class array {
     return allocator::allocator().size(buffer());
   }
 
-  // Return the shared pointer to the array::Data struct
-  const std::shared_ptr<Data>& data_shared_ptr() const {
+  // Return a copy of the shared pointer
+  // to the array::Data struct
+  std::shared_ptr<Data> data_shared_ptr() const {
     return array_desc_->data;
   }
-
   // Return a raw pointer to the arrays data
   template <typename T>
   T* data() {
@@ -356,7 +356,7 @@ class array {
   }
 
   enum Status {
-    // The output of a computation which has not been scheduled.
+    // The ouptut of a computation which has not been scheduled.
     // For example, the status of `x` in `auto x = a + b`.
     unscheduled,
 

mlx/backend/common/CMakeLists.txt

@@ -1,10 +1,8 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/broadcasting.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/transpose.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)

mlx/backend/common/binary.h

@@ -44,14 +44,14 @@ inline void set_binary_op_output_data(
   switch (bopt) {
     case BinaryOpType::ScalarScalar:
       out.set_data(
-          allocator::malloc(out.itemsize()), 1, a.strides(), a.flags());
+          allocator::malloc_or_wait(out.itemsize()), 1, a.strides(), a.flags());
       break;
     case BinaryOpType::ScalarVector:
       if (b_donatable) {
         out.copy_shared_buffer(b);
       } else {
         out.set_data(
-            allocator::malloc(b.data_size() * out.itemsize()),
+            allocator::malloc_or_wait(b.data_size() * out.itemsize()),
             b.data_size(),
             b.strides(),
             b.flags());
@@ -62,7 +62,7 @@ inline void set_binary_op_output_data(
         out.copy_shared_buffer(a);
       } else {
         out.set_data(
-            allocator::malloc(a.data_size() * out.itemsize()),
+            allocator::malloc_or_wait(a.data_size() * out.itemsize()),
             a.data_size(),
             a.strides(),
             a.flags());
@@ -75,7 +75,7 @@ inline void set_binary_op_output_data(
         out.copy_shared_buffer(b);
       } else {
         out.set_data(
-            allocator::malloc(a.data_size() * out.itemsize()),
+            allocator::malloc_or_wait(a.data_size() * out.itemsize()),
             a.data_size(),
             a.strides(),
             a.flags());
@@ -88,7 +88,7 @@ inline void set_binary_op_output_data(
           b_donatable && b.flags().row_contiguous && b.size() == out.size()) {
         out.copy_shared_buffer(b);
       } else {
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
       }
       break;
   }

mlx/backend/common/broadcasting.cpp

@@ -1,24 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include "mlx/backend/common/utils.h"
-
-namespace mlx::core {
-
-void broadcast(const array& in, array& out) {
-  if (out.size() == 0) {
-    out.set_data(nullptr);
-    return;
-  }
-  Strides strides(out.ndim(), 0);
-  int diff = out.ndim() - in.ndim();
-  for (int i = in.ndim() - 1; i >= 0; --i) {
-    strides[i + diff] = (in.shape()[i] == 1) ? 0 : in.strides()[i];
-  }
-  auto flags = in.flags();
-  if (out.size() > in.size()) {
-    flags.row_contiguous = flags.col_contiguous = false;
-  }
-  out.copy_shared_buffer(in, strides, flags, in.data_size());
-}
-
-} // namespace mlx::core

mlx/backend/common/broadcasting.h

@@ -1,11 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#include "mlx/array.h"
-
-namespace mlx::core {
-
-void broadcast(const array& in, array& out);
-
-} // namespace mlx::core

mlx/backend/common/common.cpp

@@ -1,8 +1,6 @@
 // Copyright © 2024 Apple Inc.
 #include <cassert>
 
-#include "mlx/backend/common/broadcasting.h"
-#include "mlx/backend/common/transpose.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/primitives.h"
 
@@ -20,23 +18,47 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
         "AsStrided must be used with row contiguous arrays only.");
   }
 
-  // Calculate the contiguity based on the given shape and strides
-  auto [ds, rc, cc] = check_contiguity(shape_, strides_);
+  // 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++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
+  }
   auto flags = in.flags();
-
   // TODO: Compute the contiguous flag in a better way cause now we are
   //       unnecessarily strict.
-  flags.contiguous = rc || cc;
-  flags.row_contiguous = rc;
-  flags.col_contiguous = cc;
+  flags.contiguous = row_contiguous || col_contiguous;
+  flags.row_contiguous = row_contiguous;
+  flags.col_contiguous = col_contiguous;
 
-  // There is no easy way to compute the actual data size so we use out.size()
-  // when the array is not contiguous.
-  size_t data_size = flags.contiguous ? ds : out.size();
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
 
   return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 
+void broadcast(const array& in, array& out) {
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+  Strides strides(out.ndim(), 0);
+  int diff = out.ndim() - in.ndim();
+  for (int i = in.ndim() - 1; i >= 0; --i) {
+    strides[i + diff] = (in.shape()[i] == 1) ? 0 : in.strides()[i];
+  }
+  auto flags = in.flags();
+  if (out.size() > in.size()) {
+    flags.row_contiguous = flags.col_contiguous = false;
+  }
+  out.copy_shared_buffer(in, strides, flags, in.data_size());
+}
+
 void Broadcast::eval(const std::vector<array>& inputs, array& out) {
   broadcast(inputs[0], out);
 }
@@ -81,7 +103,7 @@ void ExpandDims::eval(const std::vector<array>& inputs, array& out) {
 
 void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   double numel = 1;
   for (auto ax : axes_) {
@@ -264,7 +286,36 @@ void StopGradient::eval(const std::vector<array>& inputs, array& out) {
 
 void Transpose::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  transpose(inputs[0], out, axes_);
+  Strides out_strides(out.ndim());
+  auto& in = inputs[0];
+  for (int ax = 0; ax < axes_.size(); ++ax) {
+    out_strides[ax] = in.strides()[axes_[ax]];
+  }
+
+  // Conditions for {row/col}_contiguous
+  // - array must be contiguous (no gaps)
+  // - underlying buffer size should have the same size as the array
+  // - cumulative product of shapes is equal to the strides (we can ignore axes
+  //   with size == 1)
+  //   - in the forward direction (column contiguous)
+  //   - in the reverse direction (row contiguous)
+  // - vectors are both row and col contiguous (hence if both row/col are
+  //   true, they stay true)
+  auto flags = in.flags();
+  if (flags.contiguous && in.data_size() == in.size()) {
+    int64_t f_stride = 1;
+    int64_t b_stride = 1;
+    flags.col_contiguous = true;
+    flags.row_contiguous = true;
+    for (int i = 0, ri = out.ndim() - 1; i < out.ndim(); ++i, --ri) {
+      flags.col_contiguous &= (out_strides[i] == f_stride || out.shape(i) == 1);
+      f_stride *= out.shape(i);
+      flags.row_contiguous &=
+          (out_strides[ri] == b_stride || out.shape(ri) == 1);
+      b_stride *= out.shape(ri);
+    }
+  }
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
 }
 
 } // namespace mlx::core

mlx/backend/common/compiled.cpp

@@ -188,7 +188,7 @@ void compiled_allocate_outputs(
     }
     for (; o < outputs.size(); ++o) {
       outputs[o].set_data(
-          allocator::malloc(data_size * outputs[o].itemsize()),
+          allocator::malloc_or_wait(data_size * outputs[o].itemsize()),
           data_size,
           strides,
           flags);
@@ -211,7 +211,7 @@ void compiled_allocate_outputs(
       }
     }
     for (; o < outputs.size(); ++o) {
-      outputs[o].set_data(allocator::malloc(outputs[o].nbytes()));
+      outputs[o].set_data(allocator::malloc_or_wait(outputs[o].nbytes()));
     }
   }
 }

mlx/backend/common/copy.h

@@ -31,14 +31,14 @@ inline bool set_copy_output_data(const array& in, array& out, CopyType ctype) {
       return true;
     } else {
       out.set_data(
-          allocator::malloc(in.data_size() * out.itemsize()),
+          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
           in.data_size(),
           in.strides(),
           in.flags());
       return false;
     }
   } else {
-    out.set_data(allocator::malloc(out.nbytes()));
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
     return false;
   }
 }

mlx/backend/common/hadamard.h

@@ -99,11 +99,7 @@ inline std::pair<int, int> decompose_hadamard(int n) {
           "[hadamard] Only supports n = m*2^k where m in (1, 12, 20, 28).");
     }
   }
-  if (n > (1 << 26)) {
-    throw std::invalid_argument(
-        "[hadamard] Only supports n = m*2^k where k <= 26");
-  }
   return {n, m};
 }
 
-} // namespace mlx::core
+} // namespace mlx::core

mlx/backend/common/load.cpp

@@ -28,7 +28,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) {
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   auto read_task = [out_ptr = out.data<char>(),
                     size = out.size(),
                     itemsize = out.itemsize(),

mlx/backend/common/ternary.h

@@ -48,12 +48,12 @@ inline void set_ternary_op_output_data(
   switch (topt) {
     case TernaryOpType::ScalarScalarScalar:
       out.set_data(
-          allocator::malloc(out.itemsize()), 1, b.strides(), b.flags());
+          allocator::malloc_or_wait(out.itemsize()), 1, b.strides(), b.flags());
       break;
     case TernaryOpType::VectorVectorVector:
       if (!(maybe_donate(a) || maybe_donate(b) || maybe_donate(c))) {
         out.set_data(
-            allocator::malloc(out.itemsize() * b.data_size()),
+            allocator::malloc_or_wait(out.itemsize() * b.data_size()),
             b.data_size(),
             b.strides(),
             b.flags());
@@ -64,7 +64,7 @@ inline void set_ternary_op_output_data(
       if (!((a.flags().row_contiguous && maybe_donate(a)) ||
             (b.flags().row_contiguous && maybe_donate(b)) ||
             (c.flags().row_contiguous && maybe_donate(c)))) {
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
       }
       break;
   }

mlx/backend/common/transpose.cpp

@@ -1,57 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include <cassert>
-
-#include "mlx/backend/common/utils.h"
-
-namespace mlx::core {
-
-void transpose(const array& in, array& out, const std::vector<int>& axes) {
-  Strides out_strides(out.ndim());
-  for (int ax = 0; ax < axes.size(); ++ax) {
-    out_strides[ax] = in.strides()[axes[ax]];
-  }
-
-  // Conditions for {row/col}_contiguous
-  // - array must be contiguous (no gaps)
-  // - underlying buffer size should have the same size as the array
-  // - cumulative product of shapes is equal to the strides (we can ignore axes
-  //   with size == 1)
-  //   - in the forward direction (column contiguous)
-  //   - in the reverse direction (row contiguous)
-  // - vectors are both row and col contiguous (hence if both row/col are
-  //   true, they stay true)
-  auto flags = in.flags();
-  if (flags.contiguous && in.data_size() == in.size()) {
-    auto [_, rc, cc] = check_contiguity(out.shape(), out_strides);
-    flags.row_contiguous = rc;
-    flags.col_contiguous = cc;
-  }
-  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
-}
-
-void as_transposed(array& out, const std::vector<int>& axes) {
-  assert(out.data_size() == out.size() && out.flags().contiguous);
-
-  // Calculate the contiguous strides.
-  Strides strides(out.ndim(), 1);
-  for (int i = out.ndim() - 2; i >= 0; i--) {
-    strides[i] = strides[i + 1] * out.shape(i);
-  }
-
-  // Calculate the new strides for transposing.
-  Strides new_strides;
-  new_strides.reserve(out.ndim());
-  for (auto ax : axes) {
-    new_strides.push_back(strides[ax]);
-  }
-
-  auto [ds, rc, cc] = check_contiguity(out.shape(), new_strides);
-  auto flags = out.flags();
-  flags.row_contiguous = rc;
-  flags.col_contiguous = cc;
-
-  out.copy_shared_buffer(out, new_strides, flags, ds);
-}
-
-} // namespace mlx::core

mlx/backend/common/transpose.h

@@ -1,12 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#include "mlx/array.h"
-
-namespace mlx::core {
-
-void transpose(const array& in, array& out, const std::vector<int>& axes);
-void as_transposed(array& out, const std::vector<int>& axes);
-
-} // namespace mlx::core

mlx/backend/common/utils.h

@@ -132,11 +132,6 @@ struct ContiguousIterator {
 };
 
 inline auto check_contiguity(const Shape& shape, const Strides& strides) {
-  // Conditions for {row/col}_contiguous
-  // - cumulative product of shapes is equal to the strides (we can ignore axes
-  //   with size == 1)
-  //   - in the forward direction (column contiguous)
-  //   - in the reverse direction (row contiguous)
   size_t no_broadcast_data_size = 1;
   int64_t f_stride = 1;
   int64_t b_stride = 1;

mlx/backend/cpu/CMakeLists.txt

@@ -40,8 +40,7 @@ add_dependencies(mlx cpu_compiled_preamble)
 
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
@@ -59,7 +58,6 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
@@ -75,8 +73,8 @@ target_sources(
 if(MLX_BUILD_ACCELERATE)
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/bnns.cpp)
 else()
-  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_fp16.cpp
-                             ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_bf16.cpp)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/no_fp16.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/gemms/no_bf16.cpp)
 endif()
 
 if(IOS)

mlx/backend/cpu/arg_reduce.cpp

@@ -68,7 +68,7 @@ void arg_reduce_dispatch(
 void ArgReduce::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.set_input_array(in);
   encoder.set_output_array(out);

mlx/backend/cpu/available.cpp

@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cpu/available.h"
-
-namespace mlx::core::cpu {
-
-bool is_available() {
-  return true;
-}
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/available.h

@@ -1,9 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-namespace mlx::core::cpu {
-
-bool is_available();
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/binary.cpp

@@ -172,12 +172,9 @@ void binary_float(
       case bfloat16:
         binary_op<bfloat16_t, Op>(a, b, out, bopt);
         break;
-      case complex64:
-        binary_op<complex64_t, Op>(a, b, out, bopt);
-        break;
       default:
         throw std::runtime_error(
-            "[binary_float] Only supports floating point types.");
+            "[binary_float] Only supports non-complex floating point types.");
     }
   });
 }

mlx/backend/cpu/compiled.cpp

@@ -40,10 +40,7 @@ struct CompilerCache {
   std::shared_mutex mtx;
 };
 
-static CompilerCache& cache() {
-  static CompilerCache cache_;
-  return cache_;
-};
+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.
@@ -59,16 +56,14 @@ void* compile(
     const std::string& kernel_name,
     const std::function<std::string(void)>& source_builder) {
   {
-    std::shared_lock lock(cache().mtx);
-    if (auto it = cache().kernels.find(kernel_name);
-        it != cache().kernels.end()) {
+    std::shared_lock lock(cache.mtx);
+    if (auto it = cache.kernels.find(kernel_name); it != cache.kernels.end()) {
       return it->second;
     }
   }
 
-  std::unique_lock lock(cache().mtx);
-  if (auto it = cache().kernels.find(kernel_name);
-      it != cache().kernels.end()) {
+  std::unique_lock lock(cache.mtx);
+  if (auto it = cache.kernels.find(kernel_name); it != cache.kernels.end()) {
     return it->second;
   }
   std::string source_code = source_builder();
@@ -125,10 +120,10 @@ void* compile(
   }
 
   // load library
-  cache().libs.emplace_back(shared_lib_path);
+  cache.libs.emplace_back(shared_lib_path);
 
   // Load function
-  void* fun = dlsym(cache().libs.back().lib, kernel_name.c_str());
+  void* fun = dlsym(cache.libs.back().lib, kernel_name.c_str());
   if (!fun) {
     std::ostringstream msg;
     msg << "[Compile::eval_cpu] Failed to load compiled function "
@@ -136,7 +131,7 @@ void* compile(
         << dlerror();
     throw std::runtime_error(msg.str());
   }
-  cache().kernels.insert({kernel_name, fun});
+  cache.kernels.insert({kernel_name, fun});
   return fun;
 }
 

mlx/backend/cpu/conv.cpp

@@ -921,7 +921,7 @@ void explicit_gemm_conv_1D_cpu(
 
   if (out.dtype() != float32) {
     gemm_out = array(out.shape(), float32, nullptr, {});
-    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
+    gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
     temps.push_back(gemm_out);
   }
 
@@ -1048,7 +1048,7 @@ void explicit_gemm_conv_2D_cpu(
 
   if (out.dtype() != float32) {
     gemm_out = array(out.shape(), float32, nullptr, {});
-    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
+    gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
     temps.push_back(gemm_out);
   }
 
@@ -1214,7 +1214,7 @@ void explicit_gemm_conv_ND_cpu(
 
   if (out.dtype() != float32) {
     gemm_out = array(out.shape(), float32, nullptr, {});
-    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
+    gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
     temps.push_back(gemm_out);
   }
 
@@ -1327,7 +1327,7 @@ void conv_3D_cpu(
 } // namespace
 
 void Convolution::eval_cpu(const std::vector<array>& inputs, array& out) {
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& in = inputs[0];
   auto& wt = inputs[1];

mlx/backend/cpu/distributed.cpp

@@ -30,7 +30,7 @@ void AllReduce::eval_cpu(
       if (in.is_donatable()) {
         out.copy_shared_buffer(in);
       } else {
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
       }
       return in;
     } else {
@@ -46,15 +46,8 @@ void AllReduce::eval_cpu(
     case Sum:
       distributed::detail::all_sum(group(), in, outputs[0], stream());
       break;
-    case Max:
-      distributed::detail::all_max(group(), in, outputs[0], stream());
-      break;
-    case Min:
-      distributed::detail::all_min(group(), in, outputs[0], stream());
-      break;
     default:
-      throw std::runtime_error(
-          "Only all reduce sum, min and max are supported for now");
+      throw std::runtime_error("Only all reduce sum is supported for now");
   }
 }
 
@@ -65,7 +58,7 @@ void AllGather::eval_cpu(
   assert(outputs.size() == 1);
 
   auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
-  outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
+  outputs[0].set_data(allocator::malloc_or_wait(outputs[0].nbytes()));
   distributed::detail::all_gather(group(), in, outputs[0], stream());
   if (copied) {
     auto& enc = cpu::get_command_encoder(stream());
@@ -94,7 +87,7 @@ void Recv::eval_cpu(
   assert(inputs.size() == 0);
   assert(outputs.size() == 1);
 
-  outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
+  outputs[0].set_data(allocator::malloc_or_wait(outputs[0].nbytes()));
   distributed::detail::recv(group(), outputs[0], src_, stream());
 }
 

mlx/backend/cpu/eigh.cpp

@@ -55,8 +55,9 @@ void eigh_impl(
       liwork = iwork;
     }
 
-    auto work_buf = array::Data{allocator::malloc(sizeof(T) * lwork)};
-    auto iwork_buf = array::Data{allocator::malloc(sizeof(int) * liwork)};
+    auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
+    auto iwork_buf =
+        array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
     for (size_t i = 0; i < size / (N * N); ++i) {
       syevd<T>(
           &jobz,
@@ -97,7 +98,7 @@ void Eigh::eval_cpu(
       ? outputs[1]
       : array(a.shape(), a.dtype(), nullptr, {});
 
-  values.set_data(allocator::malloc(values.nbytes()));
+  values.set_data(allocator::malloc_or_wait(values.nbytes()));
 
   copy(
       a,

mlx/backend/cpu/encoder.h

@@ -9,9 +9,6 @@
 
 namespace mlx::core::cpu {
 
-// Number of dispatches per scheduler task
-constexpr int DISPATCHES_PER_TASK = 10;
-
 struct CommandEncoder {
   CommandEncoder(Stream stream) : stream_(stream) {}
 
@@ -42,24 +39,13 @@ struct CommandEncoder {
 
   template <class F, class... Args>
   void dispatch(F&& f, Args&&... args) {
-    num_ops_ = (num_ops_ + 1) % DISPATCHES_PER_TASK;
     auto task = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
-    if (num_ops_ == 0) {
-      scheduler::notify_new_task(stream_);
-      auto task_wrap = [s = stream_, task = std::move(task)]() mutable {
-        task();
-        scheduler::notify_task_completion(s);
-      };
-      scheduler::enqueue(stream_, std::move(task_wrap));
-    } else {
-      scheduler::enqueue(stream_, std::move(task));
-    }
+    scheduler::enqueue(stream_, std::move(task));
   }
 
  private:
   Stream stream_;
   std::vector<array> temporaries_;
-  int num_ops_{0};
 };
 
 CommandEncoder& get_command_encoder(Stream stream);

mlx/backend/cpu/eval.cpp

@@ -33,8 +33,12 @@ void eval(array& arr) {
     buffers.erase(it);
   }
   auto& encoder = cpu::get_command_encoder(s);
-  encoder.dispatch([buffers = std::move(buffers),
-                    temps = std::move(encoder.temporaries())]() {});
+  scheduler::notify_new_task(s);
+  encoder.dispatch([s,
+                    buffers = std::move(buffers),
+                    temps = std::move(encoder.temporaries())]() {
+    scheduler::notify_task_completion(s);
+  });
 }
 
 } // namespace mlx::core::cpu

mlx/backend/cpu/fft.cpp

@@ -22,7 +22,7 @@ void FFT::eval_cpu(const std::vector<array>& inputs, array& out) {
     s *= out.itemsize();
   }
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   std::vector<size_t> shape;
   if (out.dtype() == float32) {

mlx/backend/cpu/gemms/no_bf16.cpp

@@ -0,0 +1,27 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cpu/gemm.h"
+
+namespace mlx::core {
+
+template <>
+void matmul<bfloat16_t>(
+    const bfloat16_t*,
+    const bfloat16_t*,
+    bfloat16_t*,
+    bool,
+    bool,
+    size_t,
+    size_t,
+    size_t,
+    float,
+    float,
+    size_t,
+    const Shape&,
+    const Strides&,
+    const Shape&,
+    const Strides&) {
+  throw std::runtime_error("[Matmul::eval_cpu] bfloat16 not supported.");
+}
+
+} // namespace mlx::core

mlx/backend/cpu/gemms/no_fp16.cpp

@@ -0,0 +1,27 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cpu/gemm.h"
+
+namespace mlx::core {
+
+template <>
+void matmul<float16_t>(
+    const float16_t*,
+    const float16_t*,
+    float16_t*,
+    bool,
+    bool,
+    size_t,
+    size_t,
+    size_t,
+    float,
+    float,
+    size_t,
+    const Shape&,
+    const Strides&,
+    const Shape&,
+    const Strides&) {
+  throw std::runtime_error("[Matmul::eval_cpu] float16 not supported.");
+}
+
+} // namespace mlx::core

mlx/backend/cpu/gemms/simd_bf16.cpp

@@ -1,45 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/gemm.h"
-#include "mlx/backend/cpu/gemms/simd_gemm.h"
-
-namespace mlx::core {
-
-template <>
-void matmul<bfloat16_t>(
-    const bfloat16_t* a,
-    const bfloat16_t* b,
-    bfloat16_t* out,
-    bool a_transposed,
-    bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
-    float alpha,
-    float beta,
-    size_t batch_size,
-    const Shape& a_shape,
-    const Strides& a_strides,
-    const Shape& b_shape,
-    const Strides& b_strides) {
-  auto ndim = a_shape.size();
-  size_t M = a_shape[ndim - 2];
-  size_t N = b_shape[ndim - 1];
-  size_t K = a_shape[ndim - 1];
-  for (int i = 0; i < batch_size; ++i) {
-    simd_gemm<bfloat16_t, float>(
-        a + elem_to_loc(M * K * i, a_shape, a_strides),
-        b + elem_to_loc(K * N * i, b_shape, b_strides),
-        out + M * N * i,
-        a_transposed,
-        b_transposed,
-        M,
-        N,
-        K,
-        alpha,
-        beta);
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cpu/gemms/simd_fp16.cpp

@@ -1,45 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/gemm.h"
-#include "mlx/backend/cpu/gemms/simd_gemm.h"
-
-namespace mlx::core {
-
-template <>
-void matmul<float16_t>(
-    const float16_t* a,
-    const float16_t* b,
-    float16_t* out,
-    bool a_transposed,
-    bool b_transposed,
-    size_t lda,
-    size_t ldb,
-    size_t ldc,
-    float alpha,
-    float beta,
-    size_t batch_size,
-    const Shape& a_shape,
-    const Strides& a_strides,
-    const Shape& b_shape,
-    const Strides& b_strides) {
-  auto ndim = a_shape.size();
-  size_t M = a_shape[ndim - 2];
-  size_t N = b_shape[ndim - 1];
-  size_t K = a_shape[ndim - 1];
-  for (int i = 0; i < batch_size; ++i) {
-    simd_gemm<float16_t, float>(
-        a + elem_to_loc(M * K * i, a_shape, a_strides),
-        b + elem_to_loc(K * N * i, b_shape, b_strides),
-        out + M * N * i,
-        a_transposed,
-        b_transposed,
-        M,
-        N,
-        K,
-        alpha,
-        beta);
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cpu/gemms/simd_gemm.h

@@ -1,139 +0,0 @@
-// Copyright © 2025 Apple Inc.
-#pragma once
-
-#include "mlx/backend/cpu/simd/simd.h"
-
-namespace mlx::core {
-
-inline int ceildiv(int a, int b) {
-  return (a + b - 1) / b;
-}
-
-template <int block_size, typename T, typename AccT>
-void load_block(
-    const T* in,
-    AccT* out,
-    int M,
-    int N,
-    int i,
-    int j,
-    bool transpose) {
-  if (transpose) {
-    for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
-      for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
-        out[jj * block_size + ii] =
-            in[(i * block_size + ii) * N + j * block_size + jj];
-      }
-    }
-  } else {
-    for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
-      for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
-        out[ii * block_size + jj] =
-            in[(i * block_size + ii) * N + j * block_size + jj];
-      }
-    }
-  }
-}
-
-template <typename T, typename AccT>
-void simd_gemm(
-    const T* a,
-    const T* b,
-    T* c,
-    bool a_trans,
-    bool b_trans,
-    int M,
-    int N,
-    int K,
-    float alpha,
-    float beta) {
-  constexpr int block_size = 16;
-  constexpr int simd_size = simd::max_size<AccT>;
-  static_assert(
-      (block_size % simd_size) == 0,
-      "Block size must be divisible by SIMD size");
-
-  int last_k_block_size = K - block_size * (K / block_size);
-  int last_k_simd_block = (last_k_block_size / simd_size) * simd_size;
-  for (int i = 0; i < ceildiv(M, block_size); i++) {
-    for (int j = 0; j < ceildiv(N, block_size); j++) {
-      AccT c_block[block_size * block_size] = {0.0};
-      AccT a_block[block_size * block_size];
-      AccT b_block[block_size * block_size];
-
-      int k = 0;
-      for (; k < K / block_size; k++) {
-        // Load a and b blocks
-        if (a_trans) {
-          load_block<block_size>(a, a_block, K, M, k, i, true);
-        } else {
-          load_block<block_size>(a, a_block, M, K, i, k, false);
-        }
-        if (b_trans) {
-          load_block<block_size>(b, b_block, N, K, j, k, false);
-        } else {
-          load_block<block_size>(b, b_block, K, N, k, j, true);
-        }
-
-        // Multiply and accumulate
-        for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
-          for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
-            for (int kk = 0; kk < block_size; kk += simd_size) {
-              auto av =
-                  simd::load<AccT, simd_size>(a_block + ii * block_size + kk);
-              auto bv =
-                  simd::load<AccT, simd_size>(b_block + jj * block_size + kk);
-              c_block[ii * block_size + jj] += simd::sum(av * bv);
-            }
-          }
-        }
-      }
-      if (last_k_block_size) {
-        // Load a and b blocks
-        if (a_trans) {
-          load_block<block_size>(a, a_block, K, M, k, i, true);
-        } else {
-          load_block<block_size>(a, a_block, M, K, i, k, false);
-        }
-        if (b_trans) {
-          load_block<block_size>(b, b_block, N, K, j, k, false);
-        } else {
-          load_block<block_size>(b, b_block, K, N, k, j, true);
-        }
-
-        // Multiply and accumulate
-        for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
-          for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
-            int kk = 0;
-            for (; kk < last_k_simd_block; kk += simd_size) {
-              auto av =
-                  simd::load<AccT, simd_size>(a_block + ii * block_size + kk);
-              auto bv =
-                  simd::load<AccT, simd_size>(b_block + jj * block_size + kk);
-              c_block[ii * block_size + jj] += simd::sum(av * bv);
-            }
-            for (; kk < last_k_block_size; ++kk) {
-              c_block[ii * block_size + jj] +=
-                  a_block[ii * block_size + kk] * b_block[jj * block_size + kk];
-            }
-          }
-        }
-      }
-
-      // Store
-      for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
-        for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
-          auto c_idx = (i * block_size + ii) * N + j * block_size + jj;
-          if (beta != 0) {
-            c[c_idx] = static_cast<T>(
-                alpha * c_block[ii * block_size + jj] + beta * c[c_idx]);
-          } else {
-            c[c_idx] = static_cast<T>(alpha * c_block[ii * block_size + jj]);
-          }
-        }
-      }
-    }
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cpu/indexing.cpp

@@ -197,7 +197,7 @@ void dispatch_gather(
 }
 
 void Gather::eval_cpu(const std::vector<array>& inputs, array& out) {
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& src = inputs[0];
   std::vector<array> inds;
@@ -354,7 +354,7 @@ void dispatch_gather_axis(
 }
 
 void GatherAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& src = inputs[0];
   auto& inds = inputs[1];

mlx/backend/cpu/inverse.cpp

@@ -11,7 +11,7 @@ namespace mlx::core {
 template <typename T>
 void general_inv(T* inv, int N) {
   int info;
-  auto ipiv = array::Data{allocator::malloc(sizeof(int) * N)};
+  auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
   // Compute LU factorization.
   getrf<T>(
       /* m = */ &N,
@@ -49,7 +49,7 @@ void general_inv(T* inv, int N) {
   }
 
   const int lwork = workspace_size;
-  auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
+  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
 
   // Compute inverse.
   getri<T>(

mlx/backend/cpu/logsumexp.cpp

@@ -1,140 +0,0 @@
-// Copyright © 2023-2024 Apple Inc.
-
-#include <cassert>
-#include <cmath>
-
-#include "mlx/backend/cpu/copy.h"
-#include "mlx/backend/cpu/encoder.h"
-#include "mlx/backend/cpu/simd/simd.h"
-#include "mlx/primitives.h"
-#include "mlx/types/limits.h"
-
-namespace mlx::core {
-
-namespace {
-
-using namespace mlx::core::simd;
-
-template <typename T, typename AccT>
-void logsumexp(const array& in, array& out, Stream stream) {
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(in);
-  encoder.set_output_array(out);
-
-  const T* in_ptr = in.data<T>();
-  T* out_ptr = out.data<T>();
-
-  int M = in.shape().back();
-  int L = in.data_size() / M;
-
-  encoder.dispatch([in_ptr, out_ptr, M, L]() mutable {
-    constexpr int N = std::min(max_size<AccT>, max_size<T>);
-
-    const T* current_in_ptr;
-
-    for (int i = 0; i < L; i++, in_ptr += M, out_ptr += 1) {
-      // Find the maximum
-      current_in_ptr = in_ptr;
-      Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
-      size_t s = M;
-      while (s >= N) {
-        Simd<AccT, N> vals = load<T, N>(current_in_ptr);
-        vmaximum = maximum(vals, vmaximum);
-        current_in_ptr += N;
-        s -= N;
-      }
-
-      AccT maximum = max(vmaximum);
-      while (s-- > 0) {
-        maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
-        current_in_ptr++;
-      }
-
-      // Compute the normalizer and the exponentials
-      Simd<AccT, N> vnormalizer(0.0);
-      current_in_ptr = in_ptr;
-      s = M;
-      while (s >= N) {
-        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
-        vexp = exp(vexp - maximum);
-        vnormalizer = vnormalizer + vexp;
-        current_in_ptr += N;
-        s -= N;
-      }
-      AccT normalizer = sum(vnormalizer);
-      while (s-- > 0) {
-        AccT _exp = std::exp(*current_in_ptr - maximum);
-        normalizer += _exp;
-        current_in_ptr++;
-      }
-      // Normalize
-      *out_ptr = std::isinf(maximum)
-          ? static_cast<T>(maximum)
-          : static_cast<T>(std::log(normalizer) + maximum);
-    }
-  });
-}
-
-} // namespace
-
-void LogSumExp::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-
-  // Make sure that the last dimension is contiguous
-  auto s = stream();
-  auto& encoder = cpu::get_command_encoder(s);
-  auto ensure_contiguous = [&s, &encoder](const array& x) {
-    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
-      return x;
-    } else {
-      auto x_copy = array(x.shape(), x.dtype(), nullptr, {});
-      copy(x, x_copy, CopyType::General, s);
-      encoder.add_temporary(x_copy);
-      return x_copy;
-    }
-  };
-
-  auto in = ensure_contiguous(inputs[0]);
-  if (in.flags().row_contiguous) {
-    out.set_data(allocator::malloc(out.nbytes()));
-  } else {
-    auto n = in.shape(-1);
-    auto flags = in.flags();
-    auto strides = in.strides();
-    for (auto& s : strides) {
-      s /= n;
-    }
-    bool col_contig = strides[0] == 1;
-    for (int i = 1; col_contig && i < strides.size(); ++i) {
-      col_contig &=
-          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
-    }
-    flags.col_contiguous = col_contig;
-    out.set_data(
-        allocator::malloc(in.nbytes() / n),
-        in.data_size() / n,
-        std::move(strides),
-        flags);
-  }
-
-  switch (in.dtype()) {
-    case float32:
-      logsumexp<float, float>(in, out, stream());
-      break;
-    case float16:
-      logsumexp<float16_t, float>(in, out, stream());
-      break;
-    case bfloat16:
-      logsumexp<bfloat16_t, float>(in, out, stream());
-      break;
-    case float64:
-      logsumexp<double, double>(in, out, stream());
-      break;
-    default:
-      throw std::runtime_error(
-          "[logsumexp] only supports floating point types");
-      break;
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cpu/luf.cpp

@@ -30,7 +30,8 @@ void luf_impl(
   auto strides = lu.strides();
   strides[ndim - 1] = M;
   strides[ndim - 2] = 1;
-  lu.set_data(allocator::malloc(lu.nbytes()), lu.nbytes(), strides, flags);
+  lu.set_data(
+      allocator::malloc_or_wait(lu.nbytes()), lu.nbytes(), strides, flags);
   copy_inplace(
       a,
       lu,
@@ -43,8 +44,8 @@ void luf_impl(
       stream);
 
   auto a_ptr = lu.data<T>();
-  pivots.set_data(allocator::malloc(pivots.nbytes()));
-  row_indices.set_data(allocator::malloc(row_indices.nbytes()));
+  pivots.set_data(allocator::malloc_or_wait(pivots.nbytes()));
+  row_indices.set_data(allocator::malloc_or_wait(row_indices.nbytes()));
   auto pivots_ptr = pivots.data<uint32_t>();
   auto row_indices_ptr = row_indices.data<uint32_t>();
   size_t num_matrices = a.size() / (M * N);

mlx/backend/cpu/masked_mm.cpp

@@ -59,7 +59,7 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     throw std::runtime_error(
         "[BlockMaskedMM::eval] Currently only supports float32.");
   }
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& a_pre = inputs[0];
   auto& b_pre = inputs[1];
@@ -318,7 +318,7 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     throw std::runtime_error(
         "[GatherMM::eval] Currently only supports float32.");
   }
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& a_pre = inputs[0];
   auto& b_pre = inputs[1];

mlx/backend/cpu/matmul.cpp

@@ -115,7 +115,7 @@ void matmul_general(
 }
 
 void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   if (inputs[0].shape(-1) == 0) {
     auto& encoder = cpu::get_command_encoder(stream());
     encoder.set_output_array(out);

mlx/backend/cpu/primitives.cpp

@@ -21,7 +21,7 @@ namespace mlx::core {
 void reshape(const array& in, array& out) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
   if (copy_necessary) {
-    out.set_data(allocator::malloc(out.nbytes()));
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
     copy_inplace(in, out, CopyType::General, out.primitive().stream());
   } else {
     shared_buffer_reshape(in, out_strides, out);
@@ -39,7 +39,7 @@ static std::pair<array, bool> compute_dynamic_offset(
   if (donate) {
     offset.copy_shared_buffer(indices);
   } else {
-    offset.set_data(allocator::malloc(offset.itemsize()));
+    offset.set_data(allocator::malloc_or_wait(offset.itemsize()));
   }
 
   auto& encoder = cpu::get_command_encoder(stream);
@@ -124,7 +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);
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   switch (out.dtype()) {
     case bool_:
       throw std::runtime_error("Bool type unsupported for arange.");
@@ -186,7 +186,7 @@ void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
   std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto strides = out.strides();
   auto flags = out.flags();
@@ -205,10 +205,8 @@ 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;
-  if (in.buffer_size() <= out.nbytes() + extra_bytes &&
-      (in.flags().row_contiguous ||
-       (allow_col_major_ && in.flags().col_contiguous))) {
+  if (in.flags().row_contiguous ||
+      (allow_col_major_ && in.flags().col_contiguous)) {
     out.copy_shared_buffer(in);
   } else {
     copy(in, out, CopyType::General, stream());
@@ -278,7 +276,7 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   size_t elems_per_key = out.size() / num_keys;
   size_t bytes_per_key = out.itemsize() * elems_per_key;
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto kptr = inputs[0].data<uint32_t>();
   auto cptr = out.data<char>();
@@ -337,7 +335,7 @@ void DynamicSlice::eval_cpu(const std::vector<array>& inputs, array& out) {
     return;
   }
   auto& in = inputs[0];
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   auto [in_offset, donated] =
       compute_dynamic_offset(inputs[1], in.strides(), axes_, stream());
   copy_inplace(
@@ -452,7 +450,7 @@ void View::eval_cpu(const std::vector<array>& inputs, array& out) {
   } else {
     auto tmp = array(
         in.shape(), in.dtype() == bool_ ? uint8 : in.dtype(), nullptr, {});
-    tmp.set_data(allocator::malloc(tmp.nbytes()));
+    tmp.set_data(allocator::malloc_or_wait(tmp.nbytes()));
     if (in.dtype() == bool_) {
       auto in_tmp = array(in.shape(), uint8, nullptr, {});
       in_tmp.copy_shared_buffer(in);

mlx/backend/cpu/qrf.cpp

@@ -25,11 +25,12 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
   auto strides = in.strides();
   strides[in.ndim() - 2] = 1;
   strides[in.ndim() - 1] = M;
-  in.set_data(allocator::malloc(in.nbytes()), in.nbytes(), strides, flags);
+  in.set_data(
+      allocator::malloc_or_wait(in.nbytes()), in.nbytes(), strides, flags);
   copy_inplace(a, in, CopyType::GeneralGeneral, stream);
   auto& encoder = cpu::get_command_encoder(stream);
-  q.set_data(allocator::malloc(q.nbytes()));
-  r.set_data(allocator::malloc(r.nbytes()));
+  q.set_data(allocator::malloc_or_wait(q.nbytes()));
+  r.set_data(allocator::malloc_or_wait(r.nbytes()));
 
   auto in_ptr = in.data<T>();
   auto r_ptr = r.data<T>();
@@ -40,7 +41,8 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
   encoder.set_output_array(r);
   encoder.dispatch([in_ptr, q_ptr, r_ptr, M, N, lda, num_matrices]() {
     int num_reflectors = std::min(M, N);
-    auto tau = allocator::malloc(sizeof(T) * num_matrices * num_reflectors);
+    auto tau =
+        allocator::malloc_or_wait(sizeof(T) * num_matrices * num_reflectors);
 
     T optimal_work;
     int lwork = -1;
@@ -51,7 +53,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
 
     // Update workspace size
     lwork = optimal_work;
-    auto work = allocator::malloc(sizeof(T) * lwork);
+    auto work = allocator::malloc_or_wait(sizeof(T) * lwork);
 
     // Loop over matrices
     for (int i = 0; i < num_matrices; ++i) {
@@ -94,7 +96,7 @@ void qrf_impl(const array& a, array& q, array& r, Stream stream) {
         &lwork,
         &info);
     lwork = optimal_work;
-    work = allocator::malloc(sizeof(T) * lwork);
+    work = allocator::malloc_or_wait(sizeof(T) * lwork);
 
     // Loop over matrices
     for (int i = 0; i < num_matrices; ++i) {

mlx/backend/cpu/quantized.cpp

@@ -515,7 +515,7 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto scales = ensure_row_contiguous(scales_pre);
   auto biases = ensure_row_contiguous(biases_pre);
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.add_temporaries(std::move(temps));
@@ -565,7 +565,7 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto scales = ensure_row_contiguous_last_dims(scales_pre);
   auto biases = ensure_row_contiguous_last_dims(biases_pre);
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.add_temporaries(std::move(temps));
@@ -691,12 +691,12 @@ void fast::AffineQuantize::eval_cpu(
 
   auto [w, copied] = ensure_row_contiguous(inputs[0]);
   auto& out = outputs[0];
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& scales = outputs[1];
   auto& biases = outputs[2];
-  scales.set_data(allocator::malloc(scales.nbytes()));
-  biases.set_data(allocator::malloc(biases.nbytes()));
+  scales.set_data(allocator::malloc_or_wait(scales.nbytes()));
+  biases.set_data(allocator::malloc_or_wait(biases.nbytes()));
   auto& encoder = cpu::get_command_encoder(stream());
   if (copied) {
     encoder.add_temporary(w);

mlx/backend/cpu/reduce.cpp

@@ -433,7 +433,7 @@ void reduce_dispatch_min_max(
 void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.set_input_array(in);
   encoder.set_output_array(out);

mlx/backend/cpu/scan.cpp

@@ -3,7 +3,6 @@
 #include <cassert>
 
 #include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/binary_ops.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
@@ -227,16 +226,6 @@ void scan_dispatch(
       scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
       break;
     }
-    case Scan::LogAddExp: {
-      auto op = [](U a, T b) {
-        return detail::LogAddExp{}(a, static_cast<U>(b));
-      };
-      auto init = (issubdtype(in.dtype(), floating))
-          ? static_cast<U>(-std::numeric_limits<float>::infinity())
-          : std::numeric_limits<U>::min();
-      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
-      break;
-    }
   }
 }
 
@@ -255,7 +244,7 @@ void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
     in = arr_copy;
     encoder.add_temporary(arr_copy);
   }
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   encoder.set_input_array(in);
   encoder.set_output_array(out);
@@ -330,8 +319,7 @@ void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
             reduce_type_, in, out, axis_, reverse_, inclusive_);
         break;
       case complex64:
-        scan_dispatch<complex64_t, complex64_t>(
-            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        throw std::runtime_error("Scan ops do not support complex types yet");
         break;
     }
   });

mlx/backend/cpu/simd/accelerate_fp16_simd.h

@@ -17,7 +17,7 @@ struct ScalarT<float16_t, N> {
 #endif
 
 template <>
-inline constexpr int max_size<float16_t> = N;
+static constexpr int max_size<float16_t> = N;
 
 #define SIMD_FP16_DEFAULT_UNARY(op)                    \
   template <>                                          \

mlx/backend/cpu/simd/accelerate_simd.h

@@ -83,25 +83,25 @@ struct Simd {
 // Values chosen based on benchmarks on M3 Max
 // TODO: consider choosing these more optimally
 template <>
-inline constexpr int max_size<int8_t> = 16;
+static constexpr int max_size<int8_t> = 16;
 template <>
-inline constexpr int max_size<int16_t> = 16;
+static constexpr int max_size<int16_t> = 16;
 template <>
-inline constexpr int max_size<int> = 8;
+static constexpr int max_size<int> = 8;
 template <>
-inline constexpr int max_size<int64_t> = 4;
+static constexpr int max_size<int64_t> = 4;
 template <>
-inline constexpr int max_size<uint8_t> = 16;
+static constexpr int max_size<uint8_t> = 16;
 template <>
-inline constexpr int max_size<uint16_t> = 16;
+static constexpr int max_size<uint16_t> = 16;
 template <>
-inline constexpr int max_size<uint32_t> = 8;
+static constexpr int max_size<uint32_t> = 8;
 template <>
-inline constexpr int max_size<uint64_t> = 4;
+static constexpr int max_size<uint64_t> = 4;
 template <>
-inline constexpr int max_size<float> = 8;
+static constexpr int max_size<float> = 8;
 template <>
-inline constexpr int max_size<double> = 4;
+static constexpr int max_size<double> = 4;
 
 #define SIMD_DEFAULT_UNARY(name, op) \
   template <typename T, int N>       \

mlx/backend/cpu/simd/base_simd.h

@@ -87,45 +87,14 @@ DEFAULT_UNARY(cosh, std::cosh)
 DEFAULT_UNARY(expm1, std::expm1)
 DEFAULT_UNARY(floor, std::floor)
 DEFAULT_UNARY(log, std::log)
+DEFAULT_UNARY(log2, std::log2)
 DEFAULT_UNARY(log10, std::log10)
+DEFAULT_UNARY(log1p, std::log1p)
 DEFAULT_UNARY(sinh, std::sinh)
 DEFAULT_UNARY(sqrt, std::sqrt)
 DEFAULT_UNARY(tan, std::tan)
 DEFAULT_UNARY(tanh, std::tanh)
 
-template <typename T>
-Simd<T, 1> log1p(Simd<T, 1> in) {
-  if constexpr (is_complex<T>) {
-    auto x = in.value.real();
-    auto y = in.value.imag();
-    auto zabs = std::abs(in.value);
-    auto theta = std::atan2(y, x + 1);
-    if (zabs < 0.5) {
-      auto r = x * (2 + x) + y * y;
-      if (r == 0) { // handle underflow
-        return Simd<T, 1>{T{x, theta}};
-      }
-      return Simd<T, 1>{T{((typeof(x))(0.5)) * std::log1p(r), theta}};
-    } else {
-      auto z0 = std::hypot(x + 1, y);
-      return Simd<T, 1>{T{std::log(z0), theta}};
-    }
-  } else {
-    return Simd<T, 1>{std::log1p(in.value)};
-  }
-}
-
-template <typename T>
-Simd<T, 1> log2(Simd<T, 1> in) {
-  if constexpr (is_complex<T>) {
-    auto out = std::log(in.value);
-    auto scale = decltype(out.real())(M_LN2);
-    return Simd<T, 1>{T{out.real() / scale, out.imag() / scale}};
-  } else {
-    return Simd<T, 1>{std::log2(in.value)};
-  }
-}
-
 template <typename T>
 Simd<T, 1> operator~(Simd<T, 1> in) {
   return ~in.value;

mlx/backend/cpu/softmax.cpp

@@ -119,12 +119,17 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   // Make sure that the last dimension is contiguous
   auto set_output = [s = stream(), &out](const array& x) {
-    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+    bool no_copy = x.strides()[x.ndim() - 1] == 1;
+    if (x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back());
+    }
+    if (no_copy) {
       if (x.is_donatable()) {
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            allocator::malloc(x.data_size() * x.itemsize()),
+            allocator::malloc_or_wait(x.data_size() * x.itemsize()),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -141,6 +146,18 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto in = set_output(inputs[0]);
 
   switch (in.dtype()) {
+    case bool_:
+    case uint8:
+    case uint16:
+    case uint32:
+    case uint64:
+    case int8:
+    case int16:
+    case int32:
+    case int64:
+      throw std::runtime_error(
+          "Softmax is defined only for floating point types");
+      break;
     case float32:
       softmax<float, float>(in, out, stream());
       break;
@@ -161,9 +178,9 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
     case float64:
       softmax<double, double>(in, out, stream());
       break;
-    default:
-      throw std::runtime_error(
-          "[softmax] Only defined for floating point types.");
+    case complex64:
+      throw std::invalid_argument(
+          "[Softmax] Not yet implemented for complex64");
       break;
   }
 }

mlx/backend/cpu/sort.cpp

@@ -288,7 +288,7 @@ void ArgSort::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
 
   // Allocate output
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.set_input_array(in);
@@ -379,7 +379,7 @@ void ArgPartition::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
 
   // Allocate output
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
   auto& encoder = cpu::get_command_encoder(stream());
   encoder.set_input_array(in);

mlx/backend/cpu/svd.cpp

@@ -50,9 +50,9 @@ void svd_impl(
     array& s = outputs[1];
     array& vt = outputs[2];
 
-    u.set_data(allocator::malloc(u.nbytes()));
-    s.set_data(allocator::malloc(s.nbytes()));
-    vt.set_data(allocator::malloc(vt.nbytes()));
+    u.set_data(allocator::malloc_or_wait(u.nbytes()));
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
+    vt.set_data(allocator::malloc_or_wait(vt.nbytes()));
 
     encoder.set_output_array(u);
     encoder.set_output_array(s);
@@ -64,7 +64,7 @@ void svd_impl(
   } else {
     array& s = outputs[0];
 
-    s.set_data(allocator::malloc(s.nbytes()));
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
 
     encoder.set_output_array(s);
 
@@ -91,7 +91,7 @@ void svd_impl(
 
     // Will contain the indices of eigenvectors that failed to converge (not
     // used here but required by lapack).
-    auto iwork = array::Data{allocator::malloc(sizeof(int) * 12 * K)};
+    auto iwork = array::Data{allocator::malloc_or_wait(sizeof(int) * 12 * K)};
 
     static const int lwork_query = -1;
 
@@ -132,7 +132,7 @@ void svd_impl(
     }
 
     const int lwork = workspace_dimension;
-    auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
+    auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
 
     // Loop over matrices.
     for (int i = 0; i < num_matrices; i++) {

mlx/backend/cpu/unary.cpp

@@ -1,8 +1,5 @@
 // Copyright © 2024 Apple Inc.
 
-// Required for using M_LN2 in MSVC.
-#define _USE_MATH_DEFINES
-
 #include <cassert>
 
 #include "mlx/backend/cpu/unary.h"

mlx/backend/cpu/unary.h

@@ -18,13 +18,13 @@ void set_unary_output_data(const array& in, array& out) {
     } else {
       auto size = in.data_size();
       out.set_data(
-          allocator::malloc(size * out.itemsize()),
+          allocator::malloc_or_wait(size * out.itemsize()),
           size,
           in.strides(),
           in.flags());
     }
   } else {
-    out.set_data(allocator::malloc(out.nbytes()));
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
 }
 

mlx/backend/cpu/unary_ops.h

@@ -86,14 +86,13 @@ struct Sign {
   template <int N, typename T>
   Simd<T, N> operator()(Simd<T, N> x) {
     auto z = Simd<T, N>{0};
-    auto o = Simd<T, N>{1};
-    auto m = Simd<T, N>{-1};
     if constexpr (std::is_unsigned_v<T>) {
-      return simd::select(x == z, z, o);
+      return x != z;
     } else if constexpr (std::is_same_v<T, complex64_t>) {
       return simd::select(x == z, x, Simd<T, N>(x / simd::abs(x)));
     } else {
-      return simd::select(x < z, m, simd::select(x > z, o, z));
+      return simd::select(
+          x < z, Simd<T, N>{-1}, simd::select(x > z, Simd<T, N>{1}, z));
     }
   }
   SINGLE()

mlx/backend/cuda/CMakeLists.txt

@@ -1,57 +0,0 @@
-# Filename rules in cuda backend:
-#
-# * Use .cu/.cuh if code contains device code, and .cpp/.h if not.
-# * Device-only kernel code should be put in kernels/ subdir.
-# * Files in kernels/ subdir should not include files outside.
-target_sources(
-  mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/event.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
-
-target_compile_definitions(mlx PUBLIC MLX_USE_CUDA)
-
-# Enable defining device lambda functions.
-target_compile_options(mlx
-                       PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
-
-# Compute capability 7 is required for synchronization between CPU/GPU with
-# managed memory. TODO: Add more architectures for potential performance gain.
-set(MLX_CUDA_ARCHITECTURES
-    "75;80"
-    CACHE STRING "CUDA architectures")
-message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
-set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
-                                     "${MLX_CUDA_ARCHITECTURES}")
-
-# Use fixed version of CCCL.
-FetchContent_Declare(
-  cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
-FetchContent_MakeAvailable(cccl)
-target_include_directories(mlx PRIVATE BEFORE "${cccl_SOURCE_DIR}/include")
-
-# Use fixed version of NVTX.
-FetchContent_Declare(
-  nvtx3
-  GIT_REPOSITORY https://github.com/NVIDIA/NVTX.git
-  GIT_TAG v3.1.1
-  GIT_SHALLOW TRUE
-  SOURCE_SUBDIR c EXCLUDE_FROM_ALL)
-FetchContent_MakeAvailable(nvtx3)
-target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)
-
-# Make cuda runtime APIs available in non-cuda files.
-find_package(CUDAToolkit REQUIRED)
-target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-
-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)

mlx/backend/cuda/allocator.cpp

@@ -1,154 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/backend/cuda/worker.h"
-
-#include <cuda_runtime.h>
-#include <fmt/format.h>
-
-#include <cassert>
-
-namespace mlx::core {
-
-namespace cu {
-
-CudaAllocator::CudaAllocator() {
-  // TODO: Set memory limit for multi-device.
-  size_t free, total;
-  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
-  memory_limit_ = total * 0.8;
-}
-
-Buffer CudaAllocator::malloc(size_t size) {
-  // TODO: Check memory limit.
-  auto* buf = new CudaBuffer{nullptr, size};
-  cudaError_t err = cudaMallocManaged(&buf->data, size);
-  if (err != cudaSuccess && err != cudaErrorMemoryAllocation) {
-    throw std::runtime_error(
-        fmt::format("cudaMallocManaged failed: {}.", cudaGetErrorString(err)));
-  }
-  std::lock_guard lock(mutex_);
-  active_memory_ += size;
-  peak_memory_ = std::max(active_memory_, peak_memory_);
-  return Buffer{buf};
-}
-
-void CudaAllocator::free(Buffer buffer) {
-  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
-  if (!buf) {
-    return;
-  }
-
-  // If free() is called from a unregistered thread, reschedule the call to
-  // worker.
-  {
-    std::lock_guard lock(worker_mutex_);
-    if (allowed_threads_.count(std::this_thread::get_id()) == 0) {
-      if (!worker_) {
-        worker_.reset(new Worker);
-      }
-      worker_->add_task([buffer]() { allocator().free(buffer); });
-      worker_->end_batch();
-      worker_->commit();
-      return;
-    }
-  }
-
-  size_t size = buf->size;
-  cudaFree(buf->data);
-  delete buf;
-  std::lock_guard lock(mutex_);
-  active_memory_ -= size;
-}
-
-size_t CudaAllocator::size(Buffer buffer) const {
-  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
-  if (!buf) {
-    return 0;
-  }
-  return buf->size;
-}
-
-void CudaAllocator::register_this_thread() {
-  std::lock_guard lock(worker_mutex_);
-  allowed_threads_.insert(std::this_thread::get_id());
-}
-
-size_t CudaAllocator::get_active_memory() const {
-  return active_memory_;
-}
-
-size_t CudaAllocator::get_peak_memory() const {
-  return peak_memory_;
-}
-
-void CudaAllocator::reset_peak_memory() {
-  std::lock_guard lock(mutex_);
-  peak_memory_ = 0;
-}
-
-size_t CudaAllocator::get_memory_limit() {
-  return memory_limit_;
-}
-
-size_t CudaAllocator::set_memory_limit(size_t limit) {
-  std::lock_guard lock(mutex_);
-  std::swap(limit, memory_limit_);
-  return limit;
-}
-
-CudaAllocator& allocator() {
-  // By creating the |allocator_| on heap, the destructor of CudaAllocator
-  // will not be called on exit and buffers in the cache will be leaked. This
-  // can save some time at program exit.
-  static CudaAllocator* allocator_ = new CudaAllocator;
-  return *allocator_;
-}
-
-} // namespace cu
-
-namespace allocator {
-
-Allocator& allocator() {
-  return cu::allocator();
-}
-
-void* Buffer::raw_ptr() {
-  if (!ptr_) {
-    return nullptr;
-  }
-  return static_cast<cu::CudaBuffer*>(ptr_)->data;
-}
-
-} // namespace allocator
-
-size_t get_active_memory() {
-  return cu::allocator().get_active_memory();
-}
-size_t get_peak_memory() {
-  return cu::allocator().get_peak_memory();
-}
-void reset_peak_memory() {
-  return cu::allocator().reset_peak_memory();
-}
-size_t set_memory_limit(size_t limit) {
-  return cu::allocator().set_memory_limit(limit);
-}
-size_t get_memory_limit() {
-  return cu::allocator().get_memory_limit();
-}
-
-// TODO: Implement buffer cache.
-size_t get_cache_memory() {
-  return 0;
-}
-size_t set_cache_limit(size_t) {
-  return 0;
-}
-size_t set_wired_limit(size_t) {
-  return 0;
-}
-void clear_cache() {}
-
-} // namespace mlx::core

mlx/backend/cuda/allocator.h

@@ -1,58 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/allocator.h"
-
-#include <mutex>
-#include <set>
-#include <thread>
-#include <utility>
-
-namespace mlx::core::cu {
-
-class Worker;
-
-using allocator::Buffer;
-
-// Stores cuda-managed unified memory.
-struct CudaBuffer {
-  void* data;
-  size_t size;
-};
-
-class CudaAllocator : public allocator::Allocator {
- public:
-  Buffer malloc(size_t size) override;
-  void free(Buffer buffer) override;
-  size_t size(Buffer buffer) const override;
-
-  // Register current thread as safe to free buffers.
-  // In cuda freeing a buffer implicitly synchronizes stream, and for threads
-  // that may be waited by gpu stream (for example cpu stream threads), freeing
-  // buffers there would result in dead lock.
-  void register_this_thread();
-
-  size_t get_active_memory() const;
-  size_t get_peak_memory() const;
-  void reset_peak_memory();
-  size_t get_memory_limit();
-  size_t set_memory_limit(size_t limit);
-
- private:
-  CudaAllocator();
-  friend CudaAllocator& allocator();
-
-  std::mutex worker_mutex_;
-  std::unique_ptr<Worker> worker_;
-  std::set<std::thread::id> allowed_threads_;
-
-  std::mutex mutex_;
-  size_t memory_limit_;
-  size_t active_memory_{0};
-  size_t peak_memory_{0};
-};
-
-CudaAllocator& allocator();
-
-} // namespace mlx::core::cu

mlx/backend/cuda/copy.cpp

@@ -1,26 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/copy.h"
-
-namespace mlx::core {
-
-void copy_gpu_inplace(
-    const array& in,
-    array& out,
-    const Shape& data_shape,
-    const Strides& strides_in_pre,
-    const Strides& strides_out_pre,
-    int64_t inp_offset,
-    int64_t out_offset,
-    CopyType ctype,
-    const Stream& s,
-    const std::optional<array>& dynamic_i_offset /* = std::nullopt */,
-    const std::optional<array>& dynamic_o_offset /* = std::nullopt */) {
-  throw std::runtime_error("copy_gpu_inplace not implemented in CUDA backend.");
-}
-
-void fill_gpu(const array& val, array& out, const Stream& s) {
-  throw std::runtime_error("fill_gpu not implemented in CUDA backend.");
-}
-
-} // namespace mlx::core

mlx/backend/cuda/device.cpp

@@ -1,117 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/backend/metal/metal.h"
-
-#include <fmt/format.h>
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace cu {
-
-DeviceStream::DeviceStream(Device& device) : device_(device), stream_(device) {}
-
-void DeviceStream::synchronize() {
-  cudaStreamSynchronize(stream_);
-}
-
-cudaStream_t DeviceStream::schedule_cuda_stream() {
-  // TODO: Return a stream that maximizes parallelism.
-  return stream_;
-}
-
-cudaStream_t DeviceStream::last_cuda_stream() {
-  return stream_;
-}
-
-CommandEncoder& DeviceStream::get_encoder() {
-  if (!encoder_) {
-    encoder_ = std::make_unique<CommandEncoder>(*this);
-  }
-  return *encoder_;
-}
-
-Device::Device(int device) : device_(device) {
-  // Validate the requirements of device.
-  int attr = 0;
-  cudaDeviceGetAttribute(&attr, cudaDevAttrConcurrentManagedAccess, device_);
-  if (attr != 1) {
-    throw std::runtime_error(fmt::format(
-        "Device {} does not support synchronization in managed memory.",
-        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;
-  if (current != device_) {
-    CHECK_CUDA_ERROR(cudaSetDevice(device_));
-    current = device_;
-  }
-}
-
-DeviceStream& Device::get_stream(Stream s) {
-  auto it = streams_.find(s.index);
-  if (it == streams_.end()) {
-    it = streams_.try_emplace(s.index, *this).first;
-  }
-  return it->second;
-}
-
-CommandEncoder::CommandEncoder(DeviceStream& s)
-    : device_(s.device()), stream_(s) {}
-
-void CommandEncoder::add_completed_handler(std::function<void()> task) {
-  worker_.add_task(std::move(task));
-}
-
-void CommandEncoder::end_encoding() {
-  if (!temporaries_.empty()) {
-    add_completed_handler([temporaries = std::move(temporaries_)]() {});
-  }
-
-  // There is no kernel running, run completion handlers immediately.
-  if (!has_gpu_work_) {
-    worker_.consume_in_this_thread();
-    return;
-  }
-  has_gpu_work_ = false;
-
-  // Put completion handlers in a batch.
-  worker_.end_batch();
-
-  // Signaling kernel completion is expensive, delay until enough batches.
-  // TODO: This number is arbitrarily picked, profile for a better stragety.
-  if (worker_.uncommited_batches() > 8) {
-    commit();
-  }
-}
-
-void CommandEncoder::commit() {
-  worker_.commit(stream_.last_cuda_stream());
-}
-
-Device& device(mlx::core::Device device) {
-  static std::unordered_map<int, Device> devices;
-  auto it = devices.find(device.index);
-  if (it == devices.end()) {
-    it = devices.try_emplace(device.index, device.index).first;
-  }
-  return it->second;
-}
-
-DeviceStream& get_stream(Stream s) {
-  return device(s.device).get_stream(s);
-}
-
-CommandEncoder& get_command_encoder(Stream s) {
-  return get_stream(s).get_encoder();
-}
-
-} // namespace cu
-
-} // namespace mlx::core

mlx/backend/cuda/device.h

@@ -1,131 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/array.h"
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/stream.h"
-
-#include <thrust/execution_policy.h>
-
-#include <unordered_map>
-
-namespace mlx::core::cu {
-
-class Device;
-class CommandEncoder;
-
-class DeviceStream {
- public:
-  explicit DeviceStream(Device& device);
-
-  DeviceStream(const DeviceStream&) = delete;
-  DeviceStream& operator=(const DeviceStream&) = delete;
-
-  // Wait until kernels in the stream complete.
-  void synchronize();
-
-  // Return a cuda stream for launching kernels.
-  cudaStream_t schedule_cuda_stream();
-
-  // Return the last cuda stream used.
-  cudaStream_t last_cuda_stream();
-
-  CommandEncoder& get_encoder();
-
-  Device& device() {
-    return device_;
-  }
-
- private:
-  Device& device_;
-  CudaStream stream_;
-  std::unique_ptr<CommandEncoder> encoder_;
-};
-
-class Device {
- public:
-  explicit Device(int device);
-
-  Device(const Device&) = delete;
-  Device& operator=(const Device&) = delete;
-
-  // Make this device the current cuda device, required by some cuda calls.
-  void make_current();
-
-  DeviceStream& get_stream(Stream s);
-
-  int cuda_device() const {
-    return device_;
-  }
-
- private:
-  int device_;
-  std::unordered_map<int, DeviceStream> streams_;
-};
-
-class CommandEncoder {
- public:
-  explicit CommandEncoder(DeviceStream& stream);
-
-  CommandEncoder(const CommandEncoder&) = delete;
-  CommandEncoder& operator=(const CommandEncoder&) = delete;
-
-  void set_input_array(const array& arr) {}
-  void set_output_array(const array& arr) {}
-
-  void add_temporary(const array& arr) {
-    temporaries_.push_back(arr.data_shared_ptr());
-  }
-
-  void add_completed_handler(std::function<void()> task);
-  void end_encoding();
-  void commit();
-
-  // Schedule a cuda stream for |fun| to launch kernels, and check error
-  // afterwards.
-  template <typename F>
-  void launch_kernel(F&& fun) {
-    launch_kernel(stream_.schedule_cuda_stream(), std::forward<F>(fun));
-  }
-
-  template <typename F>
-  void launch_kernel(cudaStream_t stream, F&& fun) {
-    device_.make_current();
-    fun(stream);
-    check_cuda_error("kernel launch", cudaGetLastError());
-    has_gpu_work_ = true;
-  }
-
-  Device& device() {
-    return device_;
-  }
-
-  DeviceStream& stream() {
-    return stream_;
-  }
-
-  bool has_gpu_work() const {
-    return has_gpu_work_;
-  }
-
- private:
-  Device& device_;
-  DeviceStream& stream_;
-  Worker worker_;
-  bool has_gpu_work_{false};
-  std::vector<std::shared_ptr<array::Data>> temporaries_;
-};
-
-Device& device(mlx::core::Device device);
-DeviceStream& get_stream(Stream s);
-CommandEncoder& get_command_encoder(Stream s);
-
-// Return an execution policy that does not sync for result.
-// Note that not all thrust APIs support async policy, confirm before using.
-inline auto thrust_policy(cudaStream_t stream) {
-  // TODO: Connect thrust's custom allocator with mlx's allocator.
-  return thrust::cuda::par_nosync.on(stream);
-}
-
-} // namespace mlx::core::cu

mlx/backend/cuda/dtype_utils.cuh

@@ -1,35 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuComplex.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-
-namespace mlx::core {
-
-// Maps CPU types to CUDA types.
-template <typename T>
-struct CTypeToCudaType {
-  using type = T;
-};
-
-template <>
-struct CTypeToCudaType<float16_t> {
-  using type = __half;
-};
-
-template <>
-struct CTypeToCudaType<bfloat16_t> {
-  using type = __nv_bfloat16;
-};
-
-template <>
-struct CTypeToCudaType<complex64_t> {
-  using type = cuComplex;
-};
-
-template <typename T>
-using cuda_type_t = typename CTypeToCudaType<T>::type;
-
-} // namespace mlx::core

mlx/backend/cuda/eval.cpp

@@ -1,68 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/eval.h"
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/available.h"
-#include "mlx/primitives.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core::gpu {
-
-bool is_available() {
-  return true;
-}
-
-void new_stream(Stream s) {
-  // Force initalization of cuda, so cuda runtime get destroyed at last.
-  cudaFree(nullptr);
-  // Ensure the static stream objects get created.
-  cu::get_command_encoder(s);
-  // The main thread is safe to free buffers.
-  cu::allocator().register_this_thread();
-}
-
-void eval(array& arr) {
-  nvtx3::scoped_range r("gpu::eval");
-  auto outputs = arr.outputs();
-  {
-    // If the array is a tracer hold a reference
-    // to its inputs so they don't get donated
-    std::vector<array> inputs;
-    if (arr.is_tracer()) {
-      inputs = arr.inputs();
-    }
-    arr.primitive().eval_gpu(arr.inputs(), outputs);
-  }
-
-  auto& encoder = cu::get_command_encoder(arr.primitive().stream());
-  if (encoder.has_gpu_work()) {
-    // Keep used buffers alive until kernel finishes running.
-    std::unordered_set<std::shared_ptr<array::Data>> buffers;
-    for (auto& in : arr.inputs()) {
-      buffers.insert(in.data_shared_ptr());
-    }
-    for (auto& s : arr.siblings()) {
-      buffers.insert(s.data_shared_ptr());
-    }
-    // Remove the output if it was donated to by an input.
-    if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
-      buffers.erase(it);
-    }
-    encoder.add_completed_handler([buffers = std::move(buffers)]() {});
-  }
-  encoder.end_encoding();
-}
-
-void finalize(Stream s) {
-  nvtx3::scoped_range r("gpu::finalize");
-  cu::get_command_encoder(s).commit();
-}
-
-void synchronize(Stream s) {
-  nvtx3::scoped_range r("gpu::synchronize");
-  cu::get_stream(s).synchronize();
-}
-
-} // namespace mlx::core::gpu

mlx/backend/cuda/event.cu

@@ -1,265 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/event.h"
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/event.h"
-#include "mlx/scheduler.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace cu {
-
-///////////////////////////////////////////////////////////////////////////////
-// CudaEvent implementations
-///////////////////////////////////////////////////////////////////////////////
-
-// Cuda event managed with RAII.
-class CudaEventHandle {
- public:
-  CudaEventHandle() {
-    CHECK_CUDA_ERROR(cudaEventCreateWithFlags(
-        &event_, cudaEventDisableTiming | cudaEventBlockingSync));
-  }
-
-  ~CudaEventHandle() {
-    CHECK_CUDA_ERROR(cudaEventDestroy(event_));
-  }
-
-  CudaEventHandle(const CudaEventHandle&) = delete;
-  CudaEventHandle& operator=(const CudaEventHandle&) = delete;
-
-  operator cudaEvent_t() const {
-    return event_;
-  }
-
- private:
-  cudaEvent_t event_;
-};
-
-CudaEvent::CudaEvent() : event_(std::make_shared<CudaEventHandle>()) {}
-
-void CudaEvent::wait() {
-  nvtx3::scoped_range r("cu::CudaEvent::wait");
-  if (!recorded_) {
-    throw std::runtime_error("Should not wait on a CudaEvent before record.");
-  }
-  cudaEventSynchronize(*event_);
-}
-
-void CudaEvent::wait(cudaStream_t stream) {
-  if (!recorded_) {
-    throw std::runtime_error("Should not wait on a CudaEvent before record.");
-  }
-  cudaStreamWaitEvent(stream, *event_);
-}
-
-void CudaEvent::wait(Stream s) {
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this]() mutable { wait(); });
-  } else {
-    wait(cu::get_stream(s).last_cuda_stream());
-  }
-}
-
-void CudaEvent::record(cudaStream_t stream) {
-  cudaEventRecord(*event_, stream);
-  recorded_ = true;
-}
-
-void CudaEvent::record(Stream s) {
-  if (s.device == mlx::core::Device::cpu) {
-    throw std::runtime_error("CudaEvent can not wait on cpu stream.");
-  } else {
-    record(cu::get_stream(s).last_cuda_stream());
-  }
-}
-
-bool CudaEvent::completed() const {
-  return cudaEventQuery(*event_) == cudaSuccess;
-}
-
-///////////////////////////////////////////////////////////////////////////////
-// SharedEvent implementations
-///////////////////////////////////////////////////////////////////////////////
-
-namespace {
-
-__host__ __device__ void event_wait(SharedEvent::Atomic* ac, uint64_t value) {
-  uint64_t current;
-  while ((current = ac->load()) < value) {
-    ac->wait(current);
-  }
-}
-
-__host__ __device__ void event_signal(SharedEvent::Atomic* ac, uint64_t value) {
-  ac->store(value);
-  ac->notify_all();
-}
-
-__global__ void event_wait_kernel(SharedEvent::Atomic* ac, uint64_t value) {
-  event_wait(ac, value);
-}
-
-__global__ void event_signal_kernel(SharedEvent::Atomic* ac, uint64_t value) {
-  event_signal(ac, value);
-}
-
-} // namespace
-
-SharedEvent::SharedEvent() {
-  // Allocate cuda::atomic on managed memory.
-  allocator::Buffer buffer = allocator::malloc(sizeof(Atomic));
-  Atomic* ac = static_cast<Atomic*>(buffer.raw_ptr());
-  new (ac) Atomic(0);
-  ac_ = std::shared_ptr<Atomic>(ac, [buffer](Atomic* ptr) {
-    ptr->~Atomic();
-    allocator::free(buffer);
-  });
-}
-
-void SharedEvent::wait(uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::wait");
-  event_wait(ac_.get(), value);
-}
-
-void SharedEvent::wait(cudaStream_t stream, uint64_t value) {
-  event_wait_kernel<<<1, 1, 0, stream>>>(ac_.get(), value);
-}
-
-void SharedEvent::wait(Stream s, uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::wait(s)");
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
-  } else {
-    auto& encoder = get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(),
-        [this, value](cudaStream_t stream) { wait(stream, value); });
-    encoder.add_completed_handler([ac = ac_]() {});
-    encoder.end_encoding();
-  }
-}
-
-void SharedEvent::signal(uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::signal");
-  event_signal(ac_.get(), value);
-}
-
-void SharedEvent::signal(cudaStream_t stream, uint64_t value) {
-  event_signal_kernel<<<1, 1, 0, stream>>>(ac_.get(), value);
-}
-
-void SharedEvent::signal(Stream s, uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::signal(s)");
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this, value]() mutable { signal(value); });
-  } else {
-    auto& encoder = get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(),
-        [this, value](cudaStream_t stream) { signal(stream, value); });
-    encoder.add_completed_handler([ac = ac_]() {});
-    encoder.end_encoding();
-  }
-}
-
-bool SharedEvent::is_signaled(uint64_t value) const {
-  nvtx3::scoped_range r("cu::SharedEvent::is_signaled");
-  return ac_->load() >= value;
-}
-
-uint64_t SharedEvent::value() const {
-  nvtx3::scoped_range r("cu::SharedEvent::value");
-  return ac_->load();
-}
-
-} // namespace cu
-
-///////////////////////////////////////////////////////////////////////////////
-// Event implementations
-///////////////////////////////////////////////////////////////////////////////
-
-namespace {
-
-struct EventImpl {
-  // CudaEvent is preferred when possible because it is fast, however we have
-  // to fallback to SharedEvent in following cases:
-  // 1. the event is used to wait/signal a cpu stream;
-  // 2. signal value other than 1 has been specified.
-  std::unique_ptr<cu::CudaEvent> cuda;
-  std::unique_ptr<cu::SharedEvent> shared;
-
-  bool is_created() const {
-    return cuda || shared;
-  }
-
-  void ensure_created(Stream s, uint64_t signal_value) {
-    if (is_created()) {
-      return;
-    }
-    if (s.device == mlx::core::Device::cpu || signal_value > 1) {
-      nvtx3::mark("Using slow SharedEvent");
-      shared = std::make_unique<cu::SharedEvent>();
-    } else {
-      cuda = std::make_unique<cu::CudaEvent>();
-    }
-  }
-};
-
-} // namespace
-
-Event::Event(Stream s) : stream_(s) {
-  event_ = std::shared_ptr<void>(
-      new EventImpl(), [](void* ptr) { delete static_cast<EventImpl*>(ptr); });
-}
-
-void Event::wait() {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  assert(event->is_created());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->wait();
-  } else {
-    event->shared->wait(value());
-  }
-}
-
-void Event::wait(Stream s) {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  assert(event->is_created());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->wait(s);
-  } else {
-    event->shared->wait(s, value());
-  }
-}
-
-void Event::signal(Stream s) {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  event->ensure_created(s, value());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->record(s);
-  } else {
-    event->shared->signal(s, value());
-  }
-}
-
-bool Event::is_signaled() const {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  if (!event->is_created()) {
-    return false;
-  }
-  if (event->cuda) {
-    assert(value() == 1);
-    return event->cuda->recorded() && event->cuda->completed();
-  } else {
-    return event->shared->is_signaled(value());
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/event.h

@@ -1,66 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/stream.h"
-
-#include <cuda_runtime.h>
-#include <cuda/atomic>
-
-#include <memory>
-
-namespace mlx::core::cu {
-
-class CudaEventHandle;
-
-// Wrapper of native cuda event. It can synchronize between GPU streams, or wait
-// on GPU stream in CPU stream, but can not wait on CPU stream.
-class CudaEvent {
- public:
-  CudaEvent();
-
-  void wait();
-  void wait(cudaStream_t stream);
-  void wait(Stream s);
-  void record(cudaStream_t stream);
-  void record(Stream s);
-
-  // Return whether the recorded kernels have completed. Note that this method
-  // returns true if record() has not been called.
-  bool completed() const;
-
-  bool recorded() const {
-    return recorded_;
-  }
-
- private:
-  bool recorded_{false};
-  std::shared_ptr<CudaEventHandle> event_;
-};
-
-// Event that can synchronize between CPU and GPU. It is much slower than
-// CudaEvent so the latter should always be preferred when possible.
-class SharedEvent {
- public:
-  using Atomic = cuda::atomic<uint64_t>;
-
-  SharedEvent();
-
-  void wait(uint64_t value);
-  void wait(cudaStream_t stream, uint64_t value);
-  void wait(Stream s, uint64_t value);
-  void signal(uint64_t value);
-  void signal(cudaStream_t stream, uint64_t value);
-  void signal(Stream s, uint64_t value);
-  bool is_signaled(uint64_t value) const;
-  uint64_t value() const;
-
-  const std::shared_ptr<Atomic>& atomic() const {
-    return ac_;
-  }
-
- private:
-  std::shared_ptr<Atomic> ac_;
-};
-
-} // namespace mlx::core::cu

mlx/backend/cuda/fence.cu

@@ -1,70 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/event.h"
-#include "mlx/fence.h"
-#include "mlx/scheduler.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace {
-
-__host__ __device__ void busy_wait(cuda::atomic<uint64_t>* ac, uint64_t value) {
-  while (true) {
-    // In theory the atomic_thread_fence is not needed, but for CUDA 11 without
-    // it the load() may never return new value.
-    cuda::atomic_thread_fence(cuda::memory_order_seq_cst);
-    uint64_t current = ac->load();
-    if (current >= value) {
-      break;
-    }
-  }
-}
-
-__global__ void busy_wait_kernel(cuda::atomic<uint64_t>* ac, uint64_t value) {
-  busy_wait(ac, value);
-}
-
-} // namespace
-
-struct FenceImpl {
-  uint32_t count;
-  cu::SharedEvent event;
-};
-
-Fence::Fence(Stream s) {
-  fence_ = std::shared_ptr<void>(
-      new FenceImpl{0}, [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
-}
-
-void Fence::wait(Stream s, const array&) {
-  auto* fence = static_cast<FenceImpl*>(fence_.get());
-  // We can't use SharedEvent::wait because it could hang in CUDA 11, see also:
-  // https://github.com/ml-explore/mlx/issues/2137
-  const auto& ac = fence->event.atomic();
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [ac, count = fence->count]() {
-      nvtx3::scoped_range r("Fence::wait()");
-      busy_wait(ac.get(), count);
-    });
-  } else {
-    nvtx3::scoped_range r("Fence::wait(s)");
-    auto& encoder = cu::get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(), [&](cudaStream_t stream) {
-          busy_wait_kernel<<<1, 1, 0>>>(ac.get(), fence->count);
-        });
-    encoder.add_completed_handler([ac]() {});
-    encoder.end_encoding();
-  }
-}
-
-void Fence::update(Stream s, const array&) {
-  auto* fence = static_cast<FenceImpl*>(fence_.get());
-  fence->count++;
-  fence->event.signal(s, fence->count);
-}
-
-} // namespace mlx::core

mlx/backend/cuda/kernels/arange.cuh

@@ -1,15 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-namespace mlx::core::cu {
-
-template <typename T>
-struct Arange {
-  const T start;
-  const T step;
-
-  __device__ T operator()(uint32_t i) const {
-    return start + i * step;
-  }
-};
-
-} // namespace mlx::core::cu

mlx/backend/cuda/kernels/fp16_math.cuh

@@ -1,107 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuda_fp16.h>
-#include <cuda/std/limits>
-#include <cuda/std/type_traits>
-
-namespace mlx::core::cu {
-
-///////////////////////////////////////////////////////////////////////////////
-// Missing C++ operator overrides for CUDA 7.
-///////////////////////////////////////////////////////////////////////////////
-
-#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-
-#define MLX_DEFINE_BF16_OP(OP)                                           \
-  __forceinline__ __device__ __nv_bfloat16 operator OP(                  \
-      __nv_bfloat16 x, __nv_bfloat16 y) {                                \
-    return __float2bfloat16(__bfloat162float(x) OP __bfloat162float(y)); \
-  }
-
-#define MLX_DEFINE_BF16_CMP(OP)                                          \
-  __forceinline__ __device__ bool operator OP(                           \
-      __nv_bfloat16 x, __nv_bfloat16 y) {                                \
-    return __float2bfloat16(__bfloat162float(x) OP __bfloat162float(y)); \
-  }
-
-MLX_DEFINE_BF16_OP(+)
-MLX_DEFINE_BF16_OP(-)
-MLX_DEFINE_BF16_OP(*)
-MLX_DEFINE_BF16_OP(/)
-MLX_DEFINE_BF16_CMP(>)
-MLX_DEFINE_BF16_CMP(<)
-MLX_DEFINE_BF16_CMP(>=)
-MLX_DEFINE_BF16_CMP(<=)
-
-#undef MLX_DEFINE_BF16_OP
-#undef MLX_DEFINE_BF16_CMP
-
-#endif // CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-
-///////////////////////////////////////////////////////////////////////////////
-// Additional C++ operator overrides between half types and native types.
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, typename U>
-constexpr bool is_integral_except =
-    cuda::std::is_integral_v<T> && !cuda::std::is_same_v<T, U>;
-
-template <typename T, typename U>
-constexpr bool is_arithmetic_except =
-    cuda::std::is_arithmetic_v<T> && !cuda::std::is_same_v<T, U>;
-
-#define MLX_DEFINE_HALF_OP(HALF, HALF2FLOAT, FLOAT2HALF, OP)          \
-  template <                                                          \
-      typename T,                                                     \
-      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
-  __forceinline__ __device__ HALF operator OP(HALF x, T y) {          \
-    return FLOAT2HALF(HALF2FLOAT(x) OP static_cast<float>(y));        \
-  }                                                                   \
-  template <                                                          \
-      typename T,                                                     \
-      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
-  __forceinline__ __device__ HALF operator OP(T x, HALF y) {          \
-    return FLOAT2HALF(static_cast<float>(x) OP HALF2FLOAT(y));        \
-  }
-
-#define MLX_DEFINE_HALF_CMP(HALF, HALF2FLOAT, OP)                       \
-  template <                                                            \
-      typename T,                                                       \
-      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
-  __forceinline__ __device__ bool operator OP(HALF x, T y) {            \
-    return HALF2FLOAT(x) OP static_cast<float>(y);                      \
-  }                                                                     \
-  template <                                                            \
-      typename T,                                                       \
-      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
-  __forceinline__ __device__ bool operator OP(T x, HALF y) {            \
-    return static_cast<float>(y) OP HALF2FLOAT(x);                      \
-  }
-
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, +)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, -)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, *)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, /)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, +)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, -)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, *)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, /)
-MLX_DEFINE_HALF_CMP(__half, __half2float, <)
-MLX_DEFINE_HALF_CMP(__half, __half2float, >)
-MLX_DEFINE_HALF_CMP(__half, __half2float, <=)
-MLX_DEFINE_HALF_CMP(__half, __half2float, >=)
-MLX_DEFINE_HALF_CMP(__half, __half2float, ==)
-MLX_DEFINE_HALF_CMP(__half, __half2float, !=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, ==)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, !=)
-
-#undef MLX_DEFINE_HALF_OP
-#undef MLX_DEFINE_HALF_CMP
-
-} // namespace mlx::core::cu

mlx/backend/cuda/primitives.cu

@@ -1,163 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/dtype_utils.cuh"
-#include "mlx/backend/cuda/kernels/arange.cuh"
-#include "mlx/backend/cuda/kernels/fp16_math.cuh"
-#include "mlx/distributed/primitives.h"
-#include "mlx/dtype_utils.h"
-#include "mlx/fast_primitives.h"
-#include "mlx/primitives.h"
-
-#include <nvtx3/nvtx3.hpp>
-#include <thrust/device_ptr.h>
-#include <thrust/transform.h>
-
-#include <cassert>
-
-namespace mlx::core {
-
-void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Arange::eval_gpu");
-  assert(inputs.size() == 0);
-  out.set_data(allocator::malloc(out.nbytes()));
-  if (out.size() == 0) {
-    return;
-  }
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  encoder.set_output_array(out);
-  encoder.launch_kernel([&, this](cudaStream_t stream) {
-    MLX_SWITCH_INT_FLOAT_TYPES_CHECKED(out.dtype(), "Arange", CTYPE, {
-      using OutType = cuda_type_t<CTYPE>;
-      CTYPE step =
-          static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_);
-      thrust::transform(
-          cu::thrust_policy(stream),
-          thrust::counting_iterator<uint32_t>(0),
-          thrust::counting_iterator<uint32_t>(out.data_size()),
-          thrust::device_pointer_cast(out.data<OutType>()),
-          cu::Arange<OutType>{
-              static_cast<OutType>(start_), static_cast<OutType>(step)});
-    });
-  });
-}
-
-#define NO_GPU_MULTI(func)                                             \
-  void func::eval_gpu(                                                 \
-      const std::vector<array>& inputs, std::vector<array>& outputs) { \
-    throw std::runtime_error(#func " has no CUDA implementation.");    \
-  }
-
-#define NO_GPU(func)                                                  \
-  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
-    throw std::runtime_error(#func " has no CUDA implementation.");   \
-  }
-
-NO_GPU(Abs)
-NO_GPU(Add)
-NO_GPU(AddMM)
-NO_GPU(ArcCos)
-NO_GPU(ArcCosh)
-NO_GPU(ArcSin)
-NO_GPU(ArcSinh)
-NO_GPU(ArcTan)
-NO_GPU(ArcTan2)
-NO_GPU(ArcTanh)
-NO_GPU(ArgPartition)
-NO_GPU(ArgReduce)
-NO_GPU(ArgSort)
-NO_GPU(BitwiseBinary)
-NO_GPU(BitwiseInvert)
-NO_GPU(BlockMaskedMM)
-NO_GPU(Ceil)
-NO_GPU_MULTI(Compiled)
-NO_GPU(Conjugate)
-NO_GPU(Convolution)
-NO_GPU(Cos)
-NO_GPU(Cosh)
-NO_GPU(Divide)
-NO_GPU_MULTI(DivMod)
-NO_GPU(DynamicSlice)
-NO_GPU(DynamicSliceUpdate)
-NO_GPU(Remainder)
-NO_GPU(Equal)
-NO_GPU(Erf)
-NO_GPU(ErfInv)
-NO_GPU(Exp)
-NO_GPU(Expm1)
-NO_GPU(FFT)
-NO_GPU(Floor)
-NO_GPU(Gather)
-NO_GPU(GatherAxis)
-NO_GPU(GatherMM)
-NO_GPU(GatherQMM)
-NO_GPU(Greater)
-NO_GPU(GreaterEqual)
-NO_GPU(Hadamard)
-NO_GPU(Imag)
-NO_GPU(Less)
-NO_GPU(LessEqual)
-NO_GPU(Load)
-NO_GPU(Log)
-NO_GPU(Log1p)
-NO_GPU(LogicalNot)
-NO_GPU(LogicalAnd)
-NO_GPU(LogicalOr)
-NO_GPU(LogAddExp)
-NO_GPU(LogSumExp)
-NO_GPU_MULTI(LUF)
-NO_GPU(Matmul)
-NO_GPU(Maximum)
-NO_GPU(Minimum)
-NO_GPU(Multiply)
-NO_GPU(Negative)
-NO_GPU(NotEqual)
-NO_GPU(Partition)
-NO_GPU(Power)
-NO_GPU_MULTI(QRF)
-NO_GPU(QuantizedMatmul)
-NO_GPU(RandomBits)
-NO_GPU(Real)
-NO_GPU(Reduce)
-NO_GPU(Round)
-NO_GPU(Scan)
-NO_GPU(Scatter)
-NO_GPU(ScatterAxis)
-NO_GPU(Select)
-NO_GPU(Sigmoid)
-NO_GPU(Sign)
-NO_GPU(Sin)
-NO_GPU(Sinh)
-NO_GPU(SliceUpdate)
-NO_GPU(Softmax)
-NO_GPU(Sort)
-NO_GPU(Square)
-NO_GPU(Sqrt)
-NO_GPU(Subtract)
-NO_GPU_MULTI(SVD)
-NO_GPU(Tan)
-NO_GPU(Tanh)
-NO_GPU(Inverse)
-NO_GPU(Cholesky)
-NO_GPU_MULTI(Eigh)
-
-namespace fast {
-NO_GPU_MULTI(LayerNorm)
-NO_GPU_MULTI(LayerNormVJP)
-NO_GPU_MULTI(RMSNorm)
-NO_GPU_MULTI(RMSNormVJP)
-NO_GPU_MULTI(RoPE)
-NO_GPU(ScaledDotProductAttention)
-NO_GPU_MULTI(AffineQuantize)
-NO_GPU_MULTI(CustomKernel)
-} // namespace fast
-
-namespace distributed {
-NO_GPU_MULTI(AllReduce)
-NO_GPU_MULTI(AllGather)
-NO_GPU_MULTI(Send)
-NO_GPU_MULTI(Recv)
-} // namespace distributed
-
-} // namespace mlx::core

mlx/backend/cuda/slicing.cpp

@@ -1,15 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/slicing.h"
-
-namespace mlx::core {
-
-void concatenate_gpu(
-    const std::vector<array>& inputs,
-    array& out,
-    int axis,
-    const Stream& s) {
-  throw std::runtime_error("concatenate_gpu not implemented in CUDA backend.");
-}
-
-} // namespace mlx::core

mlx/backend/cuda/utils.cpp

@@ -1,26 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/backend/cuda/device.h"
-
-#include <fmt/format.h>
-
-namespace mlx::core {
-
-CudaStream::CudaStream(cu::Device& device) {
-  device.make_current();
-  CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
-}
-
-CudaStream::~CudaStream() {
-  CHECK_CUDA_ERROR(cudaStreamDestroy(stream_));
-}
-
-void check_cuda_error(const char* name, cudaError_t err) {
-  if (err != cudaSuccess) {
-    throw std::runtime_error(
-        fmt::format("{} failed: {}", name, cudaGetErrorString(err)));
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/utils.h

@@ -1,36 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuda_runtime.h>
-
-namespace mlx::core {
-
-namespace cu {
-class Device;
-}
-
-// Cuda stream managed with RAII.
-class CudaStream {
- public:
-  explicit CudaStream(cu::Device& device);
-  ~CudaStream();
-
-  CudaStream(const CudaStream&) = delete;
-  CudaStream& operator=(const CudaStream&) = delete;
-
-  operator cudaStream_t() const {
-    return stream_;
-  }
-
- private:
-  cudaStream_t stream_;
-};
-
-// Throw exception if the cuda API does not succeed.
-void check_cuda_error(const char* name, cudaError_t err);
-
-// The macro version that prints the command that failed.
-#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
-
-} // namespace mlx::core

mlx/backend/cuda/worker.cpp

@@ -1,90 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
-
-namespace mlx::core::cu {
-
-Worker::Worker()
-    : signal_stream_(device(mlx::core::Device::gpu)),
-      worker_(&Worker::thread_fn, this) {}
-
-Worker::~Worker() {
-  {
-    std::lock_guard lock(worker_mutex_);
-    stop_ = true;
-  }
-  worker_event_.signal(batch_ + 1);
-  worker_.join();
-}
-
-void Worker::add_task(std::function<void()> task) {
-  pending_tasks_.push_back(std::move(task));
-}
-
-void Worker::consume_in_this_thread() {
-  for (auto& task : pending_tasks_) {
-    task();
-  }
-  pending_tasks_.clear();
-}
-
-void Worker::end_batch() {
-  batch_++;
-  {
-    std::lock_guard lock(worker_mutex_);
-    worker_tasks_[batch_] = std::move(pending_tasks_);
-  }
-  uncommited_batches_++;
-}
-
-void Worker::commit() {
-  if (uncommited_batches_ == 0) {
-    return;
-  }
-  uncommited_batches_ = 0;
-  worker_event_.signal(batch_);
-}
-
-void Worker::commit(cudaStream_t stream) {
-  if (uncommited_batches_ == 0) {
-    return;
-  }
-  uncommited_batches_ = 0;
-  // Signal the |worker_event_| in |signal_stream_| after the kernels in
-  // |stream_| finish running.
-  signal_event_.record(stream);
-  signal_event_.wait(signal_stream_);
-  worker_event_.signal(signal_stream_, batch_);
-}
-
-void Worker::thread_fn() {
-  // The worker thread is safe to free buffers.
-  allocator().register_this_thread();
-
-  while (!stop_) {
-    uint64_t batch = worker_event_.value();
-    Tasks tasks;
-    {
-      std::lock_guard lock(worker_mutex_);
-      // Move tasks in signaled batches.
-      auto end = worker_tasks_.upper_bound(batch);
-      for (auto it = worker_tasks_.begin(); it != end; ++it) {
-        if (tasks.empty()) {
-          tasks = std::move(it->second);
-        } else {
-          std::move(
-              it->second.begin(), it->second.end(), std::back_inserter(tasks));
-        }
-      }
-      worker_tasks_.erase(worker_tasks_.begin(), end);
-    }
-    for (auto& task : tasks) {
-      task();
-    }
-    worker_event_.wait(batch + 1);
-  }
-}
-
-} // namespace mlx::core::cu

mlx/backend/cuda/worker.h

@@ -1,68 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/backend/cuda/event.h"
-#include "mlx/backend/cuda/utils.h"
-
-#include <functional>
-#include <map>
-#include <mutex>
-#include <thread>
-
-namespace mlx::core::cu {
-
-// Run tasks in worker thread, synchronized with cuda stream.
-class Worker {
- public:
-  Worker();
-  ~Worker();
-
-  Worker(const Worker&) = delete;
-  Worker& operator=(const Worker&) = delete;
-
-  // Add a pending |task| that will run when consumed or commited.
-  void add_task(std::function<void()> task);
-
-  // Run pending tasks immediately in current thread.
-  void consume_in_this_thread();
-
-  // Put pending tasks in a batch.
-  void end_batch();
-
-  // Inform worker thread to run current batches now.
-  void commit();
-
-  // Inform worker thread to run current batches after kernels in |stream|
-  // finish running.
-  void commit(cudaStream_t stream);
-
-  // Return how many batches have been added but not committed yet.
-  size_t uncommited_batches() const {
-    return uncommited_batches_;
-  }
-
- private:
-  void thread_fn();
-
-  uint64_t batch_{0};
-  size_t uncommited_batches_{0};
-
-  // Cuda stream and event for signaling kernel completion.
-  CudaStream signal_stream_;
-  CudaEvent signal_event_;
-
-  // Worker thread.
-  SharedEvent worker_event_;
-  std::thread worker_;
-  std::mutex worker_mutex_;
-  bool stop_{false};
-
-  // Tasks are put in |pending_tasks_| first, and then moved to
-  // |worker_tasks_| when end_batch() is called.
-  using Tasks = std::vector<std::function<void()>>;
-  Tasks pending_tasks_;
-  std::map<uint64_t, Tasks> worker_tasks_;
-};
-
-} // namespace mlx::core::cu

mlx/backend/gpu/CMakeLists.txt

@@ -1,5 +0,0 @@
-target_sources(
-  mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp)

mlx/backend/gpu/available.h

@@ -1,9 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-namespace mlx::core::gpu {
-
-bool is_available();
-
-} // namespace mlx::core::gpu

mlx/backend/gpu/copy.cpp

@@ -1,49 +0,0 @@
-// Copyright © 2023-2024 Apple Inc.
-
-#include "mlx/backend/gpu/copy.h"
-#include "mlx/primitives.h"
-
-#include <cassert>
-
-namespace mlx::core {
-
-void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
-  bool donated = set_copy_output_data(in, out, ctype);
-  if (donated && in.dtype() == out.dtype()) {
-    // If the output has the same type as the input then there is nothing to
-    // copy, just use the buffer.
-    return;
-  }
-  if (ctype == CopyType::GeneralGeneral) {
-    ctype = CopyType::General;
-  }
-  copy_gpu_inplace(in, out, ctype, s);
-}
-
-void copy_gpu(const array& in, array& out, CopyType ctype) {
-  copy_gpu(in, out, ctype, out.primitive().stream());
-}
-
-void copy_gpu_inplace(
-    const array& in,
-    array& out,
-    CopyType ctype,
-    const Stream& s) {
-  assert(in.shape() == out.shape());
-  return copy_gpu_inplace(
-      in, out, in.shape(), in.strides(), out.strides(), 0, 0, ctype, s);
-}
-
-void copy_gpu_inplace(
-    const array& in,
-    array& out,
-    const Strides& i_strides,
-    int64_t i_offset,
-    CopyType ctype,
-    const Stream& s) {
-  assert(in.shape() == out.shape());
-  return copy_gpu_inplace(
-      in, out, in.shape(), i_strides, out.strides(), i_offset, 0, ctype, s);
-}
-
-} // namespace mlx::core

mlx/backend/gpu/primitives.cpp

@@ -1,222 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/primitives.h"
-#include "mlx/backend/common/utils.h"
-#include "mlx/backend/gpu/copy.h"
-#include "mlx/backend/gpu/slicing.h"
-
-#include <cassert>
-
-#define MLX_PROFILER_RANGE(message)
-
-namespace mlx::core {
-
-namespace {
-
-void reshape(const array& in, array& out, Stream s) {
-  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
-  if (copy_necessary) {
-    out.set_data(allocator::malloc(out.nbytes()));
-    copy_gpu_inplace(
-        in,
-        out,
-        in.shape(),
-        in.strides(),
-        make_contiguous_strides(in.shape()),
-        0,
-        0,
-        CopyType::General,
-        s);
-  } else {
-    shared_buffer_reshape(in, out_strides, out);
-  }
-}
-
-} // namespace
-
-void AsStrided::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("AsStrided::eval_gpu");
-  eval(inputs, out);
-}
-
-void AsType::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("AsType::eval_gpu");
-  CopyType ctype =
-      inputs[0].flags().contiguous ? CopyType::Vector : CopyType::General;
-  copy_gpu(inputs[0], out, ctype);
-}
-
-void Broadcast::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Broadcast::eval_gpu");
-  eval(inputs, out);
-}
-
-void BroadcastAxes::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("BroadcastAxes::eval_gpu");
-  eval(inputs, out);
-}
-
-void Concatenate::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Concatenate::eval_gpu");
-  concatenate_gpu(inputs, out, axis_, stream());
-}
-
-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;
-  if (in.buffer_size() <= out.nbytes() + extra_bytes &&
-      (in.flags().row_contiguous ||
-       (allow_col_major_ && in.flags().col_contiguous))) {
-    out.copy_shared_buffer(in);
-  } else {
-    out.set_data(allocator::malloc(out.nbytes()));
-    copy_gpu_inplace(
-        in,
-        out,
-        in.flags().row_contiguous ? CopyType::Vector : CopyType::General,
-        stream());
-  }
-}
-
-void Copy::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Copy::eval_gpu");
-  eval(inputs, out);
-}
-
-void CustomTransforms::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  MLX_PROFILER_RANGE("CustomTransforms::eval_gpu");
-  eval(inputs, outputs);
-}
-
-void Depends::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  MLX_PROFILER_RANGE("Depends::eval_gpu");
-  eval(inputs, outputs);
-}
-
-void ExpandDims::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("ExpandDims::eval_gpu");
-  eval(inputs, out);
-}
-
-void Full::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Full::eval_gpu");
-  auto in = inputs[0];
-  CopyType ctype;
-  if (in.data_size() == 1) {
-    ctype = CopyType::Scalar;
-  } else if (in.flags().contiguous) {
-    ctype = CopyType::Vector;
-  } else {
-    ctype = CopyType::General;
-  }
-  copy_gpu(in, out, ctype);
-}
-
-void Flatten::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Flatten::eval_gpu");
-  reshape(inputs[0], out, stream());
-}
-
-void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("NumberOfElements::eval_gpu");
-  eval(inputs, out);
-}
-
-void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
-  // Inputs must be base input array and scalar val array
-  assert(inputs.size() == 2);
-  auto& in = inputs[0];
-  auto& val = inputs[1];
-
-  // Padding value must be a scalar
-  assert(val.size() == 1);
-
-  // Padding value, input and output must be of the same type
-  assert(val.dtype() == in.dtype() && in.dtype() == out.dtype());
-
-  pad_gpu(in, val, out, axes_, low_pad_size_, stream());
-}
-
-void Reshape::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Reshape::eval_gpu");
-  reshape(inputs[0], out, stream());
-}
-
-void Split::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  MLX_PROFILER_RANGE("Split::eval_gpu");
-  eval(inputs, outputs);
-}
-
-void Slice::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Slice::eval_gpu");
-  assert(inputs.size() == 1);
-  if (out.size() == 0) {
-    out.set_data(nullptr);
-    return;
-  }
-
-  auto& in = inputs[0];
-  slice_gpu(in, out, start_indices_, strides_, stream());
-}
-
-void Squeeze::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Squeeze::eval_gpu");
-  eval(inputs, out);
-}
-
-void StopGradient::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("StopGradient::eval_gpu");
-  eval(inputs, out);
-}
-
-void Transpose::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Transpose::eval_gpu");
-  eval(inputs, out);
-}
-
-void Unflatten::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("Unflatten::eval_gpu");
-  reshape(inputs[0], out, stream());
-}
-
-void View::eval_gpu(const std::vector<array>& inputs, array& out) {
-  MLX_PROFILER_RANGE("View::eval_gpu");
-  auto& in = inputs[0];
-  auto ibytes = size_of(in.dtype());
-  auto obytes = size_of(out.dtype());
-  // Conditions for buffer copying (disjunction):
-  // - type size is the same
-  // - type size is smaller and the last axis is contiguous
-  // - the entire array is row contiguous
-  if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) ||
-      in.flags().row_contiguous) {
-    auto strides = in.strides();
-    for (int i = 0; i < static_cast<int>(strides.size()) - 1; ++i) {
-      strides[i] *= ibytes;
-      strides[i] /= obytes;
-    }
-    out.copy_shared_buffer(
-        in, strides, in.flags(), in.data_size() * ibytes / obytes);
-  } else {
-    auto tmp = array(in.shape(), in.dtype(), nullptr, {});
-    tmp.set_data(allocator::malloc(tmp.nbytes()));
-    copy_gpu_inplace(in, tmp, CopyType::General, stream());
-
-    auto flags = out.flags();
-    flags.contiguous = true;
-    flags.row_contiguous = true;
-    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
-    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
-    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
-  }
-}
-
-} // namespace mlx::core

mlx/backend/gpu/slicing.cpp

@@ -1,44 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/common/slicing.h"
-#include "mlx/backend/gpu/copy.h"
-#include "mlx/backend/gpu/slicing.h"
-
-namespace mlx::core {
-
-void slice_gpu(
-    const array& in,
-    array& out,
-    const Shape& start_indices,
-    const Shape& strides,
-    const Stream& s) {
-  slice(in, out, start_indices, strides);
-}
-
-void pad_gpu(
-    const array& in,
-    const array& val,
-    array& out,
-    const std::vector<int>& axes,
-    const Shape& low_pad_size,
-    const Stream& s) {
-  // Fill output with val
-  fill_gpu(val, out, s);
-
-  // Find offset for start of input values
-  size_t data_offset = 0;
-  for (int i = 0; i < axes.size(); i++) {
-    auto ax = axes[i] < 0 ? out.ndim() + axes[i] : axes[i];
-    data_offset += out.strides()[ax] * low_pad_size[i];
-  }
-
-  // Extract slice from output where input will be pasted
-  array out_slice(in.shape(), out.dtype(), nullptr, {});
-  out_slice.copy_shared_buffer(
-      out, out.strides(), out.flags(), out_slice.size(), data_offset);
-
-  // Copy input values into the slice
-  copy_gpu_inplace(in, out_slice, CopyType::GeneralGeneral, s);
-}
-
-} // namespace mlx::core

mlx/backend/metal/CMakeLists.txt

@@ -47,7 +47,6 @@ if(MLX_METAL_JIT)
   make_jit_source(binary)
   make_jit_source(binary_two)
   make_jit_source(fft kernels/fft/radix.h kernels/fft/readwrite.h)
-  make_jit_source(logsumexp)
   make_jit_source(ternary)
   make_jit_source(softmax)
   make_jit_source(scan)
@@ -61,7 +60,6 @@ if(MLX_METAL_JIT)
     kernels/steel/gemm/transforms.h)
   make_jit_source(steel/gemm/kernels/steel_gemm_fused)
   make_jit_source(steel/gemm/kernels/steel_gemm_masked kernels/steel/defines.h)
-  make_jit_source(steel/gemm/kernels/steel_gemm_gather)
   make_jit_source(steel/gemm/kernels/steel_gemm_splitk)
   make_jit_source(
     steel/conv/conv
@@ -93,12 +91,10 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp

mlx/backend/metal/allocator.cpp

@@ -1,8 +1,8 @@
 // Copyright © 2023-2024 Apple Inc.
 #include "mlx/backend/metal/allocator.h"
 #include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/metal_impl.h"
 #include "mlx/backend/metal/resident.h"
-#include "mlx/memory.h"
 
 #include <mach/vm_page_size.h>
 #include <unistd.h>
@@ -20,9 +20,6 @@ Allocator& allocator() {
 }
 
 void* Buffer::raw_ptr() {
-  if (!ptr_) {
-    return nullptr;
-  }
   return static_cast<MTL::Buffer*>(ptr_)->contents();
 }
 
@@ -32,11 +29,8 @@ namespace metal {
 
 namespace {
 
-BufferCache::BufferCache(ResidencySet& residency_set)
-    : head_(nullptr),
-      tail_(nullptr),
-      pool_size_(0),
-      residency_set_(residency_set) {}
+BufferCache::BufferCache(MTL::Device* device)
+    : device_(device), head_(nullptr), tail_(nullptr), pool_size_(0) {}
 
 BufferCache::~BufferCache() {
   auto pool = metal::new_scoped_memory_pool();
@@ -47,9 +41,6 @@ int BufferCache::clear() {
   int n_release = 0;
   for (auto& [size, holder] : buffer_pool_) {
     if (holder->buf) {
-      if (!holder->buf->heap()) {
-        residency_set_.erase(holder->buf);
-      }
       holder->buf->release();
       n_release++;
     }
@@ -107,9 +98,6 @@ int BufferCache::release_cached_buffers(size_t min_bytes_to_free) {
     while (tail_ && (total_bytes_freed < min_bytes_to_free)) {
       if (tail_->buf) {
         total_bytes_freed += tail_->buf->length();
-        if (!tail_->buf->heap()) {
-          residency_set_.erase(tail_->buf);
-        }
         tail_->buf->release();
         tail_->buf = nullptr;
         n_release++;
@@ -164,7 +152,7 @@ void BufferCache::remove_from_list(BufferCache::BufferHolder* to_remove) {
 MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
       residency_set_(device_),
-      buffer_cache_(residency_set_) {
+      buffer_cache_(device_) {
   auto pool = metal::new_scoped_memory_pool();
   auto memsize = std::get<size_t>(device_info().at("memory_size"));
   auto max_rec_size =
@@ -201,19 +189,16 @@ size_t MetalAllocator::set_cache_limit(size_t limit) {
   return limit;
 };
 
-size_t MetalAllocator::set_memory_limit(size_t limit) {
+size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) {
   std::unique_lock lk(mutex_);
   std::swap(limit, block_limit_);
+  relaxed_ = relaxed;
   gc_limit_ = std::min(
       block_limit_,
       static_cast<size_t>(0.95 * device_->recommendedMaxWorkingSetSize()));
   return limit;
 };
 
-size_t MetalAllocator::get_memory_limit() {
-  return block_limit_;
-}
-
 size_t MetalAllocator::set_wired_limit(size_t limit) {
   std::unique_lock lk(mutex_);
   std::swap(limit, wired_limit_);
@@ -221,7 +206,7 @@ size_t MetalAllocator::set_wired_limit(size_t limit) {
   return limit;
 };
 
-Buffer MetalAllocator::malloc(size_t size) {
+Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   // Metal doesn't like empty buffers
   if (size == 0) {
     return Buffer{nullptr};
@@ -248,6 +233,11 @@ Buffer MetalAllocator::malloc(size_t size) {
   if (!buf) {
     size_t mem_required = get_active_memory() + get_cache_memory() + size;
 
+    // If there is too much memory pressure, fail (likely causes a wait).
+    if (!(allow_swap && relaxed_) && mem_required >= block_limit_) {
+      return Buffer{nullptr};
+    }
+
     auto pool = metal::new_scoped_memory_pool();
 
     // If we have a lot of memory pressure or are over the maximum cache size,
@@ -271,13 +261,9 @@ Buffer MetalAllocator::malloc(size_t size) {
     if (!buf) {
       buf = device_->newBuffer(size, resource_options);
     }
-    if (!buf) {
-      return Buffer{nullptr};
-    }
     lk.lock();
-    num_resources_++;
-    if (!buf->heap()) {
-      residency_set_.insert(buf);
+    if (buf) {
+      num_resources_++;
     }
   }
 
@@ -291,6 +277,10 @@ Buffer MetalAllocator::malloc(size_t size) {
         get_cache_memory() - max_pool_size_);
   }
 
+  if (!buf->heap()) {
+    residency_set_.insert(buf);
+  }
+
   return Buffer{static_cast<void*>(buf)};
 }
 
@@ -306,14 +296,14 @@ void MetalAllocator::free(Buffer buffer) {
     return;
   }
   std::unique_lock lk(mutex_);
+  if (!buf->heap()) {
+    residency_set_.erase(buf);
+  }
   active_memory_ -= buf->length();
   if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
     num_resources_--;
-    if (!buf->heap()) {
-      residency_set_.erase(buf);
-    }
     lk.unlock();
     auto pool = metal::new_scoped_memory_pool();
     buf->release();
@@ -332,40 +322,37 @@ MetalAllocator& allocator() {
   return *allocator_;
 }
 
-} // namespace metal
-
 size_t set_cache_limit(size_t limit) {
-  return metal::allocator().set_cache_limit(limit);
+  return allocator().set_cache_limit(limit);
 }
-size_t set_memory_limit(size_t limit) {
-  return metal::allocator().set_memory_limit(limit);
-}
-size_t get_memory_limit() {
-  return metal::allocator().get_memory_limit();
+size_t set_memory_limit(size_t limit, bool relaxed /* = true */) {
+  return allocator().set_memory_limit(limit, relaxed);
 }
 size_t set_wired_limit(size_t limit) {
-  if (limit > std::get<size_t>(metal::device_info().at(
-                  "max_recommended_working_set_size"))) {
+  if (limit >
+      std::get<size_t>(device_info().at("max_recommended_working_set_size"))) {
     throw std::invalid_argument(
         "[metal::set_wired_limit] Setting a wired limit larger than "
         "the maximum working set size is not allowed.");
   }
-  return metal::allocator().set_wired_limit(limit);
+  return allocator().set_wired_limit(limit);
 }
 size_t get_active_memory() {
-  return metal::allocator().get_active_memory();
+  return allocator().get_active_memory();
 }
 size_t get_peak_memory() {
-  return metal::allocator().get_peak_memory();
+  return allocator().get_peak_memory();
 }
 void reset_peak_memory() {
-  metal::allocator().reset_peak_memory();
+  allocator().reset_peak_memory();
 }
 size_t get_cache_memory() {
-  return metal::allocator().get_cache_memory();
+  return allocator().get_cache_memory();
 }
 void clear_cache() {
-  return metal::allocator().clear_cache();
+  return allocator().clear_cache();
 }
 
+} // namespace metal
+
 } // namespace mlx::core