format (#2700)

.circleci/config.yml

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

.github/actions/build-cuda-release/action.yml

@@ -1,20 +0,0 @@
-name: 'Build CUDA wheel'
-description: 'Build CUDA wheel'
-
-inputs:
-  nvcc-location:
-    description: 'Location of nvcc compiler'
-    required: true
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build package
-      shell: bash
-      env:
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
-      run: |
-        pip install auditwheel build patchelf setuptools
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        bash python/scripts/repair_cuda.sh

.github/actions/build-cuda/action.yml

@@ -1,45 +0,0 @@
-name: 'Build and Test with CUDA'
-description: 'Build and test MLX with CUDA'
-
-inputs:
-  nvcc-location:
-    description: 'Location of nvcc compiler'
-    required: true
-    default: '/usr/local/cuda-12.9/bin/nvcc'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      shell: bash
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
-      run: pip install -e ".[dev]" -v
-
-    - name: Run Python tests - CPU
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-        DEVICE: cpu
-      run: python -m unittest discover python/tests -v
-
-    - name: Run Python tests - GPU
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-        DEVICE: gpu
-      run: python -m tests discover python/tests -v
-
-    - name: Build CPP only
-      shell: bash
-      run: |
-        cmake . -B build \
-          -DMLX_BUILD_CUDA=ON \
-          -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }} \
-          -DCMAKE_BUILD_TYPE=DEBUG
-        cmake --build build -j $(nproc)
-    
-    - name: Run CPP tests
-      shell: bash
-      run: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"

.github/actions/build-docs/action.yml

@@ -1,38 +0,0 @@
-name: 'Build Documentation'
-description: 'Build documentation on a mac'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Setup machine
-      uses: ./.github/actions/setup-macos
-
-    - name: Install dependencies
-      shell: sh
-      run: |
-        brew install doxygen
-        uv pip install --upgrade pip cmake
-        uv pip install -r docs/requirements.txt
-        uv pip install . -v
-  
-    - name: Build documentation
-      shell: bash
-      run: |
-        source .venv/bin/activate
-        cd docs
-        doxygen
-        make html O=-W
-    
-    - name: Create artifact tar
-      shell: sh
-      run: tar -cf artifact.tar --cd docs --dereference build/html index.html
-
-    # Do it manually because upload-pages-artifact requires gtar
-    - name: Upload artifact
-      id: upload-artifact
-      uses: actions/upload-artifact@v5
-      with:
-        name: github-pages
-        path: artifact.tar
-        retention-days: 1
-        if-no-files-found: error

.github/actions/build-linux-release/action.yml

@@ -1,33 +0,0 @@
-name: 'Build Linux wheel'
-description: 'Build Linux wheel'
-
-inputs:
-  build-backend:
-    description: 'Build the backend mlx-cpu package'
-    type: boolean
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Generate package stubs
-      shell: bash
-      run: |
-        pip install -e ".[dev]" -v
-        pip install typing_extensions
-        python setup.py generate_stubs
-    - name: Build Python package
-      shell: bash
-      run: |
-        pip install auditwheel patchelf build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        bash python/scripts/repair_linux.sh
-    - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash
-      run: |
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64

.github/actions/build-linux/action.yml

@@ -1,41 +0,0 @@
-name: 'Build and Test on Linux'
-description: 'Build and test MLX on Linux'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      shell: sh
-      env:
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-        DEBUG: 1
-      run: pip install -e ".[dev]" -v
-    
-    - name: Generate package stubs
-      shell: sh
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-    
-    - name: Run Python tests
-      shell: bash
-      run: |
-        python -m unittest discover python/tests -v
-        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-        if grep -Fq '[WARN]' stderr.log ; then
-          grep -F '[WARN]' stderr.log
-          echo "Distributed ring test failed";
-          exit 1;
-        fi
-    
-    - name: Build CPP only
-      shell: bash
-      run: |
-        mkdir -p build && cd build
-        cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-        make -j $(nproc)
-    
-    - name: Run CPP tests
-      shell: sh
-      run: ./build/tests/tests

.github/actions/build-macos-release/action.yml

@@ -1,33 +0,0 @@
-name: 'Build macOS release'
-description: 'Build MLX releases macOS'
-
-inputs:
-  macos-target:
-    description: 'macOS build target'
-    required: false
-    default: '15.0'
-  build-backend:
-    description: 'Build the backend mlx-metal package'
-    type: boolean
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build Python package
-      shell: bash
-      env:
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        uv pip install build
-        uv run --no-project setup.py clean --all
-        MLX_BUILD_STAGE=1 uv run -m build -w
-    - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash
-      env:
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        uv run --no-project setup.py clean --all
-        MLX_BUILD_STAGE=2 uv run -m build -w

.github/actions/build-macos/action.yml

@@ -1,88 +0,0 @@
-name: 'Build and Test on macOS'
-description: 'Build and test MLX on macOS'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install dependencies
-      shell: sh
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      run: |
-        uv pip install --upgrade pip
-        uv pip install cmake setuptools nanobind==2.4.0
-        uv pip install -e . -v
-
-    - name: Generate package stubs
-      shell: bash
-      run: |
-        uv pip install typing_extensions
-        uv run --no-project setup.py generate_stubs
-
-    - name: Install tests dependencies
-      shell: sh
-      run: |
-        uv pip install numpy torch tensorflow unittest-xml-reporting
-
-    - name: Run Python tests
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-      run: |
-        DEVICE=cpu uv run -m xmlrunner discover -v python/tests -o test-results/cpu
-        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 uv run -m xmlrunner discover -v python/tests -o test-results/gpu
-        mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
-        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-        if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-    
-    - name: Build example extension
-      shell: bash
-      run: |
-        cd examples/extensions
-        uv pip install -r requirements.txt
-        uv run --no-project setup.py build_ext --inplace
-        uv run --no-project test.py
-    
-    - name: Build CPP only
-      shell: bash
-      run: |
-        mkdir -p build
-        cd build
-        cmake ..
-        make -j $(sysctl -n hw.ncpu)
-    
-    - name: Run CPP tests
-      shell: bash
-      env:
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
-    - name: Build small binary with JIT
-      shell: bash
-      run: |
-        mkdir -p build
-        cd build
-        cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
-          -DBUILD_SHARED_LIBS=ON \
-          -DMLX_BUILD_CPU=OFF \
-          -DMLX_BUILD_SAFETENSORS=OFF \
-          -DMLX_BUILD_GGUF=OFF \
-          -DMLX_METAL_JIT=ON
-        make -j $(sysctl -n hw.ncpu)
-    
-    - name: Run Python tests with JIT
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: |
-        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          uv pip install -e . -v
-        uv run -m xmlrunner discover \
-            -v python/tests \
-            -o test-results/gpu_jit

.github/actions/setup-linux/action.yml

@@ -1,83 +0,0 @@
-name: 'Setup Linux Environment'
-description: 'Install dependencies for Linux builds'
-
-inputs:
-  runner-type:
-    description: 'Whether to set this up as a linux or CUDA runner'
-    required: false
-    default: 'linux'
-    type: choice
-    options:
-      - linux
-      - cuda
-  python-version:
-    description: 'Version of python to set up'
-    required: false
-    default: '3.10'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Free disk space
-      shell: sh
-      if: inputs.runner-type == 'linux'
-      run: sudo rm -rf "$AGENT_TOOLSDIRECTORY"
-
-    - name: Install common dependencies
-      env:
-        TZ: Etc/UTC
-      shell: bash
-      run: |
-        sudo apt-get update
-        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev tzdata zip
-        sudo apt autoremove -y
-
-    - uses: actions/setup-python@v6
-      with:
-        python-version: ${{ inputs.python-version }}
-        cache: 'pip'
-
-    - name: setup python venv
-      shell: bash
-      run: |
-        python -m venv .venv
-        source .venv/bin/activate
-        echo PATH=$PATH >> $GITHUB_ENV
-        pip install --upgrade pip cmake
-
-    - name: Install MPI
-      if: inputs.runner-type == 'linux'
-      shell: bash
-      run: sudo apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
-
-    - name: Network CUDA installation from packages
-      id: install-cuda
-      if: inputs.runner-type == 'cuda'
-      env:
-        TZ: Etc/UTC
-      shell: bash ## Specific to Ubuntu 22.04 & Architecture x86_64
-      run: |
-        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
-        sudo dpkg -i cuda-keyring_1.1-1_all.deb
-        sudo apt-get update
-        sudo apt-get install -y libcudnn9-dev-cuda-12 libnccl2 libnccl-dev cuda-toolkit-12-9
-      # Note: This installs CUDA 12.9, which is the latest supported by cuDNN 9.x and works with the NVidia 570 drivers
-      # cuda-toolkit by itself installs version 13 (+) and requires updated drives (580+), which require a reboot to function properly.
-      # Compatibility matrix: https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
-      # This also drops `nvcc` into `/usr/local/cuda-12.9/bin/nvcc` - but it's *not* on the default PATH
-
-    - name: Package and Driver Report
-      if: inputs.runner-type == 'cuda'
-      shell: bash
-      run: |
-        sudo apt-get install -y ubuntu-drivers-common dkms
-        echo "NVIDIA Driver Packages Available:"
-        sudo ubuntu-drivers list --gpgpu
-        echo "NVIDIA Driver Version:"
-        cat /proc/driver/nvidia/version || echo "nvidia driver not found"
-        echo "Installed NVIDIA and CUDA packages:"
-        dpkg -l | egrep "cuda|nvidia" -i
-        echo "DKMS Status:"
-        dkms status || echo "dkms not found"
-        echo "NVIDIA-SMI Status:"
-        nvidia-smi || echo "nvidia-smi not found"

.github/actions/setup-macos/action.yml

@@ -1,25 +0,0 @@
-name: 'Setup macOS Environment'
-description: 'Install dependencies for macOS builds'
-
-inputs:
-  python-version:
-    description: 'Python version to use'
-    required: false
-    default: '3.10'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Homebrew packages
-      shell: sh
-      run: /opt/homebrew/bin/brew install openmpi
-    
-    - name: Verify MetalToolchain installed
-      shell: bash
-      run: xcodebuild -showComponent MetalToolchain
-    
-    - name: Setup uv
-      uses: astral-sh/setup-uv@v6
-      with:
-          python-version: ${{ inputs.python-version }}
-          activate-environment: true

.github/dependabot.yml

@@ -1,6 +0,0 @@
-version: 2
-updates:
-  - package-ecosystem: "github-actions"
-    directory: "/"
-    schedule:
-      interval: "weekly"

.github/scripts/setup+build-cpp-linux-fedora-container.sh

@@ -1,27 +0,0 @@
-#!/bin/bash
-set -ex
-
-# [Setup] Install dependencies inside the container.
-dnf update -y
-dnf install -y \
-  blas-devel \
-  lapack-devel \
-  openblas-devel \
-  make \
-  cmake \
-  clang \
-  git
-dnf clean all
-
-# [C++] CI Build Sanity Check: Verifies code compilation, not for release.
-export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-export DEBUG=1
-export CMAKE_C_COMPILER=/usr/bin/clang
-export CMAKE_CXX_COMPILER=/usr/bin/clang++
-
-mkdir -p build
-pushd build
-cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-make -j $(nproc)
-./tests/tests
-popd

.github/workflows/documentation.yml

@@ -1,28 +0,0 @@
-name: Documentation
-
-on:
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build:
-    runs-on: [self-hosted, macos]
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/build-docs
-      
-  deploy:
-    needs: build
-    permissions:
-      pages: write
-      id-token: write
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4

.github/workflows/nightly.yml

@@ -1,124 +0,0 @@
-name: Nightly Build
-
-on:
-  schedule:
-    - cron: 33 6 * * 1-5
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build_linux_release:
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.10", "3.14"]
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux-release
-        with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: linux-wheels-${{ matrix.python_version }}
-          path: wheelhouse/mlx-*.whl
-          retention-days: 7
-      - name: Upload mlx-cpu artifacts
-        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-cpu
-          path: wheelhouse/mlx_cpu-*.whl
-          retention-days: 7
-  
-  build_linux_with_tests:
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux
-
-  build_mac_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python-version: ["3.10", "3.13"]
-    runs-on: [self-hosted, macos]
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-macos
-        with:
-          python-version: ${{ matrix.python-version }}
-      - uses: ./.github/actions/build-macos
-
-      - name: Build macOS 15 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 15.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 14 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 14.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-
-  build_cuda_with_tests:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: gpu-t4-4-core
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          runner-type: 'cuda'
-      - uses: ./.github/actions/build-cuda
-
-  build_cuda_release:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          runner-type: 'cuda'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-        with:
-          nvcc-location: '/usr/local/cuda-12.9/bin/nvcc'
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
-          retention-days: 7
-
-  linux_fedora_build_cpp:
-    name: Linux Fedora CPP Build (${{ matrix.arch }})
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - host: ubuntu-22.04
-            arch: x86_64
-          - host: ubuntu-22.04-arm
-            arch: aarch64
-
-    runs-on: ${{ matrix.host }}
-    container:
-      image: fedora:42
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v5
-
-      - name: CPP Build Test - No Release
-        run: |
-          bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh

.github/workflows/pull_request.yml

@@ -1,71 +1,20 @@
-name: Build and Test
-
-on: pull_request  
-
-permissions:
-  contents: read
+on:
+  pull_request:
+    branches:
+      - main
 
 jobs:
   check_lint:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-      - uses: pre-commit/action@v3.0.1
-
-  linux_build_and_test:
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
-
-  mac_build_and_test:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: [self-hosted, macos]
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-macos
-      - uses: ./.github/actions/build-macos
-
-  cuda_build_and_test:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: gpu-t4-4-core
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
         with:
-          runner-type: 'cuda'
-      - uses: ./.github/actions/build-cuda
-
-  build_documentation:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: [self-hosted, macos]
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/build-docs
-
-  linux_fedora_build_cpp:
-    name: Linux Fedora CPP Build (${{ matrix.arch }})
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - host: ubuntu-22.04
-            arch: x86_64
-          - host: ubuntu-22.04-arm
-            arch: aarch64
-
-    runs-on: ${{ matrix.host }}
-    container:
-      image: fedora:42
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v5
-
-      - name: CPP Build Test - No Release
+          python-version: 3.8
+      - name: Install dependencies
         run: |
-          bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh
+          python -m pip install --upgrade pip
+          pip install pre-commit black isort clang-format
+      - name: Run lint
+        run: |
+          pre-commit run --all-files

.github/workflows/release.yml

@@ -1,226 +0,0 @@
-name: PyPI Release
-
-on:
-  push:
-    tags:
-      - 'v*'
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  setup:
-    runs-on: ubuntu-latest
-    outputs:
-      pypi_env: ${{ github.event_name == 'push' && 'pypi' || 'test-pypi' }}
-      pypi_url: ${{ github.event_name == 'push' && 'https://upload.pypi.org/legacy/' || 'https://test.pypi.org/legacy/' }}
-    steps:
-      - name: Set publishing variables
-        run: echo "Publishing setup complete"
-
-  build_documentation:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: [self-hosted, macos]
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/build-docs
-    
-  deploy_documentation:
-    needs: build_documentation
-    permissions:
-      pages: write
-      id-token: write
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4
-
-  build_linux_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
-    env:
-      PYPI_RELEASE: 1
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux-release
-        with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: linux-wheels-${{ matrix.python_version }}
-          path: wheelhouse/mlx-*.whl
-      - name: Upload CPU artifacts
-        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-cpu
-          path: wheelhouse/mlx_cpu-*.whl
-  
-  build_mac_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: [self-hosted, macos]
-    env:
-      PYPI_RELEASE: 1
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-macos
-        with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        shell: sh
-        run: |
-          uv pip install --upgrade pip
-          uv pip install cmake setuptools nanobind==2.4.0
-          uv pip install -e . -v
-      - name: Generate package stubs
-        shell: bash
-        run: |
-          uv pip install typing_extensions
-          uv run --no-project setup.py generate_stubs
-      - name: Build macOS 14 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 14.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 15 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 15.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: mac-wheels-${{ matrix.python-version }}
-          path: dist/mlx-*.whl
-      - name: Upload Metal artifacts
-        if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-metal
-          path: dist/mlx_metal-*.whl
-
-  build_cuda_release:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
-    env:
-      PYPI_RELEASE: 1
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          runner-type: 'cuda'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-        with:
-          nvcc-location: '/usr/local/cuda-12.9/bin/nvcc'
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
-
-  pypi-publish:
-    name: Upload release to PyPI
-    runs-on: ubuntu-latest
-    needs: [setup, build_linux_release, build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
-      url: https://pypi.org/p/mlx
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          pattern: linux-wheels-*
-          merge-multiple: true
-          path: dist
-      - uses: actions/download-artifact@v6
-        with:
-          pattern: mac-wheels-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
-  
-  pypi-publish-cuda:
-    name: Upload CUDA release to PyPI
-    runs-on: ubuntu-latest
-    needs: [setup, build_cuda_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
-      url: https://pypi.org/p/mlx-cuda
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          name: mlx-cuda
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
-
-  pypi-publish-cpu:
-    name: Upload CPU release to PyPI
-    runs-on: ubuntu-latest
-    needs: [setup, build_linux_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
-      url: https://pypi.org/p/mlx-cpu
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          name: mlx-cpu
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
-
-  pypi-publish-metal:
-    name: Upload Metal release to PyPI
-    runs-on: ubuntu-latest
-    needs: [setup, build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
-      url: https://pypi.org/p/mlx-metal
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          name: mlx-metal
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
-

.pre-commit-config.yaml

@@ -1,10 +1,4 @@
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v6.0.0
-    hooks:
-    -   id: check-yaml
-    # -   id: end-of-file-fixer
-    # -   id: trailing-whitespace
 -   repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v19.1.7
     hooks:

CMakeLists.txt

@@ -127,12 +127,9 @@ if(MLX_BUILD_METAL)
   message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}")
 
   set(METAL_CPP_URL
-      https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18.zip)
 
   if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
-    if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0)
-      message(FATAL_ERROR "MLX requires macOS >= 14.0")
-    endif()
     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
   endif()
   execute_process(
@@ -141,6 +138,7 @@ if(MLX_BUILD_METAL)
       "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
     OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+
   FetchContent_MakeAvailable(metal_cpp)
   target_include_directories(
     mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>

docs/src/install.rst

@@ -16,11 +16,12 @@ silicon computer is
 To install from PyPI your system must meet the following requirements:
 
 - Using an M series chip (Apple silicon)
-- Using a native Python >= 3.10
-- macOS >= 14.0
+- Using a native Python >= 3.9
+- macOS >= 13.5
 
 .. note::
-    MLX is only available on devices running macOS >= 14.0 and higher.
+    MLX is only available on devices running macOS >= 13.5
+    It is highly recommended to use macOS 14 (Sonoma)
 
 CUDA
 ^^^^
@@ -38,7 +39,7 @@ requirements:
 - Nvidia driver >= 550.54.14
 - CUDA toolkit >= 12.0
 - Linux distribution with glibc >= 2.35
-- Python >= 3.10
+- Python >= 3.9
 
 
 CPU-only (Linux)
@@ -54,7 +55,7 @@ To install the CPU-only package from PyPi your system must meet the following
 requirements:
 
 - Linux distribution with glibc >= 2.35
-- Python >= 3.10
+- Python >= 3.9
 
 
 Troubleshooting

docs/src/python/ops.rst

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

docs/src/usage/distributed.rst

@@ -7,13 +7,12 @@ Distributed Communication
 
 MLX supports distributed communication operations that allow the computational cost
 of training or inference to be shared across many physical machines. At the
-moment we support three different communication backends:
+moment we support two different communication backends:
 
 * `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
   full-featured and mature distributed communications library
-* A **ring** backend of our own that uses native TCP sockets. It should be
-  faster for thunderbolt connections, but it also works over Ethernet.
-* `nccl <https://developer.nvidia.com/nccl>`_, for use in CUDA environments.
+* A **ring** backend of our own that uses native TCP sockets and should be
+  faster for thunderbolt connections.
 
 The list of all currently supported operations and their documentation can be
 seen in the :ref:`API docs<distributed>`.
@@ -85,8 +84,9 @@ Selecting Backend
 ^^^^^^^^^^^^^^^^^
 
 You can select the backend you want to use when calling :func:`init` by passing
-one of ``{'any', 'ring', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all
-available backends. If they all fail then a singleton group is created.
+one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to
+initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they
+both fail then a singleton group is created.
 
 .. note::
    After a distributed backend is successfully initialized :func:`init` will
@@ -220,7 +220,7 @@ print 4 etc.
 Installing MPI
 ^^^^^^^^^^^^^^
 
-MPI can be installed with Homebrew, pip, using the Anaconda package manager, or
+MPI can be installed with Homebrew, using the Anaconda package manager or
 compiled from source. Most of our testing is done using ``openmpi`` installed
 with the Anaconda package manager as follows:
 
@@ -228,16 +228,14 @@ with the Anaconda package manager as follows:
 
     $ conda install conda-forge::openmpi
 
-Installing with Homebrew or pip requires specifying the location of ``libmpi.dyld``
+Installing with Homebrew may require specifying the location of ``libmpi.dyld``
 so that MLX can find it and load it at runtime. This can simply be achieved by
 passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
-done automatically by ``mlx.launch``. Some environments use a non-standard
-library filename that can be specified using the ``MPI_LIBNAME`` environment
-variable. This is automatically taken care of by ``mlx.launch`` as well.
+done automatically by ``mlx.launch``.
 
 .. code:: shell
 
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
     $ # or simply
     $ mlx.launch -n 2 test.py
 

mlx/allocator.h

@@ -14,7 +14,7 @@ class Buffer {
   void* ptr_;
 
  public:
-  explicit Buffer(void* ptr) : ptr_(ptr) {};
+  Buffer(void* ptr) : ptr_(ptr) {};
 
   // Get the raw data pointer from the buffer
   void* raw_ptr();

mlx/array.cpp

@@ -64,7 +64,7 @@ array array::unsafe_weak_copy(const array& other) {
       other.strides(),
       other.flags(),
       [](auto) {});
-  cpy.array_desc_->offset = other.array_desc_->offset;
+  cpy.array_desc_->data_ptr = other.array_desc_->data_ptr;
   return cpy;
 }
 
@@ -141,7 +141,7 @@ bool array::is_tracer() const {
 
 void array::set_data(allocator::Buffer buffer, Deleter d) {
   array_desc_->data = std::make_shared<Data>(buffer, d);
-  array_desc_->offset = 0;
+  array_desc_->data_ptr = buffer.raw_ptr();
   array_desc_->data_size = size();
   array_desc_->flags.contiguous = true;
   array_desc_->flags.row_contiguous = true;
@@ -156,7 +156,7 @@ void array::set_data(
     Flags flags,
     Deleter d) {
   array_desc_->data = std::make_shared<Data>(buffer, d);
-  array_desc_->offset = 0;
+  array_desc_->data_ptr = buffer.raw_ptr();
   array_desc_->data_size = data_size;
   array_desc_->strides = std::move(strides);
   array_desc_->flags = flags;
@@ -172,8 +172,9 @@ void array::copy_shared_buffer(
   array_desc_->strides = strides;
   array_desc_->flags = flags;
   array_desc_->data_size = data_size;
-  array_desc_->offset =
-      sizeof(char) * itemsize() * offset + other.array_desc_->offset;
+  auto char_offset = sizeof(char) * itemsize() * offset;
+  array_desc_->data_ptr = static_cast<void*>(
+      static_cast<char*>(other.array_desc_->data_ptr) + char_offset);
 }
 
 void array::copy_shared_buffer(const array& other) {
@@ -240,8 +241,8 @@ array::ArrayDesc::ArrayDesc(
     std::vector<array> inputs)
     : shape(std::move(shape)),
       dtype(dtype),
-      primitive(std::move(primitive)),
       status(Status::unscheduled),
+      primitive(std::move(primitive)),
       inputs(std::move(inputs)) {
   init();
 }

mlx/array.h

@@ -294,11 +294,6 @@ class array {
     return array_desc_->siblings;
   }
 
-  /** The array's position in the sibling list. */
-  int sibling_position() const {
-    return array_desc_->position;
-  }
-
   void set_siblings(std::vector<array> siblings, uint16_t position) {
     array_desc_->siblings = std::move(siblings);
     array_desc_->position = position;
@@ -354,23 +349,15 @@ class array {
     return array_desc_->data;
   }
 
-  // Return a raw pointer to the arrays data. This function may do a copy if
-  // the underlying buffer is not accessible on the CPU. When accessing the
-  // data for GPU kernels, be sure to use the correct method / function for the
-  // given backend to access the GPU pointer.
+  // Return a raw pointer to the arrays data
   template <typename T>
   T* data() {
-    return reinterpret_cast<T*>(
-        (static_cast<char*>(buffer().raw_ptr()) + array_desc_->offset));
+    return static_cast<T*>(array_desc_->data_ptr);
   }
 
   template <typename T>
   const T* data() const {
-    return const_cast<array&>(*this).data<T>();
-  }
-
-  int64_t offset() const {
-    return array_desc_->offset;
+    return static_cast<T*>(array_desc_->data_ptr);
   }
 
   enum Status {
@@ -474,8 +461,8 @@ class array {
     // can share the underlying data buffer.
     std::shared_ptr<Data> data;
 
-    // Offset from beginning of data pointer
-    int64_t offset{0};
+    // Properly offset data pointer
+    void* data_ptr{nullptr};
 
     // The size in elements of the data buffer the array accesses
     size_t data_size;

mlx/backend/common/binary.h

@@ -38,20 +38,20 @@ inline void set_binary_op_output_data(
     const array& a,
     const array& b,
     array& out,
-    BinaryOpType bopt,
-    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+    BinaryOpType bopt) {
   bool b_donatable = is_donatable(b, out);
   bool a_donatable = is_donatable(a, out);
   switch (bopt) {
     case BinaryOpType::ScalarScalar:
-      out.set_data(mallocfn(out.itemsize()), 1, a.strides(), a.flags());
+      out.set_data(
+          allocator::malloc(out.itemsize()), 1, a.strides(), a.flags());
       break;
     case BinaryOpType::ScalarVector:
       if (b_donatable) {
         out.copy_shared_buffer(b);
       } else {
         out.set_data(
-            mallocfn(b.data_size() * out.itemsize()),
+            allocator::malloc(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(
-            mallocfn(a.data_size() * out.itemsize()),
+            allocator::malloc(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(
-            mallocfn(a.data_size() * out.itemsize()),
+            allocator::malloc(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(mallocfn(out.nbytes()));
+        out.set_data(allocator::malloc(out.nbytes()));
       }
       break;
   }

mlx/backend/common/broadcasting.cpp

@@ -6,7 +6,7 @@ namespace mlx::core {
 
 void broadcast(const array& in, array& out) {
   if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
+    out.set_data(nullptr);
     return;
   }
   Strides strides(out.ndim(), 0);

mlx/backend/common/compiled.cpp

@@ -114,9 +114,7 @@ void compiled_allocate_outputs(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
     const std::function<bool(size_t)>& is_constant,
-    bool contiguous,
-    const std::function<allocator::Buffer(size_t)>&
-        mallocfn /* = allocator::malloc */) {
+    bool contiguous) {
   if (contiguous) {
     int o = 0;
     Strides strides;
@@ -142,7 +140,7 @@ void compiled_allocate_outputs(
     }
     for (; o < outputs.size(); ++o) {
       outputs[o].set_data(
-          mallocfn(data_size * outputs[o].itemsize()),
+          allocator::malloc(data_size * outputs[o].itemsize()),
           data_size,
           strides,
           flags);
@@ -165,7 +163,7 @@ void compiled_allocate_outputs(
       }
     }
     for (; o < outputs.size(); ++o) {
-      outputs[o].set_data(mallocfn(outputs[o].nbytes()));
+      outputs[o].set_data(allocator::malloc(outputs[o].nbytes()));
     }
   }
 }

mlx/backend/common/compiled.h

@@ -58,9 +58,7 @@ void compiled_allocate_outputs(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
     const std::function<bool(size_t)>& is_constant,
-    bool contiguous,
-    const std::function<allocator::Buffer(size_t)>& mallocfn =
-        allocator::malloc);
+    bool contiguous);
 
 // Collapse contiguous dims ignoring scalars and constants.
 std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(

mlx/backend/common/copy.h

@@ -22,11 +22,7 @@ enum class CopyType {
   GeneralGeneral
 };
 
-inline bool set_copy_output_data(
-    const array& in,
-    array& out,
-    CopyType ctype,
-    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+inline bool set_copy_output_data(const array& in, array& out, CopyType ctype) {
   if (ctype == CopyType::Vector) {
     // If the input is donateable, we are doing a vector copy and the types
     // have the same size, then the input buffer can hold the output.
@@ -35,14 +31,14 @@ inline bool set_copy_output_data(
       return true;
     } else {
       out.set_data(
-          mallocfn(in.data_size() * out.itemsize()),
+          allocator::malloc(in.data_size() * out.itemsize()),
           in.data_size(),
           in.strides(),
           in.flags());
       return false;
     }
   } else {
-    out.set_data(mallocfn(out.nbytes()));
+    out.set_data(allocator::malloc(out.nbytes()));
     return false;
   }
 }

mlx/backend/common/slicing.cpp

@@ -45,7 +45,7 @@ void slice(
     const Shape& start_indices,
     const Shape& strides) {
   if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
+    out.set_data(nullptr);
     return;
   }
 

mlx/backend/common/ternary.h

@@ -46,8 +46,7 @@ inline void set_ternary_op_output_data(
     const array& b,
     const array& c,
     array& out,
-    TernaryOpType topt,
-    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+    TernaryOpType topt) {
   auto maybe_donate = [&out](const array& x) {
     if (is_donatable(x, out)) {
       out.copy_shared_buffer(x);
@@ -58,12 +57,13 @@ inline void set_ternary_op_output_data(
 
   switch (topt) {
     case TernaryOpType::ScalarScalarScalar:
-      out.set_data(mallocfn(out.itemsize()), 1, b.strides(), b.flags());
+      out.set_data(
+          allocator::malloc(out.itemsize()), 1, b.strides(), b.flags());
       break;
     case TernaryOpType::VectorVectorVector:
       if (!(maybe_donate(a) || maybe_donate(b) || maybe_donate(c))) {
         out.set_data(
-            mallocfn(out.itemsize() * b.data_size()),
+            allocator::malloc(out.itemsize() * b.data_size()),
             b.data_size(),
             b.strides(),
             b.flags());
@@ -76,7 +76,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(mallocfn(out.nbytes()));
+        out.set_data(allocator::malloc(out.nbytes()));
       }
       break;
   }

mlx/backend/common/unary.h

@@ -7,22 +7,19 @@
 
 namespace mlx::core {
 
-inline void set_unary_output_data(
-    const array& in,
-    array& out,
-    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+inline void set_unary_output_data(const array& in, array& out) {
   if (in.flags().contiguous) {
     if (is_donatable(in, out)) {
       out.copy_shared_buffer(in);
     } else {
       out.set_data(
-          mallocfn(in.data_size() * out.itemsize()),
+          allocator::malloc(in.data_size() * out.itemsize()),
           in.data_size(),
           in.strides(),
           in.flags());
     }
   } else {
-    out.set_data(mallocfn(out.nbytes()));
+    out.set_data(allocator::malloc(out.nbytes()));
   }
 }
 

mlx/backend/cpu/binary.cpp

@@ -14,11 +14,233 @@
 
 namespace mlx::core {
 
+namespace {
+
+template <typename Op>
+void binary(const array& a, const array& b, array& out, Op op, Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        binary_op<bool, Op>(a, b, out, bopt);
+        break;
+      case uint8:
+        binary_op<uint8_t, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, Op>(a, b, out, bopt);
+        break;
+      case float16:
+        binary_op<float16_t, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, Op>(a, b, out, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t, Op>(a, b, out, bopt);
+        break;
+      case complex64:
+        binary_op<complex64_t, Op>(a, b, out, bopt);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void comparison_op(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (a.dtype()) {
+      case bool_:
+        binary_op<bool, bool, Op>(a, b, out, bopt);
+        break;
+      case uint8:
+        binary_op<uint8_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, bool, Op>(a, b, out, bopt);
+        break;
+      case float16:
+        binary_op<float16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, bool, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, bool, Op>(a, b, out, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case complex64:
+        binary_op<complex64_t, bool, Op>(a, b, out, bopt);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void binary_float(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case float16:
+        binary_op<float16_t, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, Op>(a, b, out, bopt);
+        break;
+      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.");
+    }
+  });
+}
+
+template <typename Op>
+void binary_int(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        binary_op<bool, Op>(a, b, out, bopt);
+      case uint8:
+        binary_op<uint8_t, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, Op>(a, b, out, bopt);
+        break;
+      default:
+        throw std::runtime_error("[binary_int] Type not supported");
+        break;
+    }
+  });
+}
+
+} // namespace
+
 void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Add(), stream());
+  binary(a, b, out, detail::Add(), stream());
 }
 
 void DivMod::eval_cpu(
@@ -102,14 +324,14 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Divide(), stream());
+  binary(a, b, out, detail::Divide(), stream());
 }
 
 void Remainder::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Remainder(), stream());
+  binary(a, b, out, detail::Remainder(), stream());
 }
 
 void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -150,90 +372,89 @@ void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
       }
     });
   } else {
-    comparison_op_cpu(a, b, out, detail::Equal(), stream());
+    comparison_op(a, b, out, detail::Equal(), stream());
   }
 }
 
 void Greater::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op_cpu(inputs[0], inputs[1], out, detail::Greater(), stream());
+  comparison_op(inputs[0], inputs[1], out, detail::Greater(), stream());
 }
 
 void GreaterEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op_cpu(
-      inputs[0], inputs[1], out, detail::GreaterEqual(), stream());
+  comparison_op(inputs[0], inputs[1], out, detail::GreaterEqual(), stream());
 }
 
 void Less::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op_cpu(inputs[0], inputs[1], out, detail::Less(), stream());
+  comparison_op(inputs[0], inputs[1], out, detail::Less(), stream());
 }
 
 void LessEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op_cpu(inputs[0], inputs[1], out, detail::LessEqual(), stream());
+  comparison_op(inputs[0], inputs[1], out, detail::LessEqual(), stream());
 }
 
 void LogAddExp::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_float_op_cpu(a, b, out, detail::LogAddExp(), stream());
+  binary_float(a, b, out, detail::LogAddExp(), stream());
 }
 
 void LogicalAnd::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2); // LogicalAnd requires two input arrays
   auto& in1 = inputs[0];
   auto& in2 = inputs[1];
-  binary_op_cpu(in1, in2, out, detail::LogicalAnd(), stream());
+  binary(in1, in2, out, detail::LogicalAnd(), stream());
 }
 
 void LogicalOr::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2); // LogicalOr requires two input arrays
   auto& in1 = inputs[0];
   auto& in2 = inputs[1];
-  binary_op_cpu(in1, in2, out, detail::LogicalOr(), stream());
+  binary(in1, in2, out, detail::LogicalOr(), stream());
 }
 
 void Maximum::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Maximum(), stream());
+  binary(a, b, out, detail::Maximum(), stream());
 }
 
 void Minimum::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Minimum(), stream());
+  binary(a, b, out, detail::Minimum(), stream());
 }
 
 void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Multiply(), stream());
+  binary(a, b, out, detail::Multiply(), stream());
 }
 
 void NotEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op_cpu(inputs[0], inputs[1], out, detail::NotEqual(), stream());
+  comparison_op(inputs[0], inputs[1], out, detail::NotEqual(), stream());
 }
 
 void Power::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Power(), stream());
+  binary(a, b, out, detail::Power(), stream());
 }
 
 void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary_op_cpu(a, b, out, detail::Subtract(), stream());
+  binary(a, b, out, detail::Subtract(), stream());
 }
 
 void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -242,19 +463,19 @@ void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
   switch (op_) {
     case BitwiseBinary::And:
-      binary_int_op_cpu(a, b, out, detail::BitwiseAnd(), stream());
+      binary_int(a, b, out, detail::BitwiseAnd(), stream());
       break;
     case BitwiseBinary::Or:
-      binary_int_op_cpu(a, b, out, detail::BitwiseOr(), stream());
+      binary_int(a, b, out, detail::BitwiseOr(), stream());
       break;
     case BitwiseBinary::Xor:
-      binary_int_op_cpu(a, b, out, detail::BitwiseXor(), stream());
+      binary_int(a, b, out, detail::BitwiseXor(), stream());
       break;
     case BitwiseBinary::LeftShift:
-      binary_int_op_cpu(a, b, out, detail::LeftShift(), stream());
+      binary_int(a, b, out, detail::LeftShift(), stream());
       break;
     case BitwiseBinary::RightShift:
-      binary_int_op_cpu(a, b, out, detail::RightShift(), stream());
+      binary_int(a, b, out, detail::RightShift(), stream());
       break;
   }
 }
@@ -263,7 +484,7 @@ void ArcTan2::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   const auto& a = inputs[0];
   const auto& b = inputs[1];
-  binary_float_op_cpu(a, b, out, detail::ArcTan2(), stream());
+  binary_float(a, b, out, detail::ArcTan2(), stream());
 }
 
 } // namespace mlx::core

mlx/backend/cpu/binary.h

@@ -7,7 +7,6 @@
 #include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/utils.h"
 
-#include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
 
 namespace mlx::core {
@@ -291,227 +290,4 @@ void binary_op(const array& a, const array& b, array& out, BinaryOpType bopt) {
   binary_op<T, T, Op>(a, b, out, bopt);
 }
 
-template <typename Op>
-void binary_op_cpu(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
-    Stream stream) {
-  auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
-
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(a);
-  encoder.set_input_array(b);
-  encoder.set_output_array(out);
-  encoder.dispatch([a = array::unsafe_weak_copy(a),
-                    b = array::unsafe_weak_copy(b),
-                    out = array::unsafe_weak_copy(out),
-                    bopt]() mutable {
-    switch (out.dtype()) {
-      case bool_:
-        binary_op<bool, Op>(a, b, out, bopt);
-        break;
-      case uint8:
-        binary_op<uint8_t, Op>(a, b, out, bopt);
-        break;
-      case uint16:
-        binary_op<uint16_t, Op>(a, b, out, bopt);
-        break;
-      case uint32:
-        binary_op<uint32_t, Op>(a, b, out, bopt);
-        break;
-      case uint64:
-        binary_op<uint64_t, Op>(a, b, out, bopt);
-        break;
-      case int8:
-        binary_op<int8_t, Op>(a, b, out, bopt);
-        break;
-      case int16:
-        binary_op<int16_t, Op>(a, b, out, bopt);
-        break;
-      case int32:
-        binary_op<int32_t, Op>(a, b, out, bopt);
-        break;
-      case int64:
-        binary_op<int64_t, Op>(a, b, out, bopt);
-        break;
-      case float16:
-        binary_op<float16_t, Op>(a, b, out, bopt);
-        break;
-      case float32:
-        binary_op<float, Op>(a, b, out, bopt);
-        break;
-      case float64:
-        binary_op<double, Op>(a, b, out, bopt);
-        break;
-      case bfloat16:
-        binary_op<bfloat16_t, Op>(a, b, out, bopt);
-        break;
-      case complex64:
-        binary_op<complex64_t, Op>(a, b, out, bopt);
-        break;
-    }
-  });
-}
-
-template <typename Op>
-void comparison_op_cpu(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
-    Stream stream) {
-  auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
-
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(a);
-  encoder.set_input_array(b);
-  encoder.set_output_array(out);
-  encoder.dispatch([a = array::unsafe_weak_copy(a),
-                    b = array::unsafe_weak_copy(b),
-                    out = array::unsafe_weak_copy(out),
-                    bopt]() mutable {
-    switch (a.dtype()) {
-      case bool_:
-        binary_op<bool, bool, Op>(a, b, out, bopt);
-        break;
-      case uint8:
-        binary_op<uint8_t, bool, Op>(a, b, out, bopt);
-        break;
-      case uint16:
-        binary_op<uint16_t, bool, Op>(a, b, out, bopt);
-        break;
-      case uint32:
-        binary_op<uint32_t, bool, Op>(a, b, out, bopt);
-        break;
-      case uint64:
-        binary_op<uint64_t, bool, Op>(a, b, out, bopt);
-        break;
-      case int8:
-        binary_op<int8_t, bool, Op>(a, b, out, bopt);
-        break;
-      case int16:
-        binary_op<int16_t, bool, Op>(a, b, out, bopt);
-        break;
-      case int32:
-        binary_op<int32_t, bool, Op>(a, b, out, bopt);
-        break;
-      case int64:
-        binary_op<int64_t, bool, Op>(a, b, out, bopt);
-        break;
-      case float16:
-        binary_op<float16_t, bool, Op>(a, b, out, bopt);
-        break;
-      case float32:
-        binary_op<float, bool, Op>(a, b, out, bopt);
-        break;
-      case float64:
-        binary_op<double, bool, Op>(a, b, out, bopt);
-        break;
-      case bfloat16:
-        binary_op<bfloat16_t, bool, Op>(a, b, out, bopt);
-        break;
-      case complex64:
-        binary_op<complex64_t, bool, Op>(a, b, out, bopt);
-        break;
-    }
-  });
-}
-
-template <typename Op>
-void binary_float_op_cpu(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
-    Stream stream) {
-  auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
-
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(a);
-  encoder.set_input_array(b);
-  encoder.set_output_array(out);
-  encoder.dispatch([a = array::unsafe_weak_copy(a),
-                    b = array::unsafe_weak_copy(b),
-                    out = array::unsafe_weak_copy(out),
-                    bopt]() mutable {
-    switch (out.dtype()) {
-      case float16:
-        binary_op<float16_t, Op>(a, b, out, bopt);
-        break;
-      case float32:
-        binary_op<float, Op>(a, b, out, bopt);
-        break;
-      case float64:
-        binary_op<double, Op>(a, b, out, bopt);
-        break;
-      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.");
-    }
-  });
-}
-
-template <typename Op>
-void binary_int_op_cpu(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
-    Stream stream) {
-  auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
-
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(a);
-  encoder.set_input_array(b);
-  encoder.set_output_array(out);
-  encoder.dispatch([a = array::unsafe_weak_copy(a),
-                    b = array::unsafe_weak_copy(b),
-                    out = array::unsafe_weak_copy(out),
-                    bopt]() mutable {
-    switch (out.dtype()) {
-      case bool_:
-        binary_op<bool, Op>(a, b, out, bopt);
-      case uint8:
-        binary_op<uint8_t, Op>(a, b, out, bopt);
-        break;
-      case uint16:
-        binary_op<uint16_t, Op>(a, b, out, bopt);
-        break;
-      case uint32:
-        binary_op<uint32_t, Op>(a, b, out, bopt);
-        break;
-      case uint64:
-        binary_op<uint64_t, Op>(a, b, out, bopt);
-        break;
-      case int8:
-        binary_op<int8_t, Op>(a, b, out, bopt);
-        break;
-      case int16:
-        binary_op<int16_t, Op>(a, b, out, bopt);
-        break;
-      case int32:
-        binary_op<int32_t, Op>(a, b, out, bopt);
-        break;
-      case int64:
-        binary_op<int64_t, Op>(a, b, out, bopt);
-        break;
-      default:
-        throw std::runtime_error("[binary_int] Type not supported");
-        break;
-    }
-  });
-}
-
 } // namespace mlx::core

mlx/backend/cpu/conv.cpp

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

mlx/backend/cpu/distributed.cpp

@@ -95,9 +95,4 @@ void Recv::eval_cpu(
   distributed::detail::recv(group(), outputs[0], src_, stream());
 }
 
-void ReduceScatter::eval_cpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  throw std::runtime_error("[ReduceScatter] Not implemented yet.");
-}
 } // namespace mlx::core::distributed

mlx/backend/cpu/eig.cpp

@@ -46,6 +46,7 @@ void eig_impl(
     int info;
     {
       T work;
+      int iwork;
       geev<T>(
           &jobl,
           &jobr,

mlx/backend/cpu/gemms/bnns.cpp

@@ -1,4 +1,5 @@
 // Copyright © 2023-2024 Apple Inc.
+
 #include <Accelerate/Accelerate.h>
 
 #include "mlx/array.h"
@@ -48,15 +49,9 @@ void matmul_bnns(
   size_t K = a_shape[ndim - 1];
 
   BNNSDataType bnns_dtype = to_bnns_dtype<T>();
+
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
-  if (beta != 1.0 && beta != 0.0) {
-    // scale the output
-    for (auto i = 0; i < batch_size * M * N; ++i) {
-      out[i] *= beta;
-    }
-    beta = 1.0;
-  }
   const BNNSLayerParametersBroadcastMatMul gemm_params{
       /* float alpha = */ alpha,
       /* float beta = */ beta,

mlx/backend/cpu/masked_mm.cpp

@@ -215,18 +215,18 @@ void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   encoder.set_input_array(a);
   encoder.set_input_array(b);
-  const void* a_mask_ptr = nullptr;
-  const void* b_mask_ptr = nullptr;
-  const void* out_mask_ptr = nullptr;
+  const void* a_mask_ptr;
+  const void* b_mask_ptr;
+  const void* out_mask_ptr;
   Shape a_mask_shape;
   Shape b_mask_shape;
   Shape out_mask_shape;
   Strides a_mask_strides;
   Strides b_mask_strides;
   Strides out_mask_strides;
-  bool a_mask_bool = false;
-  bool b_mask_bool = false;
-  bool out_mask_bool = false;
+  bool a_mask_bool;
+  bool b_mask_bool;
+  bool out_mask_bool;
   if (has_op_mask) {
     auto& a_mask = inputs[inputs.size() - 2];
     auto& b_mask = inputs[inputs.size() - 1];
@@ -423,6 +423,7 @@ void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& rhs_indices = inputs[3];
 
   auto batch_shape = get_batch_dims(out.shape());
+  int batch_ndim = batch_shape.size();
 
   auto batch_shape_A = get_batch_dims(a.shape());
   auto batch_strides_A = get_batch_dims(a.strides());

mlx/backend/cpu/matmul.cpp

@@ -2,8 +2,6 @@
 
 #include <cstring>
 #include "mlx/array.h"
-#include "mlx/backend/cpu/binary.h"
-#include "mlx/backend/cpu/binary_ops.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/gemm.h"
@@ -93,6 +91,7 @@ void matmul_general(
   auto [b_transposed, ldb, b] = check_transpose(b_pre);
   size_t M = a.shape(-2);
   size_t N = b.shape(-1);
+  size_t K = a.shape(-1);
   if (M == 0 || N == 0) {
     return;
   }
@@ -137,29 +136,15 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
-  // Handle empty matrix case (K=0)
-  if (inputs[0].shape(-1) == 0) {
-    auto& c = inputs[2];
-    if (beta_ == 1.0f) {
-      CopyType ctype = c.data_size() == 1
-          ? CopyType::Scalar
-          : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
-      copy_cpu(c, out, ctype, stream());
-    } else {
-      array beta_scalar = array(beta_, c.dtype());
-      auto& encoder = cpu::get_command_encoder(stream());
-      binary_float_op_cpu(c, beta_scalar, out, detail::Multiply(), stream());
-      encoder.add_temporary(std::move(beta_scalar));
-    }
-    return;
-  }
-
   // Fill output with C
   auto& c = inputs[2];
   CopyType ctype = c.data_size() == 1
       ? CopyType::Scalar
       : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
   copy_cpu(c, out, ctype, stream());
+  if (inputs[0].shape(-1) == 0) {
+    return;
+  }
   matmul_general(inputs[0], inputs[1], out, stream(), alpha_, beta_);
 }
 

mlx/backend/cpu/primitives.cpp

@@ -333,7 +333,7 @@ void Reshape::eval_cpu(const std::vector<array>& inputs, array& out) {
 
 void DynamicSlice::eval_cpu(const std::vector<array>& inputs, array& out) {
   if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
+    out.set_data(nullptr);
     return;
   }
   auto& in = inputs[0];
@@ -361,7 +361,7 @@ void DynamicSliceUpdate::eval_cpu(
     const std::vector<array>& inputs,
     array& out) {
   if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
+    out.set_data(nullptr);
     return;
   }
 
@@ -396,7 +396,7 @@ void DynamicSliceUpdate::eval_cpu(
 void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
+    out.set_data(nullptr);
     return;
   }
 

mlx/backend/cpu/quantized.cpp

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

mlx/backend/cpu/simd/accelerate_simd.h

@@ -217,20 +217,14 @@ Simd<T, N> atan2(Simd<T, N> a, Simd<T, N> b) {
 
 template <typename T, int N>
 Simd<T, N> maximum(Simd<T, N> a, Simd<T, N> b) {
-  auto out = Simd<T, N>(asd::max(a.value, b.value));
-  if constexpr (!std::is_integral_v<T>) {
-    out = select(isnan(b), b, select(isnan(a), a, out));
-  }
-  return out;
+  // TODO add isnan
+  return asd::max(a.value, b.value);
 }
 
 template <typename T, int N>
 Simd<T, N> minimum(Simd<T, N> a, Simd<T, N> b) {
-  auto out = Simd<T, N>(asd::min(a.value, b.value));
-  if constexpr (!std::is_integral_v<T>) {
-    out = select(isnan(b), b, select(isnan(a), a, out));
-  }
-  return out;
+  // TODO add isnan
+  return asd::min(a.value, b.value);
 }
 
 template <typename T, int N>

mlx/backend/cpu/sort.cpp

@@ -39,7 +39,7 @@ struct StridedIterator {
   StridedIterator() = default;
 
   explicit StridedIterator(T* ptr, int64_t stride, difference_type offset = 0)
-      : stride_(stride), ptr_(ptr + offset * stride) {}
+      : ptr_(ptr + offset * stride), stride_(stride) {}
 
   explicit StridedIterator(array& arr, int axis, difference_type offset = 0)
       : StridedIterator(arr.data<T>(), arr.strides()[axis], offset) {}

mlx/backend/cpu/svd.cpp

@@ -83,6 +83,8 @@ void svd_impl(
 
     auto jobz = (u_ptr) ? "A" : "N";
 
+    // Will contain the number of singular values after the call has returned.
+    int ns = 0;
     T workspace_dimension = 0;
 
     // Will contain the indices of eigenvectors that failed to converge (not

mlx/backend/cpu/unary_ops.h

@@ -120,7 +120,7 @@ Simd<uint32_t, N> fp32_to_bits(Simd<float, N> x) {
 struct ToFP8 {
   template <typename T, int N>
   Simd<uint8_t, N> operator()(Simd<T, N> f) {
-    uint32_t fp8_max = 543 << 21;
+    uint32_t fp8_max = 1087 << 20;
     auto denorm_mask = Simd<uint32_t, N>(141 << 23);
     Simd<uint32_t, N> f_bits;
     Simd<float, N> f32 = f;

mlx/backend/cuda/CMakeLists.txt

@@ -32,7 +32,6 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/layer_norm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp

mlx/backend/cuda/allocator.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/utils.h"
 #include "mlx/utils.h"
 
@@ -68,7 +67,6 @@ CudaBuffer* SmallSizePool::malloc() {
   next_free_ = next_free_->next;
   b->buf.data = static_cast<char*>(data_) + i * small_block_size;
   b->buf.size = small_block_size;
-  b->buf.device = -1;
   return &b->buf;
 }
 
@@ -90,41 +88,16 @@ CudaAllocator::CudaAllocator()
           page_size,
           [](CudaBuffer* buf) { return buf->size; },
           [this](CudaBuffer* buf) { cuda_free(buf); }) {
+  // TODO: Set memory limit for multi-device.
   size_t free, total;
   CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
   memory_limit_ = total * 0.95;
   max_pool_size_ = memory_limit_;
-
-  int device_count = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
-  int curr;
-  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
-  for (int i = 0; i < device_count; ++i) {
-    CHECK_CUDA_ERROR(cudaSetDevice(i));
-    cudaStream_t s;
-    CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));
-    free_streams_.push_back(s);
-  }
-  CHECK_CUDA_ERROR(cudaSetDevice(curr));
 }
 
-void copy_to_managed(CudaBuffer& buf) {
-  // TODO maybe make this async on a i/o stream to avoid synchronizing the
-  // device on malloc/and free
-  void* new_data;
-  CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, buf.size));
-  buf.device = -1;
-  CHECK_CUDA_ERROR(cudaMemcpy(new_data, buf.data, buf.size, cudaMemcpyDefault));
-  CHECK_CUDA_ERROR(cudaFree(buf.data));
-  buf.data = new_data;
-}
-
-Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
-  if (size == 0) {
-    return Buffer{new CudaBuffer{nullptr, 0, -1}};
-  }
-
+Buffer CudaAllocator::malloc(size_t size) {
   // Find available buffer from cache.
+  auto orig_size = size;
   std::unique_lock lock(mutex_);
   if (size <= small_block_size) {
     size = 8;
@@ -134,11 +107,6 @@ Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
     size = page_size * ((size + page_size - 1) / page_size);
   }
 
-  int device = -1;
-  if (size > small_block_size && stream != nullptr) {
-    CHECK_CUDA_ERROR(cudaStreamGetDevice(stream, &device));
-  }
-
   CudaBuffer* buf = buffer_cache_.reuse_from_cache(size);
   if (!buf) {
     // If we have a lot of memory pressure try to reclaim memory from the cache.
@@ -154,13 +122,8 @@ Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
     }
     lock.unlock();
     if (!buf) {
-      buf = new CudaBuffer{nullptr, size, device};
-      cudaError_t err;
-      if (device == -1) {
-        err = cudaMallocManaged(&buf->data, size);
-      } else {
-        err = cudaMallocAsync(&buf->data, size, stream);
-      }
+      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)));
@@ -168,37 +131,21 @@ Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
     }
     lock.lock();
   }
-  active_memory_ += buf->size;
+  active_memory_ += size;
   peak_memory_ = std::max(active_memory_, peak_memory_);
 
   // Maintain the cache below the requested limit.
   if (get_cache_memory() > max_pool_size_) {
     buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
   }
-  // Copy to managed here if the buffer is not on the right device
-  if (buf->device != device) {
-    copy_to_managed(*buf);
-  }
   return Buffer{buf};
 }
 
-Buffer CudaAllocator::malloc_async(size_t size, cudaStream_t stream) {
-  return malloc_impl(size, stream);
-}
-
-Buffer CudaAllocator::malloc(size_t size) {
-  return malloc_impl(size, nullptr);
-}
-
 void CudaAllocator::free(Buffer buffer) {
   auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
   if (!buf) {
     return;
   }
-  if (buf->size == 0) {
-    delete buf;
-    return;
-  }
 
   std::unique_lock lock(mutex_);
   active_memory_ -= buf->size;
@@ -222,11 +169,7 @@ void CudaAllocator::cuda_free(CudaBuffer* buf) {
   if (scalar_pool_.in_pool(buf)) {
     scalar_pool_.free(buf);
   } else {
-    if (buf->device >= 0) {
-      cudaFreeAsync(buf->data, free_streams_[buf->device]);
-    } else {
-      cudaFree(buf->data);
-    }
+    cudaFree(buf->data);
     delete buf;
   }
 }
@@ -277,16 +220,6 @@ CudaAllocator& allocator() {
   return *allocator_;
 }
 
-Buffer malloc_async(size_t size, cudaStream_t stream) {
-  auto buffer = allocator().malloc_async(size, stream);
-  if (size && !buffer.ptr()) {
-    std::ostringstream msg;
-    msg << "[malloc_async] Unable to allocate " << size << " bytes.";
-    throw std::runtime_error(msg.str());
-  }
-  return buffer;
-}
-
 } // namespace cu
 
 namespace allocator {
@@ -299,11 +232,7 @@ void* Buffer::raw_ptr() {
   if (!ptr_) {
     return nullptr;
   }
-  auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
-  if (cbuf.device != -1) {
-    copy_to_managed(cbuf);
-  }
-  return cbuf.data;
+  return static_cast<cu::CudaBuffer*>(ptr_)->data;
 }
 
 } // namespace allocator

mlx/backend/cuda/allocator.h

@@ -4,9 +4,7 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/buffer_cache.h"
-#include "mlx/backend/cuda/cuda_utils.h"
 
-#include <cuda_runtime.h>
 #include <mutex>
 #include <set>
 #include <utility>
@@ -19,7 +17,6 @@ using allocator::Buffer;
 struct CudaBuffer {
   void* data;
   size_t size;
-  int device; // -1 for managed
 };
 
 class SmallSizePool {
@@ -48,7 +45,6 @@ class SmallSizePool {
 class CudaAllocator : public allocator::Allocator {
  public:
   Buffer malloc(size_t size) override;
-  Buffer malloc_async(size_t size, cudaStream_t stream);
   void free(Buffer buffer) override;
   size_t size(Buffer buffer) const override;
 
@@ -62,7 +58,6 @@ class CudaAllocator : public allocator::Allocator {
   void clear_cache();
 
  private:
-  Buffer malloc_impl(size_t size, cudaStream_t stream);
   void cuda_free(CudaBuffer* buf);
 
   CudaAllocator();
@@ -74,12 +69,9 @@ class CudaAllocator : public allocator::Allocator {
   BufferCache<CudaBuffer> buffer_cache_;
   size_t active_memory_{0};
   size_t peak_memory_{0};
-  std::vector<cudaStream_t> free_streams_;
   SmallSizePool scalar_pool_;
 };
 
 CudaAllocator& allocator();
 
-Buffer malloc_async(size_t size, cudaStream_t stream);
-
 } // namespace mlx::core::cu

mlx/backend/cuda/arange.cu

@@ -41,8 +41,9 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
   if (out.size() == 0) {
     return;
   }
+  out.set_data(allocator::malloc(out.nbytes()));
+
   auto& encoder = cu::get_command_encoder(stream());
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
   encoder.set_output_array(out);
 
   dispatch_int_float_types(out.dtype(), "Arange", [&](auto type_tag) {
@@ -57,7 +58,7 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
           num_blocks,
           block_dims,
           0,
-          gpu_ptr<OutType>(out),
+          out.data<OutType>(),
           out.data_size(),
           static_cast<CTYPE>(start_),
           static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_));

mlx/backend/cuda/arg_reduce.cu

@@ -140,10 +140,8 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
   nvtx3::scoped_range r("ArgReduce::eval_gpu");
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-
+  out.set_data(allocator::malloc(out.nbytes()));
   auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
 
   // Prepare the shapes, strides and axis arguments.
   Shape shape = remove_index(in.shape(), axis_);
@@ -156,6 +154,7 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
   int32_t ndim = shape.size();
 
   // ArgReduce.
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(in);
   encoder.set_output_array(out);
   dispatch_real_types(in.dtype(), "ArgReduce", [&](auto type_tag) {
@@ -173,8 +172,8 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
           num_blocks,
           block_dim(),
           0,
-          gpu_ptr<T>(in),
-          gpu_ptr<uint32_t>(out),
+          in.data<T>(),
+          out.data<uint32_t>(),
           out.size(),
           const_param(shape),
           const_param(in_strides),

mlx/backend/cuda/binary/binary.cuh

@@ -292,9 +292,9 @@ void binary_op_gpu_inplace(
                         {num_blocks_x, num_blocks_y},
                         block_dims,
                         0,
-                        gpu_ptr<InType>(a),
-                        gpu_ptr<InType>(b),
-                        gpu_ptr<OutType>(out),
+                        a.data<InType>(),
+                        b.data<InType>(),
+                        out.data<OutType>(),
                         rest,
                         const_param<dims_constant()>(shape),
                         const_param<dims_constant()>(a_strides),
@@ -310,9 +310,9 @@ void binary_op_gpu_inplace(
                       {num_blocks_x, num_blocks_y},
                       block_dims,
                       0,
-                      gpu_ptr<InType>(a),
-                      gpu_ptr<InType>(b),
-                      gpu_ptr<OutType>(out),
+                      a.data<InType>(),
+                      b.data<InType>(),
+                      out.data<OutType>(),
                       rest,
                       const_param(shape),
                       const_param(a_strides),
@@ -339,9 +339,9 @@ void binary_op_gpu_inplace(
                 num_blocks,
                 block_dims,
                 0,
-                gpu_ptr<InType>(a),
-                gpu_ptr<InType>(b),
-                gpu_ptr<OutType>(out),
+                a.data<InType>(),
+                b.data<InType>(),
+                out.data<OutType>(),
                 out.data_size());
           });
         }
@@ -365,11 +365,7 @@ void binary_op_gpu(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  auto& encoder = cu::get_command_encoder(s);
-
-  set_binary_op_output_data(a, b, out, bopt, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
+  set_binary_op_output_data(a, b, out, bopt);
   binary_op_gpu_inplace<Op>(inputs, out, op, s);
 }
 

mlx/backend/cuda/binary_two.cu

@@ -245,18 +245,14 @@ void binary_two_op_gpu_inplace(
   auto& out_a = outputs[0];
   auto& out_b = outputs[1];
   auto bopt = get_binary_op_type(a, b);
-  auto& encoder = cu::get_command_encoder(s);
-  set_binary_op_output_data(a, b, out_a, bopt, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
-  set_binary_op_output_data(a, b, out_b, bopt, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
+  set_binary_op_output_data(a, b, out_a, bopt);
+  set_binary_op_output_data(a, b, out_b, bopt);
 
   if (out_a.size() == 0) {
     return;
   }
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(a);
   encoder.set_input_array(b);
   encoder.set_output_array(out_a);
@@ -317,10 +313,10 @@ void binary_two_op_gpu_inplace(
                         {num_blocks_x, num_blocks_y},
                         block_dims,
                         0,
-                        gpu_ptr<InType>(a),
-                        gpu_ptr<InType>(b),
-                        gpu_ptr<OutType>(out_a),
-                        gpu_ptr<OutType>(out_b),
+                        a.data<InType>(),
+                        b.data<InType>(),
+                        out_a.data<OutType>(),
+                        out_b.data<OutType>(),
                         rest,
                         const_param<dims_constant()>(shape),
                         const_param<dims_constant()>(a_strides),
@@ -336,10 +332,10 @@ void binary_two_op_gpu_inplace(
                       {num_blocks_x, num_blocks_y},
                       block_dims,
                       0,
-                      gpu_ptr<InType>(a),
-                      gpu_ptr<InType>(b),
-                      gpu_ptr<OutType>(out_a),
-                      gpu_ptr<OutType>(out_b),
+                      a.data<InType>(),
+                      b.data<InType>(),
+                      out_a.data<OutType>(),
+                      out_b.data<OutType>(),
                       rest,
                       const_param(shape),
                       const_param(a_strides),
@@ -370,10 +366,10 @@ void binary_two_op_gpu_inplace(
                 num_blocks,
                 block_dims,
                 0,
-                gpu_ptr<InType>(a),
-                gpu_ptr<InType>(b),
-                gpu_ptr<OutType>(out_a),
-                gpu_ptr<OutType>(out_b),
+                a.data<InType>(),
+                b.data<InType>(),
+                out_a.data<OutType>(),
+                out_b.data<OutType>(),
                 out_a.data_size());
           });
         }

mlx/backend/cuda/compiled.cpp

@@ -293,13 +293,8 @@ void Compiled::eval_gpu(
     }
   }
 
-  auto& encoder = cu::get_command_encoder(s);
-
   // Put outputs.
-  compiled_allocate_outputs(
-      inputs, outputs, is_constant_, contiguous, [&](auto n) {
-        return cu::malloc_async(n, encoder.stream());
-      });
+  compiled_allocate_outputs(inputs, outputs, is_constant_, contiguous);
   for (auto& x : outputs) {
     args.append(x);
   }
@@ -329,6 +324,7 @@ void Compiled::eval_gpu(
     kernel_name += fmt::format(
         "_strided<{}, {}, {}>", shape.size(), index_type, work_per_thread);
   }
+  auto& encoder = cu::get_command_encoder(s);
   for (const auto& in : inputs) {
     encoder.set_input_array(in);
   }

mlx/backend/cuda/conv.cpp

@@ -270,16 +270,17 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
   if (out_.size() == 0) {
     return;
   }
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
 
   assert(inputs.size() == 2);
   array in = inputs[0];
   array wt = inputs[1];
   array out = out_;
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
   Dtype dtype = out.dtype();
 
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
   // Search cache.
   ConvCacheKey cache_key{
       encoder.device().cuda_device(),

mlx/backend/cuda/conv/gemm_conv.cu

@@ -86,7 +86,7 @@ array unfold_inputs_nd(
     int mat_N,
     ConvParams<NDIM>& params) {
   array unfolded({mat_M, mat_K}, in.dtype(), nullptr, {});
-  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder.stream()));
+  unfolded.set_data(allocator::malloc(unfolded.nbytes()));
   encoder.add_temporary(unfolded);
 
   int filter_size = params.C;
@@ -118,8 +118,8 @@ array unfold_inputs_nd(
         num_blocks,
         block_dims,
         0,
-        gpu_ptr<DataType>(in),
-        gpu_ptr<DataType>(unfolded),
+        in.data<DataType>(),
+        unfolded.data<DataType>(),
         filter_size,
         out_pixels,
         params);

mlx/backend/cuda/conv/gemm_grouped_conv.cu

@@ -89,7 +89,7 @@ array grouped_unfold_transpose_inputs_nd(
     int mat_N,
     ConvParams<NDIM>& params) {
   array unfolded({mat_M, mat_K * params.groups}, in.dtype(), nullptr, {});
-  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder.stream()));
+  unfolded.set_data(allocator::malloc(unfolded.nbytes()));
   encoder.add_temporary(unfolded);
 
   int filter_size = params.C;
@@ -121,8 +121,8 @@ array grouped_unfold_transpose_inputs_nd(
         num_blocks,
         block_dims,
         0,
-        gpu_ptr<DataType>(in),
-        gpu_ptr<DataType>(unfolded),
+        in.data<DataType>(),
+        unfolded.data<DataType>(),
         filter_size,
         out_pixels,
         params);

mlx/backend/cuda/copy.cu

@@ -5,22 +5,6 @@
 
 namespace mlx::core {
 
-void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
-  auto& encoder = cu::get_command_encoder(s);
-  bool donated = set_copy_output_data(in, out, ctype, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
-  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_inplace(
     const array& in,
     array& out,
@@ -103,31 +87,11 @@ void fill_gpu(const array& in, array& out, const Stream& s) {
   if (out.size() == 0) {
     return;
   }
+  out.set_data(allocator::malloc(out.nbytes()));
   auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
   encoder.set_input_array(in);
   encoder.set_output_array(out);
   copy_contiguous(encoder, CopyType::Scalar, in, out, 0, 0);
 }
 
-void reshape_gpu(const array& in, array& out, Stream s) {
-  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
-  if (copy_necessary) {
-    auto& encoder = cu::get_command_encoder(s);
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
-    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 mlx::core

mlx/backend/cuda/copy/copy_contiguous.cu

@@ -77,8 +77,8 @@ void copy_contiguous(
             num_blocks,
             block_dims,
             0,
-            gpu_ptr<InType>(in) + in_offset,
-            gpu_ptr<OutType>(out) + out_offset,
+            in.data<InType>() + in_offset,
+            out.data<OutType>() + out_offset,
             out.data_size());
       });
     });

mlx/backend/cuda/copy/copy_general.cu

@@ -106,8 +106,8 @@ void copy_general(
             using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
             using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
-            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
-            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+            const InType* in_ptr = in.data<InType>() + offset_in;
+            OutType* out_ptr = out.data<OutType>() + offset_out;
             int ndim = shape.size();
             size_t data_size = 1;
             for (auto& s : shape)

mlx/backend/cuda/copy/copy_general_dynamic.cu

@@ -69,8 +69,8 @@ void copy_general_dynamic(
             using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
             using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
-            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
-            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+            const InType* in_ptr = in.data<InType>() + offset_in;
+            OutType* out_ptr = out.data<OutType>() + offset_out;
             int ndim = shape.size();
             if (ndim <= 3) {
               dispatch_1_2_3(ndim, [&](auto dims_constant) {
@@ -90,8 +90,8 @@ void copy_general_dynamic(
                     const_param<dims_constant()>(shape),
                     const_param<dims_constant()>(strides_in),
                     const_param<dims_constant()>(strides_out),
-                    gpu_ptr<int64_t>(dynamic_offset_in),
-                    gpu_ptr<int64_t>(dynamic_offset_out));
+                    dynamic_offset_in.data<int64_t>(),
+                    dynamic_offset_out.data<int64_t>());
               });
             } else { // ndim >= 4
               auto [num_blocks, block_dims] = get_launch_args(out, large());
@@ -107,8 +107,8 @@ void copy_general_dynamic(
                   const_param(strides_in),
                   const_param(strides_out),
                   ndim,
-                  gpu_ptr<int64_t>(dynamic_offset_in),
-                  gpu_ptr<int64_t>(dynamic_offset_out));
+                  dynamic_offset_in.data<int64_t>(),
+                  dynamic_offset_out.data<int64_t>());
             }
           });
     });

mlx/backend/cuda/copy/copy_general_input.cu

@@ -92,8 +92,8 @@ void copy_general_input(
             using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
             using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
-            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
-            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+            const InType* in_ptr = in.data<InType>() + offset_in;
+            OutType* out_ptr = out.data<OutType>() + offset_out;
             int ndim = shape.size();
             int work_per_thread = 1;
             auto dim0 = ndim > 0 ? shape.back() : 1;

mlx/backend/cuda/cuda_utils.h

@@ -1,82 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cublasLt.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-namespace mlx::core {
-
-// Throw exception if the cuda API does not succeed.
-void check_cublas_error(const char* name, cublasStatus_t err);
-void check_cuda_error(const char* name, cudaError_t err);
-void check_cuda_error(const char* name, CUresult err);
-
-// The macro version that prints the command that failed.
-#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
-#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
-
-// Base class for RAII managed CUDA resources.
-template <typename Handle, cudaError_t (*Destroy)(Handle)>
-class CudaHandle {
- public:
-  CudaHandle(Handle handle = nullptr) : handle_(handle) {}
-
-  CudaHandle(CudaHandle&& other) : handle_(other.handle_) {
-    assert(this != &other);
-    other.handle_ = nullptr;
-  }
-
-  ~CudaHandle() {
-    reset();
-  }
-
-  CudaHandle(const CudaHandle&) = delete;
-  CudaHandle& operator=(const CudaHandle&) = delete;
-
-  CudaHandle& operator=(CudaHandle&& other) {
-    assert(this != &other);
-    reset();
-    std::swap(handle_, other.handle_);
-    return *this;
-  }
-
-  void reset() {
-    if (handle_ != nullptr) {
-      CHECK_CUDA_ERROR(Destroy(handle_));
-      handle_ = nullptr;
-    }
-  }
-
-  operator Handle() const {
-    return handle_;
-  }
-
- protected:
-  Handle handle_;
-};
-
-namespace cu {
-class Device;
-}; // namespace cu
-
-// Wrappers of CUDA resources.
-class CudaGraph : public CudaHandle<cudaGraph_t, cudaGraphDestroy> {
- public:
-  using CudaHandle::CudaHandle;
-  explicit CudaGraph(cu::Device& device);
-  void end_capture(cudaStream_t stream);
-};
-
-class CudaGraphExec : public CudaHandle<cudaGraphExec_t, cudaGraphExecDestroy> {
- public:
-  void instantiate(cudaGraph_t graph);
-};
-
-class CudaStream : public CudaHandle<cudaStream_t, cudaStreamDestroy> {
- public:
-  explicit CudaStream(cu::Device& device);
-};
-
-} // namespace mlx::core

mlx/backend/cuda/cudnn_utils.cpp

@@ -132,18 +132,14 @@ bool prepare_cudnn_plan(
     void** data_ptrs,
     F&& execute) {
   int workspace_size = plan.getWorkspaceSize();
-  void* workspace_ptr = nullptr;
-  if (workspace_size > 0) {
-    array workspace(
-        cu::malloc_async(workspace_size, encoder.stream()),
-        {workspace_size},
-        uint8);
-    encoder.add_temporary(workspace);
-    workspace_ptr = gpu_ptr<void>(workspace);
-  }
+  array workspace(
+      workspace_size > 0 ? allocator::malloc(workspace_size)
+                         : allocator::Buffer(nullptr),
+      {workspace_size},
+      uint8);
 
   auto args = cudnn_frontend::VariantPackBuilder()
-                  .setWorkspacePointer(workspace_ptr)
+                  .setWorkspacePointer(workspace.data<void>())
                   .setDataPointers(num_args, data_ptrs)
                   .setUids(num_args, uids)
                   .build();
@@ -155,6 +151,7 @@ bool prepare_cudnn_plan(
     return false;
   }
 
+  encoder.add_temporary(workspace);
   return true;
 }
 

mlx/backend/cuda/cudnn_utils.h

@@ -3,7 +3,6 @@
 #pragma once
 
 #include "mlx/array.h"
-#include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device/config.h"
 #include "mlx/backend/cuda/utils.h"
 #include "mlx/dtype_utils.h"
@@ -24,7 +23,7 @@ class CommandEncoder;
 // Return pointer alignment of |x|'s data.
 inline uint8_t get_alignment(const array& x) {
   uint8_t alignment = 1;
-  uintptr_t address = reinterpret_cast<uintptr_t>(gpu_ptr<void>(x));
+  uintptr_t address = reinterpret_cast<uintptr_t>(x.data<void>());
   for (; alignment < 32; alignment *= 2) {
     if (address % (alignment * 2)) {
       return alignment;
@@ -57,7 +56,7 @@ inline std::array<T, MAX_NDIM> vector_key(const Vec<T>& vec) {
 
 // Helpers used by get_data_ptrs to get pointers.
 inline void* get_data_ptr(const array& arr) {
-  return const_cast<void*>(gpu_ptr<void>(arr));
+  return const_cast<void*>(arr.data<void>());
 }
 
 template <typename T, typename = std::enable_if_t<std::is_scalar_v<T>>>

mlx/backend/cuda/custom_kernel.cpp

@@ -279,7 +279,6 @@ void CustomKernel::eval_gpu(
     std::vector<array>& outputs) {
   nvtx3::scoped_range r("CustomKernel::eval_gpu");
   auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
 
   std::vector<array> copies;
 
@@ -289,7 +288,7 @@ void CustomKernel::eval_gpu(
       copies.emplace_back(init_value_.value(), out.dtype());
       fill_gpu(copies.back(), out, s);
     } else {
-      out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+      out.set_data(allocator::malloc(out.nbytes()));
     }
   }
 
@@ -357,6 +356,7 @@ void CustomKernel::eval_gpu(
   dim3 grid((gx + tx - 1) / tx, (gy + ty - 1) / ty, (gz + tz - 1) / tz);
 
   // Call the kernel
+  auto& encoder = cu::get_command_encoder(s);
   for (const auto& in : checked_inputs) {
     encoder.set_input_array(in);
   }

mlx/backend/cuda/device.h

@@ -3,7 +3,6 @@
 #pragma once
 
 #include "mlx/array.h"
-#include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/lru_cache.h"
 #include "mlx/backend/cuda/worker.h"
 #include "mlx/stream.h"

mlx/backend/cuda/distributed.cu

@@ -15,10 +15,8 @@ void AllReduce::eval_gpu(
   assert(inputs.size() == 1);
   assert(outputs.size() == 1);
 
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  auto set_input_output = [&](const array& in,
-                              array& out) -> std::pair<array, array> {
+  auto set_input_output =
+      [s = stream()](const array& in, array& out) -> std::pair<array, array> {
     if (!in.flags().row_contiguous) {
       copy_gpu(in, out, CopyType::General, s);
       return {out, out};
@@ -26,17 +24,19 @@ void AllReduce::eval_gpu(
       out.copy_shared_buffer(in);
       return {in, out};
     } else {
-      out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+      out.set_data(allocator::malloc(out.nbytes()));
       return {in, out};
     }
   };
 
   auto [input, output] = set_input_output(inputs[0], outputs[0]);
 
+  auto& encoder = cu::get_command_encoder(stream());
   encoder.set_input_array(input);
   encoder.set_output_array(output);
 
   auto capture = encoder.capture_context();
+  auto& s = stream();
 
   switch (reduce_type_) {
     case Sum:
@@ -53,69 +53,4 @@ void AllReduce::eval_gpu(
           "Only all reduce sum, max, and min are supported.");
   }
 }
-
-void AllGather::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-  assert(outputs.size() == 1);
-
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-
-  auto ensure_contiguous = [&s, &encoder](const array& x) {
-    if (x.flags().row_contiguous) {
-      return x;
-    } else {
-      array x_copy = contiguous_copy_gpu(x, s);
-      encoder.add_temporary(x_copy);
-      return x_copy;
-    }
-  };
-
-  auto input = ensure_contiguous(inputs[0]);
-  outputs[0].set_data(cu::malloc_async(outputs[0].nbytes(), encoder.stream()));
-
-  encoder.set_input_array(input);
-  encoder.set_output_array(outputs[0]);
-
-  auto capture = encoder.capture_context();
-  distributed::detail::all_gather(group(), input, outputs[0], s);
-}
-
-void ReduceScatter::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-  assert(outputs.size() == 1);
-
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-
-  auto ensure_contiguous = [&s, &encoder](const array& x) {
-    if (x.flags().row_contiguous) {
-      return x;
-    } else {
-      array x_copy = contiguous_copy_gpu(x, s);
-      encoder.add_temporary(x_copy);
-      return x_copy;
-    }
-  };
-
-  auto input = ensure_contiguous(inputs[0]);
-  outputs[0].set_data(cu::malloc_async(outputs[0].nbytes(), encoder.stream()));
-
-  encoder.set_input_array(input);
-  encoder.set_output_array(outputs[0]);
-
-  auto capture = encoder.capture_context();
-
-  switch (reduce_type_) {
-    case Sum:
-      distributed::detail::sum_scatter(group(), input, outputs[0], s);
-      break;
-    default:
-      throw std::runtime_error("Only sum scatter is supported. ");
-  }
-}
 } // namespace mlx::core::distributed

mlx/backend/cuda/fence.cpp

@@ -1,8 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/fence.h"
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/event.h"
 
 namespace mlx::core {
@@ -22,24 +20,8 @@ void Fence::wait(Stream s, const array&) {
   fence->event.wait(fence->count);
 }
 
-void Fence::update(Stream s, const array& a, bool cross_device) {
+void Fence::update(Stream s, const array&) {
   auto* fence = static_cast<FenceImpl*>(fence_.get());
-  if (cross_device) {
-    // Move to managed memory if there is a device switch
-    auto& cbuf =
-        *static_cast<cu::CudaBuffer*>(const_cast<array&>(a).buffer().ptr());
-    if (cbuf.device != -1) {
-      void* new_data;
-      CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, cbuf.size));
-      cbuf.device = -1;
-      auto& encoder = cu::device(s.device).get_command_encoder(s);
-      encoder.commit();
-      CHECK_CUDA_ERROR(cudaMemcpyAsync(
-          new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream()));
-      CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream()));
-      cbuf.data = new_data;
-    }
-  }
   fence->count++;
   fence->event.signal(s, fence->count);
 }

mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -241,7 +241,7 @@ void CublasGemm::set_bias(cu::CommandEncoder& encoder, const array& bias) {
       CUBLASLT_MATMUL_DESC_EPILOGUE,
       &epilogue,
       sizeof(epilogue)));
-  auto* bias_ptr = gpu_ptr<void>(bias);
+  auto* bias_ptr = bias.data<void>();
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
       matmul_desc_,
       CUBLASLT_MATMUL_DESC_BIAS_POINTER,
@@ -278,9 +278,9 @@ void CublasGemm::run(
 
   execute(
       encoder,
-      gpu_ptr<void>(out),
-      gpu_ptr<void>(a),
-      gpu_ptr<void>(b),
+      out.data<void>(),
+      a.data<void>(),
+      b.data<void>(),
       nullptr,
       alpha);
 }
@@ -321,10 +321,10 @@ void CublasGemm::run(
 
   execute(
       encoder,
-      gpu_ptr<void>(out),
-      gpu_ptr<void>(a),
-      gpu_ptr<void>(b),
-      gpu_ptr<void>(c),
+      out.data<void>(),
+      a.data<void>(),
+      b.data<void>(),
+      c.data<void>(),
       alpha,
       beta);
 }
@@ -370,11 +370,11 @@ void CublasGemm::execute(
     // Ensure workspace is 256-byte aligned
     int nbytes = cuda::ceil_div(heuristic_.workspaceSize, 256) * 256;
     array workspace(
-        cu::malloc_async(nbytes, encoder.stream()),
+        allocator::malloc(nbytes),
         {static_cast<int>(heuristic_.workspaceSize)},
         int8);
     encoder.add_temporary(workspace);
-    workspace_ptr = gpu_ptr<void>(workspace);
+    workspace_ptr = workspace.data<void>();
   }
 
   auto capture = encoder.capture_context();

mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp

@@ -25,10 +25,9 @@ void CublasGemm::run_batched(
   for (size_t i = 0; i < nbatch; ++i) {
     execute(
         encoder,
-        gpu_ptr<int8_t>(out) +
-            out.itemsize() * i * batch_shape.back() * M_ * N_,
-        gpu_ptr<int8_t>(a) + a.itemsize() * a_it.loc,
-        gpu_ptr<int8_t>(b) + b.itemsize() * b_it.loc,
+        out.data<int8_t>() + out.itemsize() * i * batch_shape.back() * M_ * N_,
+        a.data<int8_t>() + a.itemsize() * a_it.loc,
+        b.data<int8_t>() + b.itemsize() * b_it.loc,
         nullptr,
         alpha);
     a_it.step();
@@ -61,11 +60,10 @@ void CublasGemm::run_batched(
   for (size_t i = 0; i < nbatch; ++i) {
     execute(
         encoder,
-        gpu_ptr<int8_t>(out) +
-            out.itemsize() * i * batch_shape.back() * M_ * N_,
-        gpu_ptr<int8_t>(a) + a.itemsize() * a_it.loc,
-        gpu_ptr<int8_t>(b) + b.itemsize() * b_it.loc,
-        gpu_ptr<int8_t>(c) + c.itemsize() * c_it.loc,
+        out.data<int8_t>() + out.itemsize() * i * batch_shape.back() * M_ * N_,
+        a.data<int8_t>() + a.itemsize() * a_it.loc,
+        b.data<int8_t>() + b.itemsize() * b_it.loc,
+        c.data<int8_t>() + c.itemsize() * c_it.loc,
         alpha,
         beta);
     a_it.step();

mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu

@@ -163,7 +163,7 @@ void CublasGemm::run_batched(
 
   // Launch kernel to set device offsets
   auto pointers = array(
-      cu::malloc_async(batch_count * sizeof(void*) * 3, encoder.stream()),
+      allocator::malloc(batch_count * sizeof(void*) * 3),
       {batch_count * 3},
       uint64);
 
@@ -183,10 +183,10 @@ void CublasGemm::run_batched(
           num_blocks,
           block_dims,
           0,
-          gpu_ptr<int8_t*>(pointers),
-          gpu_ptr<int8_t>(a),
-          gpu_ptr<int8_t>(b),
-          gpu_ptr<int8_t>(out),
+          pointers.data<int8_t*>(),
+          a.data<int8_t>(),
+          b.data<int8_t>(),
+          out.data<int8_t>(),
           item_size,
           const_param<ndim_constant()>(batch_shape),
           const_param<ndim_constant()>(a_batch_strides),
@@ -200,10 +200,10 @@ void CublasGemm::run_batched(
         num_blocks,
         block_dims,
         0,
-        gpu_ptr<int8_t*>(pointers),
-        gpu_ptr<int8_t>(a),
-        gpu_ptr<int8_t>(b),
-        gpu_ptr<int8_t>(out),
+        pointers.data<int8_t*>(),
+        a.data<int8_t>(),
+        b.data<int8_t>(),
+        out.data<int8_t>(),
         item_size,
         const_param(batch_shape),
         const_param(a_batch_strides),
@@ -219,7 +219,7 @@ void CublasGemm::run_batched(
   encoder.set_input_array(b);
   encoder.set_output_array(out);
 
-  auto a_pointers = gpu_ptr<int8_t*>(pointers);
+  auto a_pointers = pointers.data<int8_t*>();
   auto b_pointers = a_pointers + batch_count;
   auto out_pointers = b_pointers + batch_count;
   execute(
@@ -251,7 +251,7 @@ void CublasGemm::run_batched(
 
   // Launch kernel to set device offsets
   auto pointers = array(
-      cu::malloc_async(batch_count * sizeof(uint64_t) * 4, encoder.stream()),
+      allocator::malloc(batch_count * sizeof(uint64_t) * 4),
       {batch_count * 4},
       uint64);
 
@@ -271,11 +271,11 @@ void CublasGemm::run_batched(
           num_blocks,
           block_dims,
           0,
-          gpu_ptr<int8_t*>(pointers),
-          gpu_ptr<int8_t>(a),
-          gpu_ptr<int8_t>(b),
-          gpu_ptr<int8_t>(c),
-          gpu_ptr<int8_t>(out),
+          pointers.data<int8_t*>(),
+          a.data<int8_t>(),
+          b.data<int8_t>(),
+          c.data<int8_t>(),
+          out.data<int8_t>(),
           item_size,
           const_param<ndim_constant()>(batch_shape),
           const_param<ndim_constant()>(a_batch_strides),
@@ -290,11 +290,11 @@ void CublasGemm::run_batched(
         num_blocks,
         block_dims,
         0,
-        gpu_ptr<int8_t*>(pointers),
-        gpu_ptr<int8_t>(a),
-        gpu_ptr<int8_t>(b),
-        gpu_ptr<int8_t>(c),
-        gpu_ptr<int8_t>(out),
+        pointers.data<int8_t*>(),
+        a.data<int8_t>(),
+        b.data<int8_t>(),
+        c.data<int8_t>(),
+        out.data<int8_t>(),
         item_size,
         const_param(batch_shape),
         const_param(a_batch_strides),
@@ -312,7 +312,7 @@ void CublasGemm::run_batched(
   encoder.set_input_array(c);
   encoder.set_output_array(out);
 
-  auto a_pointers = gpu_ptr<int8_t*>(pointers);
+  auto a_pointers = pointers.data<int8_t*>();
   auto b_pointers = a_pointers + batch_count;
   auto c_pointers = b_pointers + batch_count;
   auto out_pointers = c_pointers + batch_count;

mlx/backend/cuda/gemms/gemv.cu

@@ -149,13 +149,13 @@ void gemv(
     auto vec_strides = const_param(b_batch_strides);
 
     if (M == 1) {
-      mat = gpu_ptr<DataType>(b);
-      vec = gpu_ptr<DataType>(a);
+      mat = b.data<DataType>();
+      vec = a.data<DataType>();
       rows = N;
       std::swap(mat_strides, vec_strides);
     } else {
-      mat = gpu_ptr<DataType>(a);
-      vec = gpu_ptr<DataType>(b);
+      mat = a.data<DataType>();
+      vec = b.data<DataType>();
       rows = M;
     }
     uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
@@ -177,7 +177,7 @@ void gemv(
             0,
             mat,
             vec,
-            gpu_ptr<DataType>(out),
+            out.data<DataType>(),
             rows,
             cols);
       } else {
@@ -189,7 +189,7 @@ void gemv(
             0,
             mat,
             vec,
-            gpu_ptr<DataType>(out),
+            out.data<DataType>(),
             rows,
             cols,
             const_param(batch_shape),

mlx/backend/cuda/indexing.cpp

@@ -31,7 +31,7 @@ void append_indices_arg(
     int idx_ndim) {
   SmallVector<const void*> indices(nidx);
   for (int i = 0; i < nidx; ++i) {
-    indices[i] = gpu_ptr<void>(inputs[i + 1]);
+    indices[i] = inputs[i + 1].data<void>();
   }
   args.append(std::move(indices));
   SmallVector<int32_t> indices_shape(nidx * idx_ndim);
@@ -59,9 +59,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() > 0);
   const auto& src = inputs[0];
 
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
   if (out.size() == 0) {
     return;
   }
@@ -82,6 +80,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
       dtype_to_string(idx_dtype),
       nidx);
 
+  auto& s = stream();
   cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
     std::vector<std::string> kernel_names;
     for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
@@ -122,6 +121,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
       idx_ndim,
       large ? "int64_t" : "int32_t");
 
+  auto& encoder = cu::get_command_encoder(s);
   for (const auto& in : inputs) {
     encoder.set_input_array(in);
   }
@@ -239,9 +239,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   const auto& src = inputs[0];
   const auto& idx = inputs[1];
 
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
   if (out.size() == 0) {
     return;
   }
@@ -253,6 +251,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
       dtype_to_string(out.dtype()),
       dtype_to_string(idx.dtype()));
 
+  auto& s = stream();
   cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
     std::vector<std::string> kernel_names;
     for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
@@ -313,6 +312,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
       idx.flags().row_contiguous,
       large ? "int64_t" : "int32_t");
 
+  auto& encoder = cu::get_command_encoder(s);
   for (const auto& in : inputs) {
     encoder.set_input_array(in);
   }

mlx/backend/cuda/jit_module.h

@@ -31,7 +31,7 @@ struct KernelArgs {
   }
 
   void append(const array& a) {
-    append(reinterpret_cast<CUdeviceptr>(gpu_ptr<void>(a)));
+    append(reinterpret_cast<CUdeviceptr>(a.data<void>()));
   }
 
   template <typename T>

mlx/backend/cuda/kernel_utils.cuh

@@ -9,7 +9,6 @@
 #include <type_traits>
 
 #include "mlx/array.h"
-#include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device/utils.cuh"
 
 #include <cuda.h>

mlx/backend/cuda/layer_norm.cu

@@ -230,10 +230,9 @@ void LayerNorm::eval_gpu(
   nvtx3::scoped_range r("LayerNorm::eval_gpu");
   auto& s = stream();
   auto& out = outputs[0];
-  auto& encoder = cu::get_command_encoder(s);
 
   // Make sure that the last dimension is contiguous.
-  auto set_output = [&s, &out, &encoder](const array& x) {
+  auto set_output = [&s, &out](const array& x) {
     bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
     if (no_copy && x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
@@ -244,7 +243,7 @@ void LayerNorm::eval_gpu(
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            cu::malloc_async(x.data_size() * x.itemsize(), encoder.stream()),
+            allocator::malloc(x.data_size() * x.itemsize()),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -266,6 +265,7 @@ void LayerNorm::eval_gpu(
   int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
   int64_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0;
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(x);
   encoder.set_input_array(w);
   encoder.set_input_array(b);
@@ -280,10 +280,10 @@ void LayerNorm::eval_gpu(
           n_rows,
           block_dim(),
           0,
-          gpu_ptr<DataType>(x),
-          gpu_ptr<DataType>(w),
-          gpu_ptr<DataType>(b),
-          gpu_ptr<DataType>(out),
+          x.data<DataType>(),
+          w.data<DataType>(),
+          b.data<DataType>(),
+          out.data<DataType>(),
           eps_,
           axis_size,
           w_stride,
@@ -335,7 +335,7 @@ void LayerNormVJP::eval_gpu(
     gx.copy_shared_buffer(g);
     g_in_gx = true;
   } else {
-    gx.set_data(cu::malloc_async(gx.nbytes(), encoder.stream()));
+    gx.set_data(allocator::malloc(gx.nbytes()));
   }
   if (g_copied && !g_in_gx) {
     encoder.add_temporary(g);
@@ -355,7 +355,7 @@ void LayerNormVJP::eval_gpu(
       g_in_gw = true;
       gw_temp.copy_shared_buffer(g);
     } else {
-      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder.stream()));
+      gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
       encoder.add_temporary(gw_temp);
     }
   }
@@ -393,11 +393,11 @@ void LayerNormVJP::eval_gpu(
                 n_rows,
                 block_dim(),
                 0,
-                gpu_ptr<DataType>(x),
-                gpu_ptr<DataType>(w),
-                gpu_ptr<DataType>(g),
-                gpu_ptr<DataType>(gx),
-                gpu_ptr<DataType>(gw_temp),
+                x.data<DataType>(),
+                w.data<DataType>(),
+                g.data<DataType>(),
+                gx.data<DataType>(),
+                gw_temp.data<DataType>(),
                 eps_,
                 axis_size,
                 w_stride);

mlx/backend/cuda/load.cpp

@@ -1,60 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-#include <algorithm>
-#include <utility>
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/primitives.h"
-
-namespace {
-
-template <const uint8_t scalar_size>
-void swap_endianness(uint8_t* data_bytes, size_t N) {
-  struct Elem {
-    uint8_t bytes[scalar_size];
-  };
-
-  Elem* data = reinterpret_cast<Elem*>(data_bytes);
-
-  for (size_t i = 0; i < N; i++) {
-    for (size_t j = 0; j < (scalar_size / 2); j++) {
-      std::swap(data[i].bytes[j], data[i].bytes[scalar_size - j - 1]);
-    }
-  }
-}
-
-} // namespace
-
-namespace mlx::core {
-
-void Load::eval_gpu(const std::vector<array>& inputs, array& out) {
-  auto& encoder = cu::get_command_encoder(stream());
-  auto size = out.size();
-  auto nbytes = size * out.itemsize();
-  out.set_data(cu::malloc_async(nbytes, encoder.stream()));
-  auto out_ptr = malloc(nbytes);
-  reader_->read(static_cast<char*>(out_ptr), nbytes, offset_);
-  if (swap_endianness_) {
-    switch (out.itemsize()) {
-      case 2:
-        swap_endianness<2>(reinterpret_cast<uint8_t*>(out_ptr), size);
-        break;
-      case 4:
-        swap_endianness<4>(reinterpret_cast<uint8_t*>(out_ptr), size);
-        break;
-      case 8:
-        swap_endianness<8>(reinterpret_cast<uint8_t*>(out_ptr), size);
-        break;
-    }
-  }
-  CHECK_CUDA_ERROR(cudaMemcpyAsync(
-      gpu_ptr<void>(out),
-      out_ptr,
-      nbytes,
-      cudaMemcpyDefault,
-      encoder.stream()));
-  CHECK_CUDA_ERROR(cudaLaunchHostFunc(encoder.stream(), free, out_ptr));
-}
-
-} // namespace mlx::core

mlx/backend/cuda/logsumexp.cu

@@ -115,7 +115,7 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   auto in = ensure_contiguous(inputs[0]);
   if (in.flags().row_contiguous) {
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+    out.set_data(allocator::malloc(out.nbytes()));
   } else {
     auto n = in.shape(-1);
     auto flags = in.flags();
@@ -130,7 +130,7 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
     }
     flags.col_contiguous = col_contig;
     out.set_data(
-        cu::malloc_async(in.nbytes() / n, encoder.stream()),
+        allocator::malloc(in.nbytes() / n),
         in.data_size() / n,
         std::move(strides),
         flags);
@@ -151,8 +151,8 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
           n_rows,
           block_dim(),
           0,
-          gpu_ptr<DataType>(in),
-          gpu_ptr<DataType>(out),
+          in.data<DataType>(),
+          out.data<DataType>(),
           axis_size);
     });
   });

mlx/backend/cuda/matmul.cpp

@@ -121,7 +121,7 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
 
   int M = a_pre.shape(-2);
   int N = b_pre.shape(-1);
@@ -163,7 +163,7 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   if (beta_ == 1 && a.dtype() != complex64 && c.strides(-1) == 1 &&
       c.data_size() == out.shape(-1)) {
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+    out.set_data(allocator::malloc(out.nbytes()));
     gemm_and_bias(
         encoder,
         M,
@@ -187,10 +187,10 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
     auto sty = c.strides()[c.ndim() - 1];
     if (sty == 1 && stx == c.shape(-1)) {
       ldc = stx;
-      out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+      out.set_data(allocator::malloc(out.nbytes()));
     } else if (sty == 1 && stx == 0) {
       ldc = 0;
-      out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+      out.set_data(allocator::malloc(out.nbytes()));
     } else {
       // Copy C into out and set C to out
       ldc = c.shape(-1);

mlx/backend/cuda/primitives.cpp

@@ -28,6 +28,7 @@ NO_GPU(FFT)
 NO_GPU(GatherMM)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
+NO_GPU(Load)
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
 NO_GPU(QuantizedMatmul)
@@ -39,6 +40,7 @@ NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
 
 namespace distributed {
+NO_GPU_MULTI(AllGather)
 NO_GPU_MULTI(Send)
 NO_GPU_MULTI(Recv)
 } // namespace distributed

mlx/backend/cuda/quantized/affine_quantize.cu

@@ -262,10 +262,10 @@ void affine_quantize(
             num_blocks,
             block_dims,
             0,
-            gpu_ptr<T>(w),
-            gpu_ptr<uint8_t>(wq),
-            gpu_ptr<T>(scales),
-            gpu_ptr<T>(biases),
+            w.data<T>(),
+            wq.data<uint8_t>(),
+            scales.data<T>(),
+            biases.data<T>(),
             w.size());
       });
     });
@@ -318,10 +318,10 @@ void affine_dequantize(
             num_blocks,
             block_dims,
             0,
-            gpu_ptr<uint8_t>(wq),
-            gpu_ptr<T>(scales),
-            gpu_ptr<T>(biases),
-            gpu_ptr<T>(w),
+            wq.data<uint8_t>(),
+            scales.data<T>(),
+            biases.data<T>(),
+            w.data<T>(),
             w.size());
       });
     });

mlx/backend/cuda/quantized/cuda_fp4.h

@@ -1,33 +1,5 @@
 #pragma once
 
-struct __nv_fp8_e8m0 {
-  __device__ __nv_fp8_e8m0(float x) {
-    if (!std::isfinite(x)) {
-      __x = 0xFF;
-      return;
-    }
-    if (x < 0.0f) {
-      __x = 0x00;
-      return;
-    }
-    float le = std::log2f(x);
-    int n = static_cast<int>(std::nearbyintf(le));
-
-    n = n < -127 ? -127 : n;
-    n = n > 127 ? 127 : n;
-    __x = static_cast<uint8_t>(n + 127);
-  }
-
-  __device__ operator float() {
-    if (__x == 0xFF) {
-      return std::numeric_limits<float>::quiet_NaN();
-    }
-    return std::ldexp(1.0f, static_cast<int>(__x) - 127);
-  }
-
-  uint8_t __x{0};
-};
-
 struct __nv_fp4_e2m1 {
   __device__ __nv_fp4_e2m1(float x) {
     if (std::isnan(x)) {

mlx/backend/cuda/quantized/fp_quantize.cu

@@ -49,12 +49,13 @@ fp_quantize(const T* w, uint8_t* out, uint8_t* scales, size_t size) {
 
   auto grid_dim_x =
       cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
-  size_t index = tidx + grid_dim_x * size_t(tidy);
-  if (index >= size) {
+  size_t out_index = tidx + grid_dim_x * size_t(tidy);
+  size_t in_index = out_index;
+  if (in_index >= size) {
     return;
   }
 
-  float w_thread = w[index];
+  float w_thread = w[in_index];
 
   cg::greater<float> max_op;
   auto warp = cg::tiled_partition<group_size>(cg::this_thread_block());
@@ -69,19 +70,21 @@ fp_quantize(const T* w, uint8_t* out, uint8_t* scales, size_t size) {
   scale = float(s);
 
   // Write out the scales
-  size_t gindex = index / group_size;
-  if (index % group_size == 0) {
+  size_t gindex = in_index / group_size;
+  if (in_index % group_size == 0) {
     scales[gindex] = q_scale;
   }
 
-  uint8_t output = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
+  uint8_t output = 0;
+  uint8_t val = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
+  output = val;
   if (bits == 4) {
-    uint8_t sval = warp.shfl_down(output, 1);
+    uint8_t sval = warp.shfl_down(val, 1);
     output |= sval << bits;
   }
   constexpr int pack_factor = bits == 8 ? 1 : 2;
-  if (index % pack_factor == 0) {
-    out[index / pack_factor] = output;
+  if (out_index % pack_factor == 0) {
+    out[out_index / pack_factor] = output;
   }
 }
 
@@ -156,9 +159,9 @@ void fp_quantize(
           num_blocks,
           block_dims,
           0,
-          gpu_ptr<T>(w),
-          gpu_ptr<uint8_t>(wq),
-          gpu_ptr<uint8_t>(scales),
+          w.data<T>(),
+          wq.data<uint8_t>(),
+          scales.data<uint8_t>(),
           w.size());
     } else {
       throw std::runtime_error(
@@ -202,9 +205,9 @@ void fp_dequantize(
           num_blocks,
           block_dims,
           0,
-          gpu_ptr<uint8_t>(wq),
-          gpu_ptr<uint8_t>(scales),
-          gpu_ptr<T>(w),
+          wq.data<uint8_t>(),
+          scales.data<T>(),
+          w.data<T>(),
           w.size());
     } else {
       throw std::runtime_error(

mlx/backend/cuda/quantized/quantized.cpp

@@ -59,7 +59,7 @@ void fast::Quantize::eval_gpu(
     auto scales = ensure_row_contiguous(inputs[1], enc, s);
     auto& w = outputs[0];
 
-    w.set_data(cu::malloc_async(w.nbytes(), enc.stream()));
+    w.set_data(allocator::malloc(w.nbytes()));
 
     if (mode_ == QuantizationMode::Affine) {
       auto biases = ensure_row_contiguous(inputs[2], enc, s);
@@ -72,11 +72,11 @@ void fast::Quantize::eval_gpu(
     auto& wq = outputs[0];
     auto& scales = outputs[1];
 
-    wq.set_data(cu::malloc_async(wq.nbytes(), enc.stream()));
-    scales.set_data(cu::malloc_async(scales.nbytes(), enc.stream()));
+    wq.set_data(allocator::malloc(wq.nbytes()));
+    scales.set_data(allocator::malloc(scales.nbytes()));
     if (mode_ == QuantizationMode::Affine) {
       auto& biases = outputs[2];
-      biases.set_data(cu::malloc_async(biases.nbytes(), enc.stream()));
+      biases.set_data(allocator::malloc(biases.nbytes()));
       affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s);
     } else {
       fp_quantize(w, wq, scales, group_size_, bits_, enc, s);

mlx/backend/cuda/random.cu

@@ -143,9 +143,7 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   uint32_t elems_per_key = out.size() / num_keys;
   uint32_t bytes_per_key = out.itemsize() * elems_per_key;
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
   if (out.size() == 0) {
     return;
   }
@@ -154,6 +152,8 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
   uint32_t half_size = out_per_key / 2;
   bool odd = out_per_key % 2;
 
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(keys);
   encoder.set_output_array(out);
   dim3 grid_dims{num_keys, half_size + odd};
@@ -171,8 +171,8 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
         grid,
         block,
         0,
-        gpu_ptr<uint32_t>(keys),
-        gpu_ptr<uint8_t>(out),
+        keys.data<uint32_t>(),
+        out.data<uint8_t>(),
         grid_dims,
         odd,
         bytes_per_key);
@@ -182,8 +182,8 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
         grid,
         block,
         0,
-        gpu_ptr<uint32_t>(keys),
-        gpu_ptr<uint8_t>(out),
+        keys.data<uint32_t>(),
+        out.data<uint8_t>(),
         grid_dims,
         odd,
         bytes_per_key,

mlx/backend/cuda/reduce/all_reduce.cu

@@ -66,7 +66,7 @@ void all_reduce(
     Reduce::ReduceType reduce_type) {
   constexpr int N_READS = 8;
 
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
 
   auto get_args = [](size_t size, int N) {
     int threads = std::min(512UL, (size + N - 1) / N);
@@ -100,15 +100,14 @@ void all_reduce(
   Dtype dt = in.dtype();
 
   // Cub doesn't like const pointers for load (sigh).
-  void* indata = const_cast<void*>(gpu_ptr<void>(in));
+  void* indata = const_cast<void*>(in.data<void>());
 
   // Large array so allocate an intermediate and accumulate there
   std::tie(blocks, threads, block_step) = get_args(insize, N_READS);
   encoder.set_input_array(in);
   if (blocks > 1) {
     array intermediate({blocks}, out.dtype(), nullptr, {});
-    intermediate.set_data(
-        cu::malloc_async(intermediate.nbytes(), encoder.stream()));
+    intermediate.set_data(allocator::malloc(intermediate.nbytes()));
     encoder.add_temporary(intermediate);
     encoder.set_output_array(intermediate);
     dispatch_all_types(dt, [&](auto type_tag) {
@@ -123,14 +122,14 @@ void all_reduce(
             threads,
             0,
             static_cast<T*>(indata),
-            gpu_ptr<U>(intermediate),
+            intermediate.data<U>(),
             block_step,
             insize);
       });
     });
 
     // Set the input for the next step and recalculate the blocks
-    indata = gpu_ptr<void>(intermediate);
+    indata = intermediate.data<void>();
     dt = intermediate.dtype();
     insize = intermediate.size();
     std::tie(blocks, threads, block_step) = get_args(insize, N_READS);
@@ -150,7 +149,7 @@ void all_reduce(
           threads,
           0,
           static_cast<T*>(indata),
-          gpu_ptr<U>(out),
+          out.data<U>(),
           block_step,
           insize);
     });

mlx/backend/cuda/reduce/col_reduce.cu

@@ -250,7 +250,7 @@ void col_reduce_looped(
     const cu::ColReduceArgs& args) {
   // Allocate data for the output using in's layout to access them as
   // contiguously as possible.
-  allocate_same_layout(out, in, axes, encoder);
+  allocate_same_layout(out, in, axes);
 
   encoder.set_input_array(in);
   encoder.set_output_array(out);
@@ -261,7 +261,7 @@ void col_reduce_looped(
         using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
         using U = typename cu::ReduceResult<OP, T>::type;
         // Cub doesn't like const pointers for vectorized loads. (sigh)
-        T* indata = const_cast<T*>(gpu_ptr<T>(in));
+        T* indata = const_cast<T*>(in.data<T>());
 
         constexpr int N_READS = 4;
         constexpr int BM = 32;
@@ -276,7 +276,7 @@ void col_reduce_looped(
             blocks,
             0,
             indata,
-            gpu_ptr<U>(out),
+            out.data<U>(),
             static_cast<cu::ColReduceArgs>(args));
       });
     });
@@ -293,7 +293,7 @@ void col_reduce_small(
     const cu::ColReduceArgs& args) {
   // Allocate data for the output using in's layout to access them as
   // contiguously as possible.
-  allocate_same_layout(out, in, axes, encoder);
+  allocate_same_layout(out, in, axes);
 
   encoder.set_input_array(in);
   encoder.set_output_array(out);
@@ -312,8 +312,8 @@ void col_reduce_small(
           grid,
           block,
           0,
-          gpu_ptr<T>(in),
-          gpu_ptr<U>(out),
+          in.data<T>(),
+          out.data<U>(),
           static_cast<cu::ColReduceArgs>(args),
           out.size());
     });

mlx/backend/cuda/reduce/init_reduce.cu

@@ -28,7 +28,7 @@ void init_reduce(
     Reduce::ReduceType reduce_type) {
   // Allocate if needed
   if (out.data_shared_ptr() == nullptr) {
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+    out.set_data(allocator::malloc(out.nbytes()));
   }
 
   encoder.set_output_array(out);
@@ -42,7 +42,7 @@ void init_reduce(
       dim3 block(grid.x < 1024 ? grid.x : 1024, 1, 1);
       grid.x = (grid.x + 1023) / 1024;
       encoder.add_kernel_node(
-          kernel, grid, block, 0, gpu_ptr<U>(out), out.size());
+          kernel, grid, block, 0, out.data<U>(), out.size());
     });
   });
 }

mlx/backend/cuda/reduce/reduce_utils.cuh

@@ -5,7 +5,6 @@
 #include <numeric>
 
 #include "mlx/backend/common/utils.h"
-#include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/device/utils.cuh"
 
 #include <cooperative_groups.h>
@@ -93,10 +92,9 @@ block_reduce(Block block, Warp warp, T (&vals)[N], T* smem, Op op, T init) {
 inline void allocate_same_layout(
     array& out,
     const array& in,
-    const std::vector<int>& axes,
-    cu::CommandEncoder& encoder) {
+    const std::vector<int>& axes) {
   if (in.flags().row_contiguous) {
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+    out.set_data(allocator::malloc(out.nbytes()));
     return;
   }
 
@@ -135,7 +133,7 @@ inline void allocate_same_layout(
   fl.col_contiguous = cc;
   fl.contiguous = true;
   out.set_data(
-      cu::malloc_async(out.nbytes(), encoder.stream()),
+      allocator::malloc(out.nbytes()),
       data_size,
       final_strides,
       fl,

mlx/backend/cuda/reduce/row_reduce.cu

@@ -238,7 +238,7 @@ void row_reduce_simple(
     const ReductionPlan& plan) {
   // Allocate data for the output using in's layout to avoid elem_to_loc in the
   // kernel.
-  allocate_same_layout(out, in, axes, encoder);
+  allocate_same_layout(out, in, axes);
 
   // TODO: If out.size() < 1024 which will be a common case then write this in
   //       2 passes. Something like 32 * out.size() and then do a warp reduce.
@@ -268,10 +268,10 @@ void row_reduce_simple(
         kernel = cu::row_reduce_simple<T, U, OP, N_READS, 2>;
       }
 
-      T* indata = const_cast<T*>(gpu_ptr<T>(in));
+      T* indata = const_cast<T*>(in.data<T>());
       int size = plan.shape.back();
       encoder.add_kernel_node(
-          kernel, grid, block, 0, indata, gpu_ptr<U>(out), out.size(), size);
+          kernel, grid, block, 0, indata, out.data<U>(), out.size(), size);
     });
   });
 }
@@ -286,7 +286,7 @@ void row_reduce_looped(
     cu::RowReduceArgs args) {
   // Allocate data for the output using in's layout to access them as
   // contiguously as possible.
-  allocate_same_layout(out, in, axes, encoder);
+  allocate_same_layout(out, in, axes);
 
   encoder.set_input_array(in);
   encoder.set_output_array(out);
@@ -315,7 +315,7 @@ void row_reduce_looped(
       });
 
       encoder.add_kernel_node(
-          kernel, grid, block, 0, gpu_ptr<T>(in), gpu_ptr<U>(out), args);
+          kernel, grid, block, 0, in.data<T>(), out.data<U>(), args);
     });
   });
 }

mlx/backend/cuda/rms_norm.cu

@@ -176,10 +176,9 @@ void RMSNorm::eval_gpu(
   nvtx3::scoped_range r("RMSNorm::eval_gpu");
   auto& s = stream();
   auto& out = outputs[0];
-  auto& encoder = cu::get_command_encoder(s);
 
   // Make sure that the last dimension is contiguous.
-  auto set_output = [&s, &out, &encoder](const array& x) {
+  auto set_output = [&s, &out](const array& x) {
     bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
     if (no_copy && x.ndim() > 1) {
       auto s = x.strides()[x.ndim() - 2];
@@ -190,7 +189,7 @@ void RMSNorm::eval_gpu(
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            cu::malloc_async(x.data_size() * x.itemsize(), encoder.stream()),
+            allocator::malloc(x.data_size() * x.itemsize()),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -210,6 +209,7 @@ void RMSNorm::eval_gpu(
   int32_t n_rows = x.data_size() / axis_size;
   int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(x);
   encoder.set_input_array(w);
   encoder.set_output_array(out);
@@ -223,9 +223,9 @@ void RMSNorm::eval_gpu(
           n_rows,
           block_dim(),
           0,
-          gpu_ptr<DataType>(x),
-          gpu_ptr<DataType>(w),
-          gpu_ptr<DataType>(out),
+          x.data<DataType>(),
+          w.data<DataType>(),
+          out.data<DataType>(),
           eps_,
           axis_size,
           w_stride);
@@ -274,7 +274,7 @@ void RMSNormVJP::eval_gpu(
     gx.copy_shared_buffer(g);
     g_in_gx = true;
   } else {
-    gx.set_data(cu::malloc_async(gx.nbytes(), encoder.stream()));
+    gx.set_data(allocator::malloc(gx.nbytes()));
   }
   if (g_copied && !g_in_gx) {
     encoder.add_temporary(g);
@@ -292,7 +292,7 @@ void RMSNormVJP::eval_gpu(
     if (!g_in_gx && donate_g) {
       gw_temp.copy_shared_buffer(g);
     } else {
-      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder.stream()));
+      gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
       encoder.add_temporary(gw_temp);
     }
   }
@@ -318,11 +318,11 @@ void RMSNormVJP::eval_gpu(
                 n_rows,
                 block_dim(),
                 0,
-                gpu_ptr<DataType>(x),
-                gpu_ptr<DataType>(w),
-                gpu_ptr<DataType>(g),
-                gpu_ptr<DataType>(gx),
-                gpu_ptr<DataType>(gw_temp),
+                x.data<DataType>(),
+                w.data<DataType>(),
+                g.data<DataType>(),
+                gx.data<DataType>(),
+                gw_temp.data<DataType>(),
                 eps_,
                 axis_size,
                 w_stride);

mlx/backend/cuda/rope.cu

@@ -250,7 +250,6 @@ void RoPE::eval_gpu(
   nvtx3::scoped_range r("RoPE::eval_gpu");
 
   auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
   auto& in = inputs[0];
   auto& offset = inputs[1];
   auto& out = outputs[0];
@@ -292,14 +291,14 @@ void RoPE::eval_gpu(
       donated = true;
       out.copy_shared_buffer(in);
     } else {
-      out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+      out.set_data(allocator::malloc(out.nbytes()));
     }
     strides[0] = mat_size;
     strides[1] = in.strides()[ndim - 2];
     strides[2] = in.strides()[ndim - 1];
   } else if (dispatch_ndim == 3) {
     // Handle non-contiguous 3D inputs
-    out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+    out.set_data(allocator::malloc(out.nbytes()));
     strides[0] = in.strides()[ndim - 3];
     strides[1] = in.strides()[ndim - 2];
     strides[2] = in.strides()[ndim - 1];
@@ -320,6 +319,7 @@ void RoPE::eval_gpu(
   bool single = in.flags().row_contiguous && B == 1 && T == 1;
   bool with_freqs = inputs.size() == 3;
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(donated ? out : in);
   encoder.set_input_array(offset);
   if (with_freqs) {
@@ -340,9 +340,9 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              gpu_ptr<DataType>(donated ? out : in),
-              gpu_ptr<DataType>(out),
-              gpu_ptr<int32_t>(offset),
+              (donated ? out : in).data<DataType>(),
+              out.data<DataType>(),
+              offset.data<int32_t>(),
               scale_,
               std::log2(base_),
               mat_size,
@@ -357,10 +357,10 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              gpu_ptr<DataType>(donated ? out : in),
-              gpu_ptr<DataType>(out),
-              gpu_ptr<int32_t>(offset),
-              gpu_ptr<float>(inputs[2]),
+              (donated ? out : in).data<DataType>(),
+              out.data<DataType>(),
+              offset.data<int32_t>(),
+              inputs[2].data<float>(),
               scale_,
               mat_size,
               dims,
@@ -381,10 +381,10 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              gpu_ptr<DataType>(donated ? out : in),
-              gpu_ptr<DataType>(out),
-              gpu_ptr<int32_t>(offset),
-              gpu_ptr<float>(inputs[2]),
+              (donated ? out : in).data<DataType>(),
+              out.data<DataType>(),
+              offset.data<int32_t>(),
+              inputs[2].data<float>(),
               scale_,
               std::log2(base_),
               strides,
@@ -408,9 +408,9 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              gpu_ptr<DataType>(donated ? out : in),
-              gpu_ptr<DataType>(out),
-              gpu_ptr<int32_t>(offset),
+              (donated ? out : in).data<DataType>(),
+              out.data<DataType>(),
+              offset.data<int32_t>(),
               scale_,
               std::log2(base_),
               strides,

mlx/backend/cuda/scaled_dot_product_attention.cu

@@ -513,11 +513,11 @@ void sdpa_vector_1pass_fallback(
             grid_dim,
             block_dim,
             0,
-            gpu_ptr<DataType>(q),
-            gpu_ptr<DataType>(k),
-            gpu_ptr<DataType>(v),
-            gpu_ptr<DataType>(o),
-            sinks ? gpu_ptr<DataType>(*sinks) : nullptr,
+            q.data<DataType>(),
+            k.data<DataType>(),
+            v.data<DataType>(),
+            o.data<DataType>(),
+            sinks ? (*sinks).data<DataType>() : nullptr,
             params);
       });
     });
@@ -565,10 +565,9 @@ void sdpa_vector_2pass_fallback(
   array sums(intermediate_shape, float32, nullptr, {});
   array maxs(std::move(intermediate_shape), float32, nullptr, {});
 
-  intermediate.set_data(
-      cu::malloc_async(intermediate.nbytes(), encoder.stream()));
-  sums.set_data(cu::malloc_async(sums.nbytes(), encoder.stream()));
-  maxs.set_data(cu::malloc_async(maxs.nbytes(), encoder.stream()));
+  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+  sums.set_data(allocator::malloc(sums.nbytes()));
+  maxs.set_data(allocator::malloc(maxs.nbytes()));
 
   encoder.add_temporary(intermediate);
   encoder.add_temporary(sums);
@@ -602,13 +601,13 @@ void sdpa_vector_2pass_fallback(
               grid_dim,
               block_dim,
               0,
-              gpu_ptr<DataType>(q),
-              gpu_ptr<DataType>(k),
-              gpu_ptr<DataType>(v),
-              sinks ? gpu_ptr<DataType>(*sinks) : nullptr,
-              gpu_ptr<float>(intermediate),
-              gpu_ptr<float>(sums),
-              gpu_ptr<float>(maxs),
+              q.data<DataType>(),
+              k.data<DataType>(),
+              v.data<DataType>(),
+              sinks ? (*sinks).data<DataType>() : nullptr,
+              intermediate.data<float>(),
+              sums.data<float>(),
+              maxs.data<float>(),
               params);
         }
 
@@ -629,10 +628,10 @@ void sdpa_vector_2pass_fallback(
               grid_dim,
               block_dim,
               0,
-              gpu_ptr<float>(intermediate),
-              gpu_ptr<float>(sums),
-              gpu_ptr<float>(maxs),
-              gpu_ptr<DataType>(o),
+              intermediate.data<float>(),
+              sums.data<float>(),
+              maxs.data<float>(),
+              o.data<DataType>(),
               params);
         }
       });
@@ -788,7 +787,7 @@ void ScaledDotProductAttention::eval_gpu(
       };
 
       o.set_data(
-          cu::malloc_async(o.nbytes(), encoder.stream()),
+          allocator::malloc(o.nbytes()),
           o.size(),
           {str_oB, str_oH, str_oL, str_oD},
           flags);

mlx/backend/cuda/scan.cu

@@ -367,14 +367,13 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto in = inputs[0];
   auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
 
   if (in.flags().contiguous && in.strides()[axis_] != 0) {
     if (in.is_donatable() && in.itemsize() == out.itemsize()) {
       out.copy_shared_buffer(in);
     } else {
       out.set_data(
-          cu::malloc_async(in.data_size() * out.itemsize(), encoder.stream()),
+          allocator::malloc(in.data_size() * out.itemsize()),
           in.data_size(),
           in.strides(),
           in.flags());
@@ -388,6 +387,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   int32_t axis_size = in.shape(axis_);
   bool contiguous = in.strides()[axis_] == 1;
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(in);
   encoder.set_output_array(out);
 
@@ -415,8 +415,8 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
                   in.data_size() / axis_size,
                   block_dim,
                   0,
-                  gpu_ptr<T>(in),
-                  gpu_ptr<U>(out),
+                  in.data<T>(),
+                  out.data<U>(),
                   axis_size);
             } else {
               constexpr int BM = WARP_SIZE;
@@ -445,8 +445,8 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
                   num_blocks,
                   block_dim,
                   0,
-                  gpu_ptr<T>(in),
-                  gpu_ptr<U>(out),
+                  in.data<T>(),
+                  out.data<U>(),
                   axis_size,
                   stride,
                   stride_blocks);

mlx/backend/cuda/slicing.cpp

@@ -23,15 +23,14 @@ void concatenate_gpu(
   }
   std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
 
-  auto& encoder = cu::get_command_encoder(s);
-  out.set_data(cu::malloc_async(out.nbytes(), encoder.stream()));
+  out.set_data(allocator::malloc(out.nbytes()));
 
   auto strides = out.strides();
   auto flags = out.flags();
   flags.row_contiguous = false;
   flags.col_contiguous = false;
   flags.contiguous = false;
-  auto concurrent = encoder.concurrent_context();
+  auto concurrent = cu::get_command_encoder(s).concurrent_context();
   for (int i = 0; i < inputs.size(); i++) {
     array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
     size_t data_offset = strides[axis] * sizes[i];
@@ -81,7 +80,6 @@ array compute_dynamic_offset(
     return std::make_tuple(false, std::move(source), std::vector{kernel_name});
   });
 
-  auto& encoder = cu::get_command_encoder(s);
   // Prepare output.
   array offset({1}, int64, nullptr, {});
   bool donate = indices.is_donatable() &&
@@ -89,9 +87,10 @@ array compute_dynamic_offset(
   if (donate) {
     offset.copy_shared_buffer(indices);
   } else {
-    offset.set_data(cu::malloc_async(offset.itemsize(), encoder.stream()));
+    offset.set_data(allocator::malloc(offset.itemsize()));
   }
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.add_temporary(offset);
   encoder.set_input_array(indices);
   encoder.set_output_array(offset);

mlx/backend/cuda/softmax.cu

@@ -109,16 +109,15 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
   nvtx3::scoped_range r("Softmax::eval_gpu");
   assert(inputs.size() == 1);
   auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
 
   // Make sure that the last dimension is contiguous.
-  auto set_output = [&s, &out, &encoder](const array& x) {
+  auto set_output = [&s, &out](const array& x) {
     if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
       if (x.is_donatable()) {
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            cu::malloc_async(x.data_size() * x.itemsize(), encoder.stream()),
+            allocator::malloc(x.data_size() * x.itemsize()),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -137,6 +136,7 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
   int axis_size = in.shape().back();
   int n_rows = in.data_size() / axis_size;
 
+  auto& encoder = cu::get_command_encoder(s);
   encoder.set_input_array(in);
   encoder.set_output_array(out);
   dispatch_float_types(out.dtype(), "softmax", [&](auto type_tag) {
@@ -152,8 +152,8 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
           n_rows,
           block_dim(),
           0,
-          gpu_ptr<DataType>(in),
-          gpu_ptr<DataType>(out),
+          in.data<DataType>(),
+          out.data<DataType>(),
           axis_size);
     });
   });

mlx/backend/cuda/sort.cu

@@ -49,14 +49,11 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
     array trans = swapaxes_in_eval(in, axis, last_dim);
     in = contiguous_copy_gpu(trans, s);
     encoder.add_temporary(in);
-    out = array(
-        cu::malloc_async(out.nbytes(), encoder.stream()),
-        in.shape(),
-        out.dtype());
+    out = array(allocator::malloc(out.nbytes()), in.shape(), out.dtype());
     encoder.add_temporary(out);
   } else {
     out.set_data(
-        cu::malloc_async(in.data_size() * out.itemsize(), encoder.stream()),
+        allocator::malloc(in.data_size() * out.itemsize()),
         in.data_size(),
         in.strides(),
         in.flags());
@@ -73,28 +70,22 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
           thrust::make_counting_iterator(0), OffsetTransform{nsort});
       if (argsort) {
         // Indices in the sorted dimension.
-        array indices(
-            cu::malloc_async(out.nbytes(), encoder.stream()),
-            in.shape(),
-            out.dtype());
+        array indices(allocator::malloc(out.nbytes()), in.shape(), out.dtype());
         encoder.add_temporary(indices);
 
         // In argsort though we don't need the result of sorted values, the
         // API requires us to provide an array to store it.
-        array discard(
-            cu::malloc_async(in.nbytes(), encoder.stream()),
-            in.shape(),
-            in.dtype());
+        array discard(allocator::malloc(in.nbytes()), in.shape(), in.dtype());
         encoder.add_temporary(discard);
 
         size_t size;
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortPairs(
             nullptr,
             size,
-            gpu_ptr<Type>(in),
-            gpu_ptr<Type>(discard),
-            gpu_ptr<uint32_t>(indices),
-            gpu_ptr<uint32_t>(out),
+            in.data<Type>(),
+            discard.data<Type>(),
+            indices.data<uint32_t>(),
+            out.data<uint32_t>(),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -103,10 +94,7 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
             sizeof(Type) * 8,
             stream));
 
-        array temp(
-            cu::malloc_async(size, encoder.stream()),
-            {static_cast<int>(size)},
-            uint8);
+        array temp(allocator::malloc(size), {static_cast<int>(size)}, uint8);
         encoder.add_temporary(temp);
 
         // Start capturing after allocations
@@ -115,16 +103,16 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
             cu::thrust_policy(stream),
             thrust::counting_iterator<uint32_t>(0),
             thrust::counting_iterator<uint32_t>(indices.data_size()),
-            thrust::device_pointer_cast(gpu_ptr<uint32_t>(indices)),
+            thrust::device_pointer_cast(indices.data<uint32_t>()),
             ModOp<uint32_t>{static_cast<uint32_t>(nsort)});
 
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortPairs(
-            gpu_ptr<void>(temp),
+            temp.data<void>(),
             size,
-            gpu_ptr<Type>(in),
-            gpu_ptr<Type>(discard),
-            gpu_ptr<uint32_t>(indices),
-            gpu_ptr<uint32_t>(out),
+            in.data<Type>(),
+            discard.data<Type>(),
+            indices.data<uint32_t>(),
+            out.data<uint32_t>(),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -137,8 +125,8 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortKeys(
             nullptr,
             size,
-            gpu_ptr<Type>(in),
-            gpu_ptr<Type>(out),
+            in.data<Type>(),
+            out.data<Type>(),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -147,19 +135,16 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
             sizeof(Type) * 8,
             stream));
 
-        array temp(
-            cu::malloc_async(size, encoder.stream()),
-            {static_cast<int>(size)},
-            uint8);
+        array temp(allocator::malloc(size), {static_cast<int>(size)}, uint8);
         encoder.add_temporary(temp);
 
         // Start capturing after allocations
         auto capture = encoder.capture_context();
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortKeys(
-            gpu_ptr<void>(temp),
+            temp.data<void>(),
             size,
-            gpu_ptr<Type>(in),
-            gpu_ptr<Type>(out),
+            in.data<Type>(),
+            out.data<Type>(),
             in.data_size(),
             in.data_size() / nsort,
             offsets,

mlx/backend/cuda/ternary.cu

@@ -168,10 +168,10 @@ void ternary_op_gpu_inplace(
             num_blocks,
             block_dims,
             0,
-            gpu_ptr<bool>(a),
-            gpu_ptr<DType>(b),
-            gpu_ptr<DType>(c),
-            gpu_ptr<DType>(out),
+            a.data<bool>(),
+            b.data<DType>(),
+            c.data<DType>(),
+            out.data<DType>(),
             out.data_size());
       });
     } else {
@@ -211,10 +211,10 @@ void ternary_op_gpu_inplace(
                     {num_blocks_x, num_blocks_y},
                     block_dims,
                     0,
-                    gpu_ptr<bool>(a),
-                    gpu_ptr<DType>(b),
-                    gpu_ptr<DType>(c),
-                    gpu_ptr<DType>(out),
+                    a.data<bool>(),
+                    b.data<DType>(),
+                    c.data<DType>(),
+                    out.data<DType>(),
                     rest,
                     const_param<dims_constant()>(shape),
                     const_param<dims_constant()>(a_strides),
@@ -231,10 +231,10 @@ void ternary_op_gpu_inplace(
                   {num_blocks_x, num_blocks_y},
                   block_dims,
                   0,
-                  gpu_ptr<bool>(a),
-                  gpu_ptr<DType>(b),
-                  gpu_ptr<DType>(c),
-                  gpu_ptr<DType>(out),
+                  a.data<bool>(),
+                  b.data<DType>(),
+                  c.data<DType>(),
+                  out.data<DType>(),
                   rest,
                   const_param(shape),
                   const_param(a_strides),
@@ -256,10 +256,7 @@ void ternary_op_gpu(
   auto& b = inputs[1];
   auto& c = inputs[2];
   auto topt = get_ternary_op_type(a, b, c);
-  auto& encoder = cu::get_command_encoder(s);
-  set_ternary_op_output_data(a, b, c, out, topt, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
+  set_ternary_op_output_data(a, b, c, out, topt);
   ternary_op_gpu_inplace<Op>(inputs, out, s);
 }
 

mlx/backend/cuda/unary/unary.cuh

@@ -158,8 +158,8 @@ void unary_op_gpu_inplace(
                 num_blocks,
                 block_dims,
                 0,
-                gpu_ptr<InType>(in),
-                gpu_ptr<OutType>(out),
+                in.data<InType>(),
+                out.data<OutType>(),
                 out.data_size());
           } else {
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
@@ -182,8 +182,8 @@ void unary_op_gpu_inplace(
                 {num_blocks_x, num_blocks_y},
                 block_dims,
                 0,
-                gpu_ptr<InType>(in),
-                gpu_ptr<OutType>(out),
+                in.data<InType>(),
+                out.data<OutType>(),
                 rest,
                 const_param(shape),
                 const_param(strides),
@@ -207,10 +207,7 @@ void unary_op_gpu(
     array& out,
     const char* op,
     const Stream& s) {
-  auto& encoder = cu::get_command_encoder(s);
-  set_unary_output_data(inputs[0], out, [&](auto n) {
-    return cu::malloc_async(n, encoder.stream());
-  });
+  set_unary_output_data(inputs[0], out);
   unary_op_gpu_inplace<Op>(inputs, out, op, s);
 }