Add Neural Accelerator Support (#2772)

Co-authored-by: Awni Hannun <awni@apple.com>
Co-authored-by: Angelos Katharopoulos <katharas@gmail.com>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

.circleci/config.yml

@@ -1,579 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
-
-parameters:
-  nightly_build:
-    type: boolean
-    default: false
-  test_release:
-    type: boolean
-    default: false
-
-jobs:
-  build_documentation:
-    parameters:
-      upload-docs:
-        type: boolean
-        default: false
-    macos:
-      xcode: "26.0.0"
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            brew install python@3.10
-            brew install doxygen
-            python3.10 -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install -r docs/requirements.txt
-            pip install . -v
-      - when:
-          condition:
-            not: << parameters.upload-docs >>
-          steps:
-            - run:
-               name: Build documentation
-               command: |
-                 source env/bin/activate
-                 cd docs && doxygen && make html O=-W
-      - when:
-          condition: << parameters.upload-docs >>
-          steps:
-            - add_ssh_keys:
-                fingerprints:
-                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
-            - run:
-               name: Upload documentation
-               command: |
-                 source env/bin/activate
-                 git config user.email "mlx@group.apple.com"
-                 git config user.name "CircleCI Docs"
-                 git checkout gh-pages
-                 git rebase main
-                 cd docs
-                 git rm -rf build/html
-                 doxygen && make html O=-W
-                 git add -f build/html
-                 git commit -m "rebase"
-                 git push -f origin gh-pages
-
-  linux_build_and_test:
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Run style checks
-          command: |
-            pip install pre-commit
-            pre-commit run --all
-            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
-      - run:
-          name: Install dependencies
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            python -m unittest discover python/tests -v
-            mpirun --bind-to none -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-            if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build
-            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-            make -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests
-
-  mac_build_and_test:
-    parameters:
-      xcode_version:
-        type: string
-        default: "26.0.0"
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
-              brew install openmpi uv
-      - run:
-          name: Install Python package
-          command: |
-            uv venv --python 3.10
-            uv pip install \
-              nanobind==2.4.0 \
-              cmake \
-              numpy \
-              torch \
-              tensorflow \
-              unittest-xml-reporting
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -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
-      - run:
-          name: Build example extension
-          command: |
-            source .venv/bin/activate
-            cd examples/extensions
-            uv pip install -r requirements.txt
-            uv run --no-project setup.py build_ext --inplace
-            uv run --no-project python test.py
-      - store_test_results:
-          path: test-results
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
-      - run:
-          name: Run CPP tests
-          command: |
-            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
-      - run:
-          name: Build small binary
-          command: |
-            source .venv/bin/activate
-            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`
-      - run:
-          name: Run Python tests with JIT
-          command: |
-            CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-              uv pip install -e . -v
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
-              METAL_DEBUG_ERROR_MODE=0 \
-              uv run --no-project python -m xmlrunner discover \
-                -v python/tests \
-                -o test-results/gpu_jit
-
-  cuda_build_and_test:
-    parameters:
-      image_date:
-        type: string
-        default: "2023.11.1"
-    machine:
-      image: "linux-cuda-12:<< parameters.image_date >>"
-      resource_class: gpu.nvidia.small.gen2
-    steps:
-      - checkout
-      - restore_cache:
-          keys:
-            - cuda-<< parameters.image_date >>-{{ arch }}-
-      - run:
-          name: Install dependencies
-          command: |
-            sudo apt-get update
-            sudo apt-get install libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install libnccl2 libnccl-dev
-            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
-            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
-            rm -rf ccache-4.11.3-linux-x86_64
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Set CCache size
-          command: ccache --max-size 1G
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
-            LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            cmake . -B build \
-              -DMLX_BUILD_CUDA=ON \
-              -DCMAKE_CUDA_COMPILER=`which nvcc` \
-              -DCMAKE_BUILD_TYPE=DEBUG
-            cmake --build build -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-      - run:
-          name: CCache report
-          command: |
-            ccache --show-stats
-            ccache --zero-stats
-            ccache --cleanup
-      - save_cache:
-          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
-          paths:
-            - /home/circleci/.cache/ccache
-
-  build_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.10"
-      xcode_version:
-        type: string
-        default: "26.0.0"
-      build_env:
-        type: string
-        default: ""
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    resource_class: m4pro.medium
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            mkdir -p ~/miniconda3
-            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o ~/miniconda3/miniconda.sh
-            bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
-            rm ~/miniconda3/miniconda.sh
-            source ~/miniconda3/bin/activate
-            conda init --all
-            conda create -n env python=<< parameters.python_version >> -y
-            conda activate env
-            pip install --upgrade cmake
-            pip install nanobind==2.4.0
-            pip install --upgrade setuptools
-            pip install numpy
-            pip install twine
-            pip install build
-      - run:
-          name: Install Python package
-          command: |
-            conda activate env
-            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
-              pip install . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            conda activate env
-            pip install typing_extensions
-            python setup.py generate_stubs
-      - run:
-          name: Build Python package
-          command: |
-            conda activate env
-            python setup.py clean --all
-            << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
-      - when:
-          condition:
-            equal: ["3.10", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  conda activate env
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 python -m build -w
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  conda activate env
-                  twine upload dist/*
-      - store_artifacts:
-          path: dist/
-
-  build_linux_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.10"
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            PYTHON=python<< parameters.python_version >>
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            TZ=Etc/UTC sudo apt-get -y install tzdata
-            sudo add-apt-repository -y ppa:deadsnakes/ppa
-            sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            $PYTHON -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            << parameters.build_env >> pip install ".[dev]" -v
-            pip install typing_extensions
-            python setup.py generate_stubs
-            python setup.py clean --all
-            MLX_BUILD_STAGE=1 << parameters.build_env >> python -m build -w
-            bash python/scripts/repair_linux.sh
-      - when:
-          condition:
-            equal: ["3.10", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  source env/bin/activate
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 \
-                    python -m build -w
-                  auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload packages
-                command: |
-                  source env/bin/activate
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-  build_cuda_release:
-    parameters:
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: xlarge
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/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 cuda-toolkit-12-9 libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install zip
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
-            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-            << parameters.build_env >> MLX_BUILD_STAGE=2 \
-              CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              python -m build -w
-            bash python/scripts/repair_cuda.sh
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-workflows:
-  build_and_test:
-    when:
-      and:
-        - matches:
-            pattern: "^(?!pull/)[-\\w]+$"
-            value: << pipeline.git.branch >>
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - mac_build_and_test:
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test
-      - cuda_build_and_test:
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-      - build_documentation 
-
-  build_pypi_release:
-    when:
-      and:
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["PYPI_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_documentation:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          upload-docs: true
-      - build_linux_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              build_env: ["PYPI_RELEASE=1"]
-      - build_cuda_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              build_env: ["PYPI_RELEASE=1"]
-
-  prb:
-    when:
-      matches:
-        pattern: "^pull/\\d+(/head)?$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - hold:
-          type: approval
-      - apple/authenticate:
-          context: pr-approval
-      - mac_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test:
-          requires: [ hold ]
-      - cuda_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-  nightly_build:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.nightly_build >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-      - build_cuda_release
-
-  build_dev_release:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["DEV_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              build_env: ["DEV_RELEASE=1"]
-      - build_cuda_release:
-          matrix:
-            parameters:
-              build_env: ["DEV_RELEASE=1"]

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

@@ -2,8 +2,8 @@ name: 'Build CUDA wheel'
 description: 'Build CUDA wheel'
 
 inputs:
-  nvcc-location:
-    description: 'Location of nvcc compiler'
+  toolkit:
+    description: 'The CUDA toolkit'
     required: true
 
 runs:
@@ -12,13 +12,9 @@ runs:
     - name: Build package
       shell: bash
       env:
-        MLX_BUILD_STAGE: 2
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
+        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
       run: |
         pip install auditwheel build patchelf setuptools
         python setup.py clean --all
-        python -m build -w
-
-        if [ -f "python/scripts/repair_cuda.sh" ]; then
+        MLX_BUILD_STAGE=2 python -m build -w
         bash python/scripts/repair_cuda.sh
-        fi

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

@@ -2,19 +2,9 @@ name: 'Build and Test with CUDA'
 description: 'Build and test MLX with CUDA'
 
 inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-  nvcc-location:
-    description: 'Location of nvcc compiler'
+  toolkit:
+    description: 'The CUDA toolkit'
     required: true
-    default: '/usr/local/cuda-12.9/bin/nvcc'
-    # this value is dependent on the CUDA tools installed in the setup-linux workflow
 
 runs:
   using: "composite"
@@ -23,46 +13,14 @@ runs:
       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: Check if build actually worked
-      shell: bash
-      run: python -c "import mlx.core"
-
-    - name: Run Python tests - CPU
-      if: inputs.run-tests == 'true'
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-        DEVICE: cpu
-      run: python -m unittest discover python/tests -v
-
-    - name: Run Python tests - GPU
-      if: inputs.run-tests == 'true'
-      shell: bash
-      env:
-        LOW_MEMORY: 1
-        DEVICE: gpu
-      run: python -m tests discover python/tests -v
+        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
+      run: pip install --no-build-isolation -e ".[dev]" -v
 
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         cmake . -B build \
           -DMLX_BUILD_CUDA=ON \
-          -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }} \
+          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc \
           -DCMAKE_BUILD_TYPE=DEBUG
         cmake --build build -j $(nproc)
-    
-    - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
-      shell: bash
-      run: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-cuda-release
-      with:
-        nvcc-location: ${{ inputs.nvcc-location }}

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

@@ -1,19 +1,19 @@
 name: 'Build Documentation'
-description: 'Build documentation on a mac'
+description: 'Build documentation'
 
 runs:
   using: "composite"
   steps:
     - name: Setup machine
-      uses: ./.github/actions/setup-macos
+      uses: ./.github/actions/setup-linux
 
     - name: Install dependencies
-      shell: sh
+      shell: bash
       run: |
-        brew install doxygen
-        uv pip install --upgrade pip cmake
-        uv pip install -r docs/requirements.txt
-        uv pip install . -v
+        sudo apt-get install -y doxygen
+        source .venv/bin/activate
+        pip install -r docs/requirements.txt
+        pip install . -v
   
     - name: Build documentation
       shell: bash
@@ -24,8 +24,8 @@ runs:
         make html O=-W
     
     - name: Create artifact tar
-      shell: sh
-      run: tar -cf artifact.tar --cd docs --dereference build/html index.html
+      shell: bash
+      run: tar -cf artifact.tar -C docs --dereference build/html index.html
 
     # Do it manually because upload-pages-artifact requires gtar
     - name: Upload artifact

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

@@ -0,0 +1,40 @@
+name: 'Build Linux wheel'
+description: 'Build Linux wheel'
+
+inputs:
+  build-backend:
+    description: 'Build the backend mlx-cpu package'
+    type: boolean
+    required: false
+    default: false
+  arch:
+    description: 'Platform architecture tag'
+    required: true
+    type: choice
+    options:
+      - x86_64
+      - aarch64
+
+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 ${{ inputs.arch }}
+    - 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_${{ inputs.arch }}

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

@@ -1,34 +1,15 @@
 name: 'Build and Test on Linux'
 description: 'Build and test MLX on Linux'
 
-inputs:
-  build-type:
-    description: 'Build type'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-    type: boolean
-
 runs:
   using: "composite"
   steps:
-    - name: Set DEBUG
-      shell: sh
-      if: inputs.build-type == 'debug'
-      run: echo "DEBUG=1" >> $GITHUB_ENV
-
     - name: Install Python package
       shell: sh
       env:
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      run: pip install -e ".[dev]" -v
+        DEBUG: 1
+      run: pip install --no-build-isolation -e ".[dev]" -v
 
     - name: Generate package stubs
       shell: sh
@@ -36,43 +17,9 @@ runs:
         pip install typing_extensions
         python setup.py generate_stubs
     
-    - name: Run Python tests
-      if: inputs.run-tests == 'true'
-      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
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         mkdir -p build && cd build
         cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
         make -j $(nproc)
-    
-    - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
-      shell: sh
-      run: ./build/tests/tests
-    
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      shell: bash
-      run: |
-        pip install auditwheel patchelf build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        if [ -f "python/scripts/repair_linux.sh" ]; then
-          bash python/scripts/repair_linux.sh
-        fi
-
-        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-macos-release/action.yml

@@ -6,17 +6,25 @@ inputs:
     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(s)
-      shell: bash
-      env:
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
+    - name: Build Python package
+      shell: bash -l {0}
       run: |
-        uv pip install build
-        uv run --no-project setup.py clean --all
-        MLX_BUILD_STAGE=1 uv run -m build -w
-        uv run --no-project setup.py clean --all
-        MLX_BUILD_STAGE=2 uv run -m build -w
+        pip install build
+        python setup.py clean --all
+        MLX_BUILD_STAGE=1 python -m build -w
+
+    - name: Build backend package
+      if: ${{ inputs.build-backend }}
+      shell: bash -l {0}
+      run: |
+        python setup.py clean --all
+        MLX_BUILD_STAGE=2 python -m build -w

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

@@ -1,68 +1,51 @@
 name: 'Build and Test on macOS'
 description: 'Build and test MLX on macOS'
 
-inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-  build-jit:
-    description: 'Whether to build with JIT'
-    required: false
-    default: 'true'
-
 runs:
   using: "composite"
   steps:
     - name: Install dependencies
-      shell: sh
       env:
         DEBUG: 1
-        DEV_RELEASE: 1
+        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+      shell: bash -l {0}
       run: |
-        uv pip install --upgrade pip cmake setuptools
-        uv pip install nanobind==2.4.0 \
-          numpy torch tensorflow unittest-xml-reporting
-        uv pip install -e . -v
+        pip install --upgrade pip
+        pip install cmake setuptools nanobind==2.4.0
+        pip install -e . -v
 
     - name: Generate package stubs
-      shell: bash
+      shell: bash -l {0}
       run: |
-        uv pip install typing_extensions
-        uv run --no-project setup.py generate_stubs
+        pip install typing_extensions
+        python setup.py generate_stubs
+
+    - name: Install tests dependencies
+      shell: bash -l {0}
+      run: |
+        pip install numpy torch tensorflow unittest-xml-reporting
 
     - name: Run Python tests
-      if: inputs.run-tests == 'true'
-      shell: bash
+      shell: bash -l {0}
       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
+        DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
+        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -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
-      if: inputs.run-tests == 'true'
-      shell: bash
+      shell: bash -l {0}
       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
+        pip install -r requirements.txt
+        python setup.py build_ext --inplace
+        python test.py
     
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
-      shell: bash
+      shell: bash -l {0}
       run: |
         mkdir -p build
         cd build
@@ -70,8 +53,7 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
-      shell: bash
+      shell: bash -l {0}
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
@@ -79,8 +61,7 @@ runs:
       run: ./build/tests/tests
     
     - name: Build small binary with JIT
-      if: inputs.build-jit == 'true'
-      shell: bash
+      shell: bash -l {0}
       run: |
         mkdir -p build
         cd build
@@ -93,8 +74,7 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run Python tests with JIT
-      if: ${{ inputs.build-jit == 'true' && inputs.run-tests == 'true' }}
-      shell: bash
+      shell: bash -l {0}
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
@@ -102,23 +82,7 @@ runs:
         METAL_DEBUG_ERROR_MODE: 0
       run: |
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          uv pip install -e . -v
-        uv run -m xmlrunner discover \
+          pip install -e . -v
+        python -m xmlrunner discover \
             -v python/tests \
             -o test-results/gpu_jit
-
-    - name: Build macOS 13 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 13.0
-    - name: Build macOS 14 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 14.0
-    - name: Build macOS 15 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 15.0

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

@@ -2,14 +2,10 @@ 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'
+  toolkit:
+    description: 'Which toolkit to install'
     required: false
-    default: 'linux'
-    type: choice
-    options:
-      - linux
-      - cuda
+    default: 'cpu'
   python-version:
     description: 'Version of python to set up'
     required: false
@@ -18,56 +14,62 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Free disk space
-      shell: sh
-      if: inputs.runner-type == 'linux'
-      run: sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+    - name: Use ccache
+      uses: hendrikmuhs/ccache-action@v1.2
+      with:
+        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-py${{ inputs.python-version }}
+        max-size: 1GB
 
     - 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
+        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip
 
     - uses: actions/setup-python@v6
       with:
         python-version: ${{ inputs.python-version }}
-        cache: 'pip'
 
-    - name: setup python venv
+    - name: Setup Python venv
       shell: bash
       run: |
         python -m venv .venv
         source .venv/bin/activate
+        pip install setuptools cmake nanobind==2.4.0
         echo PATH=$PATH >> $GITHUB_ENV
-        pip install --upgrade pip cmake
+        # Make cmake search .venv for nanobind
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
 
     - 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'
+    - name: Install CUDA toolkit
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
+      shell: bash
       env:
-        TZ: Etc/UTC
-      shell: bash ## Specific to Ubuntu 22.04 & Architecture x86_64
+        # Note: the CI machine does not meet CUDA 13's driver requirement.
+        # Compatibility matrix:
+        # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
+        # The `nvcc` is installed into `/usr/local/cuda-VERSION/bin/nvcc` - but
+        # it's *not* on the default toolkit path.
+        PACKAGES: |
+          {
+            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-toolkit-12-6",
+            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-toolkit-12-9",
+            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
+          }
       run: |
-        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
+        export ARCH=${{ runner.arch == 'arm64' && 'arm64' || 'x86_64' }}
+        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/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
+        sudo apt-get install -y \
+            libnccl2 libnccl-dev \
+            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
 
-    - name: Package and Driver Report
-      if: inputs.runner-type == 'cuda'
+    - name: CUDA packages and driver report
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       shell: bash
       run: |
         sudo apt-get install -y ubuntu-drivers-common dkms

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

@@ -2,11 +2,6 @@ name: 'Setup macOS Environment'
 description: 'Install dependencies for macOS builds'
 
 inputs:
-  install-mpi:
-    description: 'Whether to install MPI'
-    required: false
-    default: 'true'
-    type: boolean
   python-version:
     description: 'Python version to use'
     required: false
@@ -17,15 +12,13 @@ runs:
   steps:
     - name: Install Homebrew packages
       shell: sh
-      if: inputs.install-mpi == 'true'
       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
+    - uses: conda-incubator/setup-miniconda@v3
       with:
+        miniconda-version: "latest"
         python-version: ${{ inputs.python-version }}
-          activate-environment: true

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

@@ -0,0 +1,69 @@
+name: 'Run Linux tests'
+
+inputs:
+  cpu-only:
+    description: 'Skip GPU tests'
+    required: false
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Run MPI tests
+      shell: bash
+      run: |
+        echo "::group::MPI tests"
+        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
+        echo "::endgroup::"
+
+    - name: Run distributed tests
+      if: ${{ inputs.cpu-only == 'true' }}
+      shell: bash
+      run: |
+        echo "::group::Distributed tests"
+        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
+        echo "::endgroup::"
+
+    - name: Run Python tests - CPU
+      if: ${{ inputs.cpu-only == 'true' }}
+      shell: bash
+      env:
+        DEVICE: cpu
+      run: |
+        echo "::group::Python tests - CPU"
+        python -m unittest discover python/tests -v
+        echo "::endgroup::"
+
+    - name: Run Python tests - GPU
+      if: ${{ inputs.cpu-only == 'false' }}
+      shell: bash
+      env:
+        DEVICE: gpu
+      run: |
+        echo "::group::Python tests - GPU"
+        python -m tests discover python/tests -v
+        echo "::endgroup::"
+
+    - name: Run CPP tests - CPU
+      shell: bash
+      env:
+        DEVICE: cpu
+      run: |
+        echo "::group::CPP tests - CPU"
+        ./build/tests/tests
+        echo "::endgroup::"
+
+    - name: Run CPP tests - GPU
+      if: ${{ inputs.cpu-only == 'false' }}
+      shell: bash
+      env:
+        DEVICE: gpu
+      run: |
+        echo "::group::CPP tests - GPU"
+        ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
+        echo "::endgroup::"

.github/workflows/documentation.yml

@@ -8,7 +8,7 @@ permissions:
 
 jobs:
   build:
-    runs-on: [self-hosted, macos]
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/build-docs

.github/workflows/nightly.yml

@@ -18,10 +18,10 @@ jobs:
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
-          run-tests: false
+          build-backend: ${{ matrix.python-version == '3.10' }}
+          arch: "x86_64"
       - name: Upload mlx artifacts
         uses: actions/upload-artifact@v5
         with:
@@ -40,23 +40,26 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
+        python_version: ["3.11", "3.12", "3.13", "3.14"]
+        runner:
+          - ubuntu-22.04
+          - ubuntu-22.04-arm
+    runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
         with:
           python-version: ${{ matrix.python_version }}
       - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/test-linux
+        with:
+          cpu-only: true
 
   build_mac_release:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
         python-version: ["3.10", "3.13"]
-        # TODO: 3.14 had issues finding a compatible tensorflow
-    env:
-      MACOSX_DEPLOYMENT_TARGET: "15.0"
     runs-on: [self-hosted, macos]
     steps:
       - uses: actions/checkout@v5
@@ -64,16 +67,16 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - uses: ./.github/actions/build-macos
-
-  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
+      - name: Build macOS 15 package
+        uses: ./.github/actions/build-macos-release
         with:
-          runner-type: 'cuda'
-      - uses: ./.github/actions/build-cuda
+          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_release:
     if: github.repository == 'ml-explore/mlx'
@@ -82,36 +85,14 @@ jobs:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
         with:
-          runner-type: 'cuda'
+          toolkit: 'cuda-12.9'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          nvcc-location: '/usr/local/cuda-12.9/bin/nvcc'
+          toolkit: 'cuda-12.9'
       - 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,28 +1,52 @@
 name: Build and Test
 
-on: pull_request  
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      # For testing CI without starting a pull request:
+      - test/*
 
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/head/main' }}
+
 jobs:
   check_lint:
     runs-on: ubuntu-22.04
     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
+    needs: check_lint
+    strategy:
+      matrix:
+        runner:
+          - ubuntu-22.04
+          - ubuntu-22.04-arm
+      fail-fast: false
+    runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
       - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/test-linux
+        with:
+          cpu-only: true
 
   mac_build_and_test:
     if: github.repository == 'ml-explore/mlx'
+    strategy:
+      matrix:
+        macos-target: ["14.0", "15.0"]
     runs-on: [self-hosted, macos]
+    env:
+      MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }}
     needs: check_lint
     steps:
       - uses: actions/checkout@v5
@@ -31,18 +55,25 @@ jobs:
 
   cuda_build_and_test:
     if: github.repository == 'ml-explore/mlx'
+    strategy:
+      fail-fast: false
+      matrix:
+        toolkit: ['cuda-12.6', 'cuda-12.9']
     runs-on: gpu-t4-4-core
     needs: check_lint
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
         with:
-          runner-type: 'cuda'
+          toolkit: ${{ matrix.toolkit }}
       - uses: ./.github/actions/build-cuda
+        with:
+          toolkit: ${{ matrix.toolkit }}
+      - uses: ./.github/actions/test-linux
 
   build_documentation:
     if: github.repository == 'ml-explore/mlx'
-    runs-on: [self-hosted, macos]
+    runs-on: ubuntu-22.04
     needs: check_lint
     steps:
       - uses: actions/checkout@v5
@@ -50,6 +81,7 @@ jobs:
 
   linux_fedora_build_cpp:
     name: Linux Fedora CPP Build (${{ matrix.arch }})
+    needs: check_lint
     strategy:
       fail-fast: false
       matrix:

.github/workflows/release.yml

@@ -5,6 +5,11 @@ on:
     tags:
       - 'v*'
   workflow_dispatch:
+    inputs:
+      dev_release:
+        description: "Do a dev release or regular release"
+        required: true
+        default: "false"
 
 permissions:
   contents: read
@@ -12,9 +17,6 @@ permissions:
 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"
@@ -45,27 +47,31 @@ jobs:
     strategy:
       matrix:
         python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
+        arch: ['x86_64', 'aarch64']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
       PYPI_RELEASE: 1
+      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
         with:
           python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
-          run-tests: false
+          build-backend: ${{ matrix.python-version == '3.10' }}
+          arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v5
         with:
+          overwrite: true
           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:
+          overwrite: true
           name: mlx-cpu
           path: wheelhouse/mlx_cpu-*.whl
   
@@ -73,28 +79,50 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
-        # TODO: 3.14 had issues finding a compatible tensorflow
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
+      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-      - uses: ./.github/actions/build-macos
+
+      - name: Install dependencies
+        shell: bash -l {0}
+        run: |
+          pip install --upgrade pip
+          pip install cmake setuptools nanobind==2.4.0
+          pip install -e . -v
+      - name: Generate package stubs
+        shell: bash -l {0}
+        run: |
+          pip install typing_extensions
+          python setup.py generate_stubs
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
         with:
-          build-type: release
+          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:
+          overwrite: true
           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:
+          overwrite: true
           name: mlx-metal
           path: dist/mlx_metal-*.whl
 
@@ -103,18 +131,20 @@ jobs:
     runs-on: ubuntu-22-large
     env:
       PYPI_RELEASE: 1
+      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
         with:
-          runner-type: 'cuda'
+          toolkit: 'cuda-12.9'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          nvcc-location: '/usr/local/cuda-12.9/bin/nvcc'
+          toolkit: 'cuda-12.9'
       - name: Upload artifacts
         uses: actions/upload-artifact@v5
         with:
+          overwrite: true
           name: mlx-cuda
           path: wheelhouse/mlx_cuda-*.whl
 
@@ -125,25 +155,25 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
+      name: pypi
       url: https://pypi.org/p/mlx
     steps:
       - uses: actions/download-artifact@v6
         with:
           pattern: linux-wheels-*
-          merge-multiples: true
+          merge-multiple: true
           path: dist
       - uses: actions/download-artifact@v6
         with:
           pattern: mac-wheels-*
-          merge-multiples: true
+          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 }}
+          repository-url: https://upload.pypi.org/legacy/
   
   pypi-publish-cuda:
     name: Upload CUDA release to PyPI
@@ -152,7 +182,7 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
+      name: pypi
       url: https://pypi.org/p/mlx-cuda
     steps:
       - uses: actions/download-artifact@v6
@@ -164,7 +194,7 @@ jobs:
       - name: Publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
+          repository-url: https://upload.pypi.org/legacy/
 
   pypi-publish-cpu:
     name: Upload CPU release to PyPI
@@ -173,7 +203,7 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
+      name: pypi
       url: https://pypi.org/p/mlx-cpu
     steps:
       - uses: actions/download-artifact@v6
@@ -185,7 +215,7 @@ jobs:
       - name: Publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
+          repository-url: https://upload.pypi.org/legacy/
 
   pypi-publish-metal:
     name: Upload Metal release to PyPI
@@ -194,7 +224,7 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ needs.setup.outputs.pypi_env }}
+      name: pypi
       url: https://pypi.org/p/mlx-metal
     steps:
       - uses: actions/download-artifact@v6
@@ -206,5 +236,4 @@ jobs:
       - name: Publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          repository-url: ${{ needs.setup.outputs.pypi_url }}
-
+          repository-url: https://upload.pypi.org/legacy/

CMakeLists.txt

@@ -74,6 +74,7 @@ endif()
 if(MLX_USE_CCACHE)
   find_program(CCACHE_PROGRAM ccache)
   if(CCACHE_PROGRAM)
+    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
     set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
     set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
     set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
@@ -127,9 +128,12 @@ 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_macOS15_iOS18.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_26.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(
@@ -138,7 +142,6 @@ 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}>

benchmarks/cpp/irregular_strides.cpp

@@ -75,7 +75,7 @@ void time_irregular_binary_ops_3D() {
 
 void time_irregular_binary_ops_4D() {
   auto device = mx::default_device();
-  std::vector<int> shape = {8, 8, 512, 512};
+  mx::Shape shape = {8, 8, 512, 512};
   auto a = mx::random::uniform(shape);
   auto b = mx::random::uniform(shape);
 
@@ -115,7 +115,7 @@ void time_irregular_binary_ops_4D() {
 
 void time_irregular_reshape() {
   auto device = mx::default_device();
-  std::vector<int> shape;
+  mx::Shape shape;
   auto reshape_fn = [&shape, device](const mx::array& a) {
     return mx::reshape(a, shape, device);
   };
@@ -170,7 +170,7 @@ void time_irregular_astype_1D() {
 void time_irregular_astype_2D() {
   auto device = mx::default_device();
   int size = 2048;
-  std::vector<int> shape = {size, size};
+  mx::Shape shape = {size, size};
 
   auto a = mx::random::uniform(shape);
   TIMEM("2D regular", mx::astype, a, mx::int32, device);

benchmarks/python/blas/bench_gemv.py

@@ -1,6 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
-import argparse
 import os
 import subprocess
 import time

benchmarks/python/masked_scatter.py

@@ -0,0 +1,212 @@
+import math
+import os
+import subprocess
+import time
+from copy import copy
+from functools import partial
+
+import matplotlib.pyplot as plt
+import mlx.core as mx
+import numpy as np
+import torch
+from matplotlib.ticker import FuncFormatter
+
+RESULTS_DIR = "./results"
+
+
+if not os.path.isdir(RESULTS_DIR):
+    os.mkdir(RESULTS_DIR)
+
+DEVICE_NAME = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+DEVICE_NAME = DEVICE_NAME.decode("utf-8").strip("\n")
+
+TORCH_DEVICE = torch.device(
+    "mps"
+    if torch.backends.mps.is_available()
+    else ("cuda" if torch.cuda.is_available() else "cpu")
+)
+
+
+N_WARMUP = 5
+N_ITER_BENCH = 50
+N_ITER_FUNC = 20
+
+VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)]
+MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
+D_TYPES = ("float32", "float16")
+
+
+def _power_of_two_formatter(value, _position):
+    if value <= 0:
+        return ""
+    exponent = int(round(math.log2(value)))
+    if abs(value - (1 << exponent)) / value > 1e-6:
+        return f"{value:g}"
+    return f"$2^{{{exponent}}}$"
+
+
+def torch_sync():
+    if TORCH_DEVICE.type == "cuda":
+        torch.cuda.synchronize()
+    elif TORCH_DEVICE.type == "mps":
+        torch.mps.synchronize()
+
+
+def masked_scatter_mlx(self_arr, mask_arr, src_arr):
+    outs = []
+    for _ in range(N_ITER_FUNC):
+        out = copy(self_arr)
+        out[mask_arr] = src_arr
+        outs.append(out)
+    mx.eval(outs)
+    return outs
+
+
+@torch.no_grad()
+def masked_scatter_torch(self_tensor, mask_tensor, src_tensor):
+    outs = []
+    for _ in range(N_ITER_FUNC):
+        out = self_tensor.clone()
+        out.masked_scatter_(mask_tensor, src_tensor)
+        outs.append(out)
+    torch_sync()
+    return outs
+
+
+def measure(fn):
+    for _ in range(N_WARMUP):
+        fn()
+    start = time.perf_counter_ns()
+    for _ in range(N_ITER_BENCH):
+        fn()
+    end = time.perf_counter_ns()
+    return (end - start) * 1e-9
+
+
+def bytes_touched(length, true_count, item_size):
+    mask_bytes = length
+    self_bytes = length * item_size * 2  # read + write
+    src_bytes = true_count * item_size
+    return (mask_bytes + self_bytes + src_bytes) * N_ITER_FUNC * N_ITER_BENCH
+
+
+def build_case(length, density, np_dtype, torch_dtype):
+    true_count = max(1, int(round(length * density)))
+
+    rng = np.random.default_rng()
+    self_np = rng.normal(0.0, 1.0, length).astype(np_dtype)
+    mask_np = np.zeros(length, dtype=bool)
+    mask_np[:true_count] = True
+    rng.shuffle(mask_np)
+    src_np = rng.normal(0.0, 1.0, true_count).astype(np_dtype)
+
+    self_mlx = mx.array(self_np)
+    mask_mlx = mx.array(mask_np)
+    src_mlx = mx.array(src_np)
+
+    self_torch = torch.from_numpy(self_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
+    mask_torch = torch.from_numpy(mask_np).to(device=TORCH_DEVICE)
+    src_torch = torch.from_numpy(src_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
+
+    # Correctness check once per configuration
+    mx_out = mx.array(self_np)
+    mx_out[mask_mlx] = src_mlx
+    mx.eval(mx_out)
+    torch_out = self_torch.clone()
+    torch_out.masked_scatter_(mask_torch, src_torch)
+
+    atol = 5e-3 if np_dtype == np.float16 else 1e-5
+    if not np.allclose(np.array(mx_out), torch_out.cpu().numpy(), atol=atol):
+        raise AssertionError("masked_scatter results diverged between MLX and Torch")
+
+    return (self_mlx, mask_mlx, src_mlx, self_torch, mask_torch, src_torch, true_count)
+
+
+def bench_case(length, density, dtype):
+    np_dtype = getattr(np, dtype)
+    torch_dtype = getattr(torch, dtype)
+    (
+        self_mlx,
+        mask_mlx,
+        src_mlx,
+        self_torch,
+        mask_torch,
+        src_torch,
+        true_count,
+    ) = build_case(length, density, np_dtype, torch_dtype)
+
+    time_mlx = measure(partial(masked_scatter_mlx, self_mlx, mask_mlx, src_mlx))
+    time_torch = measure(
+        partial(masked_scatter_torch, self_torch, mask_torch, src_torch)
+    )
+
+    total_bytes = bytes_touched(length, true_count, np_dtype().itemsize)
+    bytes_per_gb = float(1024**3)
+    mlx_gbps = (total_bytes / bytes_per_gb) / time_mlx
+    torch_gbps = (total_bytes / bytes_per_gb) / time_torch
+
+    return time_mlx, time_torch, mlx_gbps, torch_gbps
+
+
+def plot_density(ax_perf, ax_speedup, density, dtype):
+    mlx_gbps = []
+    torch_gbps = []
+    mlx_times = []
+    torch_times = []
+
+    for length in VECTOR_LENGTHS:
+        t_mlx, t_torch, gbps_mlx, gbps_torch = bench_case(length, density, dtype)
+        mlx_gbps.append(gbps_mlx)
+        torch_gbps.append(gbps_torch)
+        mlx_times.append(t_mlx)
+        torch_times.append(t_torch)
+
+    ax_perf.plot(VECTOR_LENGTHS, mlx_gbps, "tab:blue", label="MLX")
+    ax_perf.plot(VECTOR_LENGTHS, torch_gbps, "tab:red", label="Torch")
+    ax_perf.set_xscale("log", base=2)
+    ax_perf.set_xticks(VECTOR_LENGTHS)
+    formatter = FuncFormatter(_power_of_two_formatter)
+    ax_perf.xaxis.set_major_formatter(formatter)
+    ax_perf.set_title(f"density={density:.2f}")
+    ax_perf.set_ylabel("GB/s")
+    ax_perf.grid(True, which="both", linestyle=":", alpha=0.4)
+    ax_perf.legend()
+
+    speedup = np.array(torch_times) / np.array(mlx_times)
+    ax_speedup.plot(VECTOR_LENGTHS, speedup, "tab:green")
+    ax_speedup.axhline(1.0, color="tab:gray", linestyle="--")
+    ax_speedup.set_xscale("log", base=2)
+    ax_speedup.set_xticks(VECTOR_LENGTHS)
+    ax_speedup.xaxis.set_major_formatter(formatter)
+    ax_speedup.set_ylabel("Speedup (Torch_t / MLX_t)")
+    ax_speedup.grid(True, which="both", linestyle=":", alpha=0.4)
+
+
+def main():
+    for dtype in D_TYPES:
+        fig, axs = plt.subplots(
+            len(MASK_DENSITIES),
+            2,
+            figsize=(10, 12),
+            layout="constrained",
+            sharex=True,
+        )
+
+        for i, density in enumerate(MASK_DENSITIES):
+            plot_density(axs[i][0], axs[i][1], density, dtype)
+            axs[i][0].set_xlabel("vector length")
+            axs[i][1].set_xlabel("vector length")
+
+        fig.suptitle(
+            f"{DEVICE_NAME.replace('Apple ', '')} ({TORCH_DEVICE.type}) | dtype={dtype}"
+        )
+        output_path = os.path.join(
+            RESULTS_DIR,
+            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf",
+        )
+        fig.savefig(output_path)
+        plt.close(fig)
+
+
+if __name__ == "__main__":
+    main()

cmake/Findnvpl.cmake

@@ -0,0 +1,3 @@
+# This file does nothing but to suppress the cmake warning: "By not providing
+# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
+# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.

docs/src/install.rst

@@ -17,11 +17,10 @@ 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 >= 13.5
+- macOS >= 14.0
 
 .. note::
-    MLX is only available on devices running macOS >= 13.5
-    It is highly recommended to use macOS 14 (Sonoma)
+    MLX is only available on devices running macOS >= 14.0 and higher.
 
 CUDA
 ^^^^

docs/src/usage/distributed.rst

@@ -7,12 +7,13 @@ 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 two different communication backends:
+moment we support three different communication backends:
 
 * `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
   full-featured and mature distributed communications library
-* A **ring** backend of our own that uses native TCP sockets and should be
-  faster for thunderbolt connections.
+* 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.
 
 The list of all currently supported operations and their documentation can be
 seen in the :ref:`API docs<distributed>`.
@@ -84,9 +85,8 @@ Selecting Backend
 ^^^^^^^^^^^^^^^^^
 
 You can select the backend you want to use when calling :func:`init` by passing
-one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to
-initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they
-both fail then a singleton group is created.
+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.
 
 .. 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, using the Anaconda package manager or
+MPI can be installed with Homebrew, pip, 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,14 +228,16 @@ with the Anaconda package manager as follows:
 
     $ conda install conda-forge::openmpi
 
-Installing with Homebrew may require specifying the location of ``libmpi.dyld``
+Installing with Homebrew or pip requires 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``.
+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.
 
 .. code:: shell
 
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
     $ # or simply
     $ mlx.launch -n 2 test.py
 

docs/src/usage/indexing.rst

@@ -70,7 +70,8 @@ Differences from NumPy
 
   * Indexing does not perform bounds checking. Indexing out of bounds is
     undefined behavior.
-  * Boolean mask based indexing is not yet supported.
+  * Boolean mask based indexing is supported for assignment only (see
+    :ref:`boolean-mask-assignment`).
 
 The reason for the lack of bounds checking is that exceptions cannot propagate
 from the GPU. Performing bounds checking for array indices before launching the
@@ -143,3 +144,51 @@ expected. For example:
 
 In the above ``dfdx`` will have the correct gradient, namely zeros at ``idx``
 and ones elsewhere.
+
+.. _boolean-mask-assignment:
+
+Boolean Mask Assignment
+-----------------------
+
+MLX supports boolean indices using NumPy syntax. A mask must already be
+a :class:`bool_` MLX :class:`array` or a NumPy ``ndarray`` with ``dtype=bool``.
+Other index types are routed through the standard scatter code.
+
+.. code-block:: shell
+
+   >>> a = mx.array([1.0, 2.0, 3.0])
+   >>> mask = mx.array([True, False, True])
+   >>> updates = mx.array([5.0, 6.0])
+   >>> a[mask] = updates
+   >>> a
+   array([5.0, 2.0, 6.0], dtype=float32)
+
+Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar
+assignments, ``updates`` must provide at least as many elements as there are
+``True`` entries in ``mask``.
+
+.. code-block:: shell
+
+   >>> a = mx.zeros((2, 3))
+   >>> mask = mx.array([[True, False, True],
+                        [False, False, True]])
+   >>> a[mask] = 1.0
+   >>> a
+   array([[1.0, 0.0, 1.0],
+          [0.0, 0.0, 1.0]], dtype=float32)
+
+Boolean masks follow NumPy semantics:
+
+- The mask shape must match the shape of the axes it indexes exactly. No mask
+  broadcasting occurs.
+- Any axes not covered by the mask are taken in full.
+
+.. code-block:: shell
+
+   >>> a = mx.arange(1000).reshape(10, 10, 10)
+   >>> a[mx.random.randn(10, 10) > 0.0] = 0  # valid: mask covers axes 0 and 1
+
+The mask of shape ``(10, 10)`` applies to the first two axes, so ``a[mask]``
+selects the 1-D slices ``a[i, j, :]`` where ``mask[i, j]`` is ``True``.
+Shapes such as ``(1, 10, 10)`` or ``(10, 10, 1)`` do not match the indexed
+axes and therefore raise errors.

mlx/allocator.h

@@ -14,7 +14,7 @@ class Buffer {
   void* ptr_;
 
  public:
-  Buffer(void* ptr) : ptr_(ptr) {};
+  explicit 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_->data_ptr = other.array_desc_->data_ptr;
+  cpy.array_desc_->offset = other.array_desc_->offset;
   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_->data_ptr = buffer.raw_ptr();
+  array_desc_->offset = 0;
   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_->data_ptr = buffer.raw_ptr();
+  array_desc_->offset = 0;
   array_desc_->data_size = data_size;
   array_desc_->strides = std::move(strides);
   array_desc_->flags = flags;
@@ -167,14 +167,13 @@ void array::copy_shared_buffer(
     const Strides& strides,
     Flags flags,
     size_t data_size,
-    size_t offset /* = 0 */) {
+    int64_t offset /* = 0 */) {
   array_desc_->data = other.array_desc_->data;
   array_desc_->strides = strides;
   array_desc_->flags = flags;
   array_desc_->data_size = data_size;
-  auto char_offset = sizeof(char) * itemsize() * offset;
-  array_desc_->data_ptr = static_cast<void*>(
-      static_cast<char*>(other.array_desc_->data_ptr) + char_offset);
+  array_desc_->offset =
+      sizeof(char) * itemsize() * offset + other.array_desc_->offset;
 }
 
 void array::copy_shared_buffer(const array& other) {

mlx/array.h

@@ -354,15 +354,23 @@ class array {
     return array_desc_->data;
   }
 
-  // Return a raw pointer to the arrays 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.
   template <typename T>
   T* data() {
-    return static_cast<T*>(array_desc_->data_ptr);
+    return reinterpret_cast<T*>(
+        (static_cast<char*>(buffer().raw_ptr()) + array_desc_->offset));
   }
 
   template <typename T>
   const T* data() const {
-    return static_cast<T*>(array_desc_->data_ptr);
+    return const_cast<array&>(*this).data<T>();
+  }
+
+  int64_t offset() const {
+    return array_desc_->offset;
   }
 
   enum Status {
@@ -431,7 +439,7 @@ class array {
       const Strides& strides,
       Flags flags,
       size_t data_size,
-      size_t offset = 0);
+      int64_t offset = 0);
 
   void copy_shared_buffer(const array& other);
 
@@ -466,8 +474,8 @@ class array {
     // can share the underlying data buffer.
     std::shared_ptr<Data> data;
 
-    // Properly offset data pointer
-    void* data_ptr{nullptr};
+    // Offset from beginning of data pointer
+    int64_t offset{0};
 
     // 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) {
+    BinaryOpType bopt,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
   bool b_donatable = is_donatable(b, out);
   bool a_donatable = is_donatable(a, out);
   switch (bopt) {
     case BinaryOpType::ScalarScalar:
-      out.set_data(
-          allocator::malloc(out.itemsize()), 1, a.strides(), a.flags());
+      out.set_data(mallocfn(out.itemsize()), 1, a.strides(), a.flags());
       break;
     case BinaryOpType::ScalarVector:
       if (b_donatable) {
         out.copy_shared_buffer(b);
       } else {
         out.set_data(
-            allocator::malloc(b.data_size() * out.itemsize()),
+            mallocfn(b.data_size() * out.itemsize()),
             b.data_size(),
             b.strides(),
             b.flags());
@@ -62,7 +62,7 @@ inline void set_binary_op_output_data(
         out.copy_shared_buffer(a);
       } else {
         out.set_data(
-            allocator::malloc(a.data_size() * out.itemsize()),
+            mallocfn(a.data_size() * out.itemsize()),
             a.data_size(),
             a.strides(),
             a.flags());
@@ -75,7 +75,7 @@ inline void set_binary_op_output_data(
         out.copy_shared_buffer(b);
       } else {
         out.set_data(
-            allocator::malloc(a.data_size() * out.itemsize()),
+            mallocfn(a.data_size() * out.itemsize()),
             a.data_size(),
             a.strides(),
             a.flags());
@@ -88,7 +88,7 @@ inline void set_binary_op_output_data(
           b_donatable && b.flags().row_contiguous && b.size() == out.size()) {
         out.copy_shared_buffer(b);
       } else {
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(mallocfn(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(nullptr);
+    out.set_data(allocator::malloc(0));
     return;
   }
   Strides strides(out.ndim(), 0);

mlx/backend/common/compiled.cpp

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

mlx/backend/common/compiled.h

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

mlx/backend/common/copy.h

@@ -22,7 +22,11 @@ enum class CopyType {
   GeneralGeneral
 };
 
-inline bool set_copy_output_data(const array& in, array& out, CopyType ctype) {
+inline bool set_copy_output_data(
+    const array& in,
+    array& out,
+    CopyType ctype,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
   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.
@@ -31,14 +35,14 @@ inline bool set_copy_output_data(const array& in, array& out, CopyType ctype) {
       return true;
     } else {
       out.set_data(
-          allocator::malloc(in.data_size() * out.itemsize()),
+          mallocfn(in.data_size() * out.itemsize()),
           in.data_size(),
           in.strides(),
           in.flags());
       return false;
     }
   } else {
-    out.set_data(allocator::malloc(out.nbytes()));
+    out.set_data(mallocfn(out.nbytes()));
     return false;
   }
 }

mlx/backend/common/slicing.cpp

@@ -14,17 +14,13 @@ std::tuple<int64_t, Strides> prepare_slice(
     data_offset += start_indices[i] * in.strides()[i];
     inp_strides[i] = in.strides()[i] * strides[i];
   }
-  // Normalize the offset
-  if (data_offset < 0) {
-    data_offset += in.data_size();
-  }
   return std::make_tuple(data_offset, inp_strides);
 }
 
 void shared_buffer_slice(
     const array& in,
     const Strides& out_strides,
-    size_t data_offset,
+    int64_t data_offset,
     size_t data_size,
     array& out) {
   // Compute row/col contiguity
@@ -45,23 +41,30 @@ void slice(
     const Shape& start_indices,
     const Shape& strides) {
   if (out.size() == 0) {
-    out.set_data(nullptr);
+    out.set_data(allocator::malloc(0));
     return;
   }
 
   // Calculate out strides, initial offset
   auto [data_offset, inp_strides] = prepare_slice(in, start_indices, strides);
-  int64_t data_end = 1;
-  for (int i = 0; i < start_indices.size(); ++i) {
-    if (in.shape()[i] > 1) {
-      auto end_idx = start_indices[i] + out.shape()[i] * strides[i] - 1;
-      data_end += end_idx * in.strides()[i];
+
+  // Get the location of the end based on the inp strides and out.shape()
+  int64_t low_idx = 0;
+  int64_t high_idx = 0;
+  for (int i = 0; i < inp_strides.size(); ++i) {
+    auto delta = inp_strides[i] * (out.shape()[i] - 1);
+    if (inp_strides[i] > 0) {
+      high_idx += delta;
+    } else {
+      low_idx += delta;
     }
   }
-  if (data_end < 0) {
-    data_end += in.data_size();
+  int64_t data_size = (high_idx - low_idx) + 1;
+  if (data_size < 0) {
+    std::ostringstream msg;
+    msg << "[slice] Computed invalid data size: " << data_size << ".";
+    throw std::runtime_error(msg.str());
   }
-  size_t data_size = (data_end - data_offset);
   shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
 }
 

mlx/backend/common/ternary.h

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

mlx/backend/common/unary.h

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

mlx/backend/cpu/binary.cpp

@@ -14,233 +14,11 @@
 
 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(a, b, out, detail::Add(), stream());
+  binary_op_cpu(a, b, out, detail::Add(), stream());
 }
 
 void DivMod::eval_cpu(
@@ -324,14 +102,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(a, b, out, detail::Divide(), stream());
+  binary_op_cpu(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(a, b, out, detail::Remainder(), stream());
+  binary_op_cpu(a, b, out, detail::Remainder(), stream());
 }
 
 void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -372,89 +150,90 @@ void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
       }
     });
   } else {
-    comparison_op(a, b, out, detail::Equal(), stream());
+    comparison_op_cpu(a, b, out, detail::Equal(), stream());
   }
 }
 
 void Greater::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::Greater(), stream());
+  comparison_op_cpu(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(inputs[0], inputs[1], out, detail::GreaterEqual(), stream());
+  comparison_op_cpu(
+      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(inputs[0], inputs[1], out, detail::Less(), stream());
+  comparison_op_cpu(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(inputs[0], inputs[1], out, detail::LessEqual(), stream());
+  comparison_op_cpu(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(a, b, out, detail::LogAddExp(), stream());
+  binary_float_op_cpu(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(in1, in2, out, detail::LogicalAnd(), stream());
+  binary_op_cpu(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(in1, in2, out, detail::LogicalOr(), stream());
+  binary_op_cpu(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(a, b, out, detail::Maximum(), stream());
+  binary_op_cpu(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(a, b, out, detail::Minimum(), stream());
+  binary_op_cpu(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(a, b, out, detail::Multiply(), stream());
+  binary_op_cpu(a, b, out, detail::Multiply(), stream());
 }
 
 void NotEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(inputs[0], inputs[1], out, detail::NotEqual(), stream());
+  comparison_op_cpu(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(a, b, out, detail::Power(), stream());
+  binary_op_cpu(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(a, b, out, detail::Subtract(), stream());
+  binary_op_cpu(a, b, out, detail::Subtract(), stream());
 }
 
 void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
@@ -463,19 +242,19 @@ void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
   switch (op_) {
     case BitwiseBinary::And:
-      binary_int(a, b, out, detail::BitwiseAnd(), stream());
+      binary_int_op_cpu(a, b, out, detail::BitwiseAnd(), stream());
       break;
     case BitwiseBinary::Or:
-      binary_int(a, b, out, detail::BitwiseOr(), stream());
+      binary_int_op_cpu(a, b, out, detail::BitwiseOr(), stream());
       break;
     case BitwiseBinary::Xor:
-      binary_int(a, b, out, detail::BitwiseXor(), stream());
+      binary_int_op_cpu(a, b, out, detail::BitwiseXor(), stream());
       break;
     case BitwiseBinary::LeftShift:
-      binary_int(a, b, out, detail::LeftShift(), stream());
+      binary_int_op_cpu(a, b, out, detail::LeftShift(), stream());
       break;
     case BitwiseBinary::RightShift:
-      binary_int(a, b, out, detail::RightShift(), stream());
+      binary_int_op_cpu(a, b, out, detail::RightShift(), stream());
       break;
   }
 }
@@ -484,7 +263,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(a, b, out, detail::ArcTan2(), stream());
+  binary_float_op_cpu(a, b, out, detail::ArcTan2(), stream());
 }
 
 } // namespace mlx::core

mlx/backend/cpu/binary.h

@@ -7,6 +7,7 @@
 #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 {
@@ -290,4 +291,227 @@ 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/distributed.cpp

@@ -95,4 +95,9 @@ 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/indexing.cpp

@@ -747,4 +747,108 @@ void ScatterAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
   });
 }
 
+template <typename T>
+void masked_scatter_impl(const array& mask, const array& src, array& out) {
+  ContiguousIterator mask_it(mask);
+  ContiguousIterator src_it(src);
+  ContiguousIterator out_it(out);
+
+  const bool* mask_ptr = mask.data<bool>();
+  const T* src_ptr = src.data<T>();
+  T* dst_ptr = out.data<T>();
+
+  const size_t batch_count = mask.shape(0);
+  const size_t mask_batch_size = mask.size() / batch_count;
+  const size_t src_batch_size = src.size() / batch_count;
+
+  for (uint b = 0; b < batch_count; ++b) {
+    size_t src_consumed = 0;
+    src_it.seek(b * src_batch_size);
+
+    for (size_t i = 0; i < mask_batch_size; ++i) {
+      if (mask_ptr[mask_it.loc]) {
+        if (src_consumed >= src_batch_size) {
+          throw std::runtime_error(
+              "[MaskedScatter::eval_cpu] Source does not have enough elements for mask.");
+        }
+        dst_ptr[out_it.loc] = src_ptr[src_it.loc];
+        src_it.step();
+        ++src_consumed;
+      }
+      mask_it.step();
+      out_it.step();
+    }
+  }
+}
+
+void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+
+  auto& dst = inputs[0];
+  auto& mask = inputs[1];
+  auto& src = inputs[2];
+
+  // Copy src into out (copy allocates memory for out)
+  auto ctype =
+      dst.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(dst, out, ctype, stream());
+
+  if (mask.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(mask);
+  encoder.set_input_array(src);
+  encoder.set_output_array(out);
+  encoder.dispatch([mask = array::unsafe_weak_copy(mask),
+                    src = array::unsafe_weak_copy(src),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        masked_scatter_impl<bool>(mask, src, out);
+        break;
+      case uint8:
+        masked_scatter_impl<uint8_t>(mask, src, out);
+        break;
+      case uint16:
+        masked_scatter_impl<uint16_t>(mask, src, out);
+        break;
+      case uint32:
+        masked_scatter_impl<uint32_t>(mask, src, out);
+        break;
+      case uint64:
+        masked_scatter_impl<uint64_t>(mask, src, out);
+        break;
+      case int8:
+        masked_scatter_impl<int8_t>(mask, src, out);
+        break;
+      case int16:
+        masked_scatter_impl<int16_t>(mask, src, out);
+        break;
+      case int32:
+        masked_scatter_impl<int32_t>(mask, src, out);
+        break;
+      case int64:
+        masked_scatter_impl<int64_t>(mask, src, out);
+        break;
+      case float16:
+        masked_scatter_impl<float16_t>(mask, src, out);
+        break;
+      case float32:
+        masked_scatter_impl<float>(mask, src, out);
+        break;
+      case float64:
+        masked_scatter_impl<double>(mask, src, out);
+        break;
+      case bfloat16:
+        masked_scatter_impl<bfloat16_t>(mask, src, out);
+        break;
+      case complex64:
+        masked_scatter_impl<complex64_t>(mask, src, out);
+        break;
+    }
+  });
+}
+
 } // namespace mlx::core

mlx/backend/cpu/matmul.cpp

@@ -147,37 +147,8 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
       copy_cpu(c, out, ctype, stream());
     } else {
       array beta_scalar = array(beta_, c.dtype());
-      auto bopt = get_binary_op_type(c, beta_scalar);
-      set_binary_op_output_data(c, beta_scalar, out, bopt);
       auto& encoder = cpu::get_command_encoder(stream());
-      encoder.set_input_array(c);
-      encoder.set_input_array(beta_scalar);
-      encoder.set_output_array(out);
-      encoder.dispatch([c = array::unsafe_weak_copy(c),
-                        beta_scalar = array::unsafe_weak_copy(beta_scalar),
-                        out = array::unsafe_weak_copy(out),
-                        bopt]() mutable {
-        switch (out.dtype()) {
-          case float16:
-            binary_op<float16_t, detail::Multiply>(c, beta_scalar, out, bopt);
-            break;
-          case float32:
-            binary_op<float, detail::Multiply>(c, beta_scalar, out, bopt);
-            break;
-          case float64:
-            binary_op<double, detail::Multiply>(c, beta_scalar, out, bopt);
-            break;
-          case bfloat16:
-            binary_op<bfloat16_t, detail::Multiply>(c, beta_scalar, out, bopt);
-            break;
-          case complex64:
-            binary_op<complex64_t, detail::Multiply>(c, beta_scalar, out, bopt);
-            break;
-          default:
-            throw std::runtime_error(
-                "[AddMM::eval_cpu] Unsupported dtype for beta scaling");
-        }
-      });
+      binary_float_op_cpu(c, beta_scalar, out, detail::Multiply(), stream());
       encoder.add_temporary(std::move(beta_scalar));
     }
     return;

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(nullptr);
+    out.set_data(allocator::malloc(0));
     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(nullptr);
+    out.set_data(allocator::malloc(0));
     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(nullptr);
+    out.set_data(allocator::malloc(0));
     return;
   }
 

mlx/backend/cuda/CMakeLists.txt

@@ -32,6 +32,7 @@ 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
@@ -43,6 +44,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/reduce/row_reduce.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/rms_norm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/rope.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/scan.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
@@ -124,7 +126,11 @@ endif()
 # Compute capability >= 7.0 is required for synchronization between CPU/GPU with
 # managed memory.
 if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
-  set(MLX_CUDA_ARCHITECTURES "native")
+  execute_process(
+    COMMAND bash detect_cuda_arch.sh
+    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+    OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
 endif()
 message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
@@ -136,6 +142,7 @@ FetchContent_Declare(
   URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
+set_target_properties(mlx PROPERTIES CCCL_DIR "${cccl_SOURCE_DIR}/include")
 
 # Use fixed version of NVTX.
 FetchContent_Declare(
@@ -161,7 +168,7 @@ target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver)
 FetchContent_Declare(
   cudnn
   GIT_REPOSITORY https://github.com/NVIDIA/cudnn-frontend.git
-  GIT_TAG v1.14.0
+  GIT_TAG v1.16.0
   GIT_SHALLOW TRUE
   EXCLUDE_FROM_ALL)
 set(CUDNN_FRONTEND_SKIP_JSON_LIB ON)

mlx/backend/cuda/allocator.cpp

@@ -1,6 +1,7 @@
 // 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"
 
@@ -67,6 +68,7 @@ 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;
 }
 
@@ -88,14 +90,41 @@ 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));
 }
 
-Buffer CudaAllocator::malloc(size_t size) {
+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_async(size_t size, int device, cudaStream_t stream) {
+  if (size == 0) {
+    return Buffer{new CudaBuffer{nullptr, 0, -1}};
+  }
+
   // Find available buffer from cache.
   std::unique_lock lock(mutex_);
   if (size <= small_block_size) {
@@ -106,6 +135,10 @@ Buffer CudaAllocator::malloc(size_t size) {
     size = page_size * ((size + page_size - 1) / page_size);
   }
 
+  if (size <= small_block_size || stream == nullptr) {
+    device = -1;
+  }
+
   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.
@@ -121,8 +154,13 @@ Buffer CudaAllocator::malloc(size_t size) {
     }
     lock.unlock();
     if (!buf) {
-      buf = new CudaBuffer{nullptr, size};
-      cudaError_t err = cudaMallocManaged(&buf->data, size);
+      buf = new CudaBuffer{nullptr, size, device};
+      cudaError_t err;
+      if (device == -1) {
+        err = cudaMallocManaged(&buf->data, size);
+      } else {
+        err = cudaMallocAsync(&buf->data, size, stream);
+      }
       if (err != cudaSuccess && err != cudaErrorMemoryAllocation) {
         throw std::runtime_error(fmt::format(
             "cudaMallocManaged failed: {}.", cudaGetErrorString(err)));
@@ -137,14 +175,26 @@ Buffer CudaAllocator::malloc(size_t size) {
   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(size_t size) {
+  return malloc_async(size, -1, 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;
@@ -167,8 +217,12 @@ size_t CudaAllocator::size(Buffer buffer) const {
 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);
+    }
     delete buf;
   }
 }
@@ -219,6 +273,17 @@ CudaAllocator& allocator() {
   return *allocator_;
 }
 
+Buffer malloc_async(size_t size, CommandEncoder& encoder) {
+  auto buffer = allocator().malloc_async(
+      size, encoder.device().cuda_device(), encoder.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 {
@@ -231,7 +296,11 @@ void* Buffer::raw_ptr() {
   if (!ptr_) {
     return nullptr;
   }
-  return static_cast<cu::CudaBuffer*>(ptr_)->data;
+  auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
+  if (cbuf.device != -1) {
+    copy_to_managed(cbuf);
+  }
+  return cbuf.data;
 }
 
 } // namespace allocator

mlx/backend/cuda/allocator.h

@@ -4,19 +4,24 @@
 
 #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>
 
 namespace mlx::core::cu {
 
+class CommandEncoder;
+
 using allocator::Buffer;
 
 // Stores cuda-managed unified memory.
 struct CudaBuffer {
   void* data;
   size_t size;
+  int device; // -1 for managed
 };
 
 class SmallSizePool {
@@ -45,6 +50,7 @@ class SmallSizePool {
 class CudaAllocator : public allocator::Allocator {
  public:
   Buffer malloc(size_t size) override;
+  Buffer malloc_async(size_t size, int device, cudaStream_t stream);
   void free(Buffer buffer) override;
   size_t size(Buffer buffer) const override;
 
@@ -69,9 +75,12 @@ 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, CommandEncoder& encoder);
+
 } // namespace mlx::core::cu

mlx/backend/cuda/arange.cu

@@ -41,9 +41,8 @@ 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));
   encoder.set_output_array(out);
 
   dispatch_int_float_types(out.dtype(), "Arange", [&](auto type_tag) {
@@ -58,7 +57,7 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
           num_blocks,
           block_dims,
           0,
-          out.data<OutType>(),
+          gpu_ptr<OutType>(out),
           out.data_size(),
           static_cast<CTYPE>(start_),
           static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_));

mlx/backend/cuda/arg_reduce.cu

@@ -140,8 +140,10 @@ 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));
 
   // Prepare the shapes, strides and axis arguments.
   Shape shape = remove_index(in.shape(), axis_);
@@ -154,7 +156,6 @@ 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) {
@@ -172,8 +173,8 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
           num_blocks,
           block_dim(),
           0,
-          in.data<T>(),
-          out.data<uint32_t>(),
+          gpu_ptr<T>(in),
+          gpu_ptr<uint32_t>(out),
           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,
-                        a.data<InType>(),
-                        b.data<InType>(),
-                        out.data<OutType>(),
+                        gpu_ptr<InType>(a),
+                        gpu_ptr<InType>(b),
+                        gpu_ptr<OutType>(out),
                         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,
-                      a.data<InType>(),
-                      b.data<InType>(),
-                      out.data<OutType>(),
+                      gpu_ptr<InType>(a),
+                      gpu_ptr<InType>(b),
+                      gpu_ptr<OutType>(out),
                       rest,
                       const_param(shape),
                       const_param(a_strides),
@@ -339,9 +339,9 @@ void binary_op_gpu_inplace(
                 num_blocks,
                 block_dims,
                 0,
-                a.data<InType>(),
-                b.data<InType>(),
-                out.data<OutType>(),
+                gpu_ptr<InType>(a),
+                gpu_ptr<InType>(b),
+                gpu_ptr<OutType>(out),
                 out.data_size());
           });
         }
@@ -365,7 +365,10 @@ void binary_op_gpu(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
+  auto& encoder = cu::get_command_encoder(s);
+
+  set_binary_op_output_data(
+      a, b, out, bopt, [&](auto n) { return cu::malloc_async(n, encoder); });
   binary_op_gpu_inplace<Op>(inputs, out, op, s);
 }
 

mlx/backend/cuda/binary_two.cu

@@ -245,14 +245,16 @@ void binary_two_op_gpu_inplace(
   auto& out_a = outputs[0];
   auto& out_b = outputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out_a, bopt);
-  set_binary_op_output_data(a, b, out_b, bopt);
+  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); });
+  set_binary_op_output_data(
+      a, b, out_b, bopt, [&](auto n) { return cu::malloc_async(n, encoder); });
 
   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);
@@ -313,10 +315,10 @@ void binary_two_op_gpu_inplace(
                         {num_blocks_x, num_blocks_y},
                         block_dims,
                         0,
-                        a.data<InType>(),
-                        b.data<InType>(),
-                        out_a.data<OutType>(),
-                        out_b.data<OutType>(),
+                        gpu_ptr<InType>(a),
+                        gpu_ptr<InType>(b),
+                        gpu_ptr<OutType>(out_a),
+                        gpu_ptr<OutType>(out_b),
                         rest,
                         const_param<dims_constant()>(shape),
                         const_param<dims_constant()>(a_strides),
@@ -332,10 +334,10 @@ void binary_two_op_gpu_inplace(
                       {num_blocks_x, num_blocks_y},
                       block_dims,
                       0,
-                      a.data<InType>(),
-                      b.data<InType>(),
-                      out_a.data<OutType>(),
-                      out_b.data<OutType>(),
+                      gpu_ptr<InType>(a),
+                      gpu_ptr<InType>(b),
+                      gpu_ptr<OutType>(out_a),
+                      gpu_ptr<OutType>(out_b),
                       rest,
                       const_param(shape),
                       const_param(a_strides),
@@ -366,10 +368,10 @@ void binary_two_op_gpu_inplace(
                 num_blocks,
                 block_dims,
                 0,
-                a.data<InType>(),
-                b.data<InType>(),
-                out_a.data<OutType>(),
-                out_b.data<OutType>(),
+                gpu_ptr<InType>(a),
+                gpu_ptr<InType>(b),
+                gpu_ptr<OutType>(out_a),
+                gpu_ptr<OutType>(out_b),
                 out_a.data_size());
           });
         }

mlx/backend/cuda/compiled.cpp

@@ -293,8 +293,13 @@ void Compiled::eval_gpu(
     }
   }
 
+  auto& encoder = cu::get_command_encoder(s);
+
   // Put outputs.
-  compiled_allocate_outputs(inputs, outputs, is_constant_, contiguous);
+  compiled_allocate_outputs(
+      inputs, outputs, is_constant_, contiguous, [&](auto n) {
+        return cu::malloc_async(n, encoder);
+      });
   for (auto& x : outputs) {
     args.append(x);
   }
@@ -324,7 +329,6 @@ 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,19 +270,19 @@ 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(allocator::malloc(out.nbytes()));
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
   Dtype dtype = out.dtype();
 
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-
   // Search cache.
-  ConvCacheKey cache_key{
+  BytesKey<ConvCacheKey> cache_key;
+  cache_key.pod = {
       encoder.device().cuda_device(),
       dtype_to_cudnn_type(dtype),
       vector_key(in.shape()),

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(allocator::malloc(unfolded.nbytes()));
+  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder));
   encoder.add_temporary(unfolded);
 
   int filter_size = params.C;
@@ -118,8 +118,8 @@ array unfold_inputs_nd(
         num_blocks,
         block_dims,
         0,
-        in.data<DataType>(),
-        unfolded.data<DataType>(),
+        gpu_ptr<DataType>(in),
+        gpu_ptr<DataType>(unfolded),
         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(allocator::malloc(unfolded.nbytes()));
+  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder));
   encoder.add_temporary(unfolded);
 
   int filter_size = params.C;
@@ -121,8 +121,8 @@ array grouped_unfold_transpose_inputs_nd(
         num_blocks,
         block_dims,
         0,
-        in.data<DataType>(),
-        unfolded.data<DataType>(),
+        gpu_ptr<DataType>(in),
+        gpu_ptr<DataType>(unfolded),
         filter_size,
         out_pixels,
         params);

mlx/backend/cuda/copy.cu

@@ -5,6 +5,21 @@
 
 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); });
+  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,
@@ -87,11 +102,31 @@ 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));
   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));
+    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,
-            in.data<InType>() + in_offset,
-            out.data<OutType>() + out_offset,
+            gpu_ptr<InType>(in) + in_offset,
+            gpu_ptr<OutType>(out) + 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 = in.data<InType>() + offset_in;
-            OutType* out_ptr = out.data<OutType>() + offset_out;
+            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+            OutType* out_ptr = gpu_ptr<OutType>(out) + 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 = in.data<InType>() + offset_in;
-            OutType* out_ptr = out.data<OutType>() + offset_out;
+            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+            OutType* out_ptr = gpu_ptr<OutType>(out) + 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),
-                    dynamic_offset_in.data<int64_t>(),
-                    dynamic_offset_out.data<int64_t>());
+                    gpu_ptr<int64_t>(dynamic_offset_in),
+                    gpu_ptr<int64_t>(dynamic_offset_out));
               });
             } 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,
-                  dynamic_offset_in.data<int64_t>(),
-                  dynamic_offset_out.data<int64_t>());
+                  gpu_ptr<int64_t>(dynamic_offset_in),
+                  gpu_ptr<int64_t>(dynamic_offset_out));
             }
           });
     });

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 = in.data<InType>() + offset_in;
-            OutType* out_ptr = out.data<OutType>() + offset_out;
+            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
             int ndim = shape.size();
             int work_per_thread = 1;
             auto dim0 = ndim > 0 ? shape.back() : 1;

mlx/backend/cuda/cuda_utils.h

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

mlx/backend/cuda/cudnn_utils.h

@@ -3,6 +3,7 @@
 #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"
@@ -23,7 +24,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>(x.data<void>());
+  uintptr_t address = reinterpret_cast<uintptr_t>(gpu_ptr<void>(x));
   for (; alignment < 32; alignment *= 2) {
     if (address % (alignment * 2)) {
       return alignment;
@@ -43,20 +44,20 @@ inline SmallVector<T> convert_vector(const Vec& vec) {
 // There are 2 differences from the const_param util from kernel_utils.cuh:
 // 1. The rest of array is filled with 0.
 // 2. This util can be used in .cpp files.
-template <typename T, template <typename U> class Vec>
-inline std::array<T, MAX_NDIM> vector_key(const Vec<T>& vec) {
-  if (vec.size() > MAX_NDIM) {
+template <int NDIM = MAX_NDIM, typename T, template <typename U> class Vec>
+inline std::array<T, NDIM> vector_key(const Vec<T>& vec) {
+  if (vec.size() > NDIM) {
     throw std::runtime_error(
-        fmt::format("ndim can not be larger than {}.", MAX_NDIM));
+        fmt::format("ndim can not be larger than {}.", NDIM));
   }
-  std::array<T, MAX_NDIM> result = {};
+  std::array<T, NDIM> result = {};
   std::copy_n(vec.begin(), vec.size(), result.begin());
   return result;
 }
 
 // Helpers used by get_data_ptrs to get pointers.
 inline void* get_data_ptr(const array& arr) {
-  return const_cast<void*>(arr.data<void>());
+  return const_cast<void*>(gpu_ptr<void>(arr));
 }
 
 template <typename T, typename = std::enable_if_t<std::is_scalar_v<T>>>

mlx/backend/cuda/custom_kernel.cpp

@@ -57,7 +57,7 @@ std::string build_kernel(
     const std::vector<std::string>& output_names,
     const std::vector<Dtype>& output_dtypes,
     const std::vector<std::pair<std::string, TemplateArg>>& template_args,
-    const std::vector<CustomKernelShapeInfo>& shape_infos) {
+    const std::vector<std::tuple<bool, bool, bool>>& shape_infos) {
   std::string kernel_source;
   kernel_source.reserve(header.size() + source.size() + 8192);
   kernel_source += default_header;
@@ -81,17 +81,17 @@ std::string build_kernel(
     kernel_source += ",\n";
     // Add input shape, strides and ndim if present in the source
     if (arr.ndim() > 0) {
-      if (shape_infos[i].shape) {
+      if (std::get<0>(shape_infos[i])) {
         kernel_source += "    const __grid_constant__ Shape ";
         kernel_source += name;
         kernel_source += "_shape,\n";
       }
-      if (shape_infos[i].strides) {
+      if (std::get<1>(shape_infos[i])) {
         kernel_source += "    const __grid_constant__ Strides ";
         kernel_source += name;
         kernel_source += "_strides,\n";
       }
-      if (shape_infos[i].ndim) {
+      if (std::get<2>(shape_infos[i])) {
         kernel_source += "    const __grid_constant__ int ";
         kernel_source += name;
         kernel_source += "_ndim,\n";
@@ -154,12 +154,12 @@ CustomKernelFunction cuda_kernel(
         "[custom_kernel] Must specify at least one output.");
   }
 
-  std::vector<CustomKernelShapeInfo> shape_infos;
+  std::vector<std::tuple<bool, bool, bool>> shape_infos;
   for (auto& n : input_names) {
-    CustomKernelShapeInfo shape_info;
-    shape_info.shape = source.find(n + "_shape") != std::string::npos;
-    shape_info.strides = source.find(n + "_strides") != std::string::npos;
-    shape_info.ndim = source.find(n + "_ndim") != std::string::npos;
+    std::tuple<bool, bool, bool> shape_info;
+    std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos;
+    std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos;
+    std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos;
     shape_infos.push_back(shape_info);
   }
 
@@ -254,8 +254,8 @@ std::vector<array> precompiled_cuda_kernel(
     std::optional<float> init_value,
     bool ensure_row_contiguous,
     StreamOrDevice s) {
-  std::vector<CustomKernelShapeInfo> shape_infos(
-      inputs.size(), CustomKernelShapeInfo{false, false, false});
+  std::vector<std::tuple<bool, bool, bool>> shape_infos(
+      inputs.size(), {false, false, false});
   return array::make_arrays(
       output_shapes,
       output_dtypes,
@@ -279,6 +279,7 @@ 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;
 
@@ -288,7 +289,7 @@ void CustomKernel::eval_gpu(
       copies.emplace_back(init_value_.value(), out.dtype());
       fill_gpu(copies.back(), out, s);
     } else {
-      out.set_data(allocator::malloc(out.nbytes()));
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
     }
   }
 
@@ -326,13 +327,13 @@ void CustomKernel::eval_gpu(
     const array& in = checked_inputs[i];
     auto& shape_info = shape_infos_[i];
     args.append(in);
-    if (shape_info.shape) {
+    if (std::get<0>(shape_info)) {
       args.append_ndim(in.shape());
     }
-    if (shape_info.strides) {
+    if (std::get<1>(shape_info)) {
       args.append_ndim(in.strides());
     }
-    if (shape_info.ndim) {
+    if (std::get<2>(shape_info)) {
       args.append<int32_t>(in.ndim());
     }
   }
@@ -356,7 +357,6 @@ 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/detect_cuda_arch.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+arch=`__nvcc_device_query`
+case "$arch" in
+    "90")
+        echo "90a" ;;
+    "100")
+        echo "100a" ;;
+    "121")
+        echo "121a" ;;
+    *)
+        echo "native" ;;
+esac

mlx/backend/cuda/device.cpp

@@ -46,6 +46,7 @@ Device::Device(int device) : device_(device) {
         "Device {} does not support synchronization in managed memory.",
         device_));
   }
+
   // The cublasLt handle is used by matmul.
   make_current();
   CHECK_CUBLAS_ERROR(cublasLtCreate(&lt_));
@@ -189,12 +190,41 @@ void CommandEncoder::insert_graph_dependencies(std::vector<GraphNode> nodes) {
   }
 }
 
+// Can be tuned with MLX_MAX_OPS_PER_BUFFER, MLX_MAX_MB_PER_BUFFER
+std::pair<int, int> get_graph_limits(Device& d) {
+  auto cc =
+      d.compute_capability_major() * 100 + d.compute_capability_minor() * 10;
+  int ops = 20;
+  int mb = 100;
+  switch (cc) {
+    case 800: // A100
+      ops = 20;
+      mb = 400;
+      break;
+    case 900: // H100
+      ops = 30;
+      mb = 400;
+      break;
+    case 1000: // B200
+      ops = 50;
+      mb = 500;
+      break;
+    case 1210: // DGX Spark
+      ops = 20;
+      mb = 25;
+      break;
+  }
+  return {env::max_ops_per_buffer(ops), env::max_mb_per_buffer(mb)};
+}
+
 CommandEncoder::CommandEncoder(Device& d)
     : device_(d),
       stream_(d),
       graph_(d),
       worker_(d),
-      graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {}
+      graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {
+  std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d);
+}
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
   worker_.add_task(std::move(task));
@@ -204,6 +234,7 @@ void CommandEncoder::set_input_array(const array& arr) {
   if (!use_cuda_graphs()) {
     return;
   }
+  bytes_in_graph_ += arr.data_size();
   auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
   active_deps_.push_back(id);
 }
@@ -301,8 +332,9 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
   insert_graph_dependencies(GraphNode{node, 'G'});
 }
 
-int CommandEncoder::get_num_ops() {
-  return node_count_;
+bool CommandEncoder::needs_commit() {
+  return (node_count_ > max_ops_per_graph_) ||
+      ((bytes_in_graph_ >> 20) > max_mb_per_graph_);
 }
 
 void CommandEncoder::commit() {
@@ -365,10 +397,11 @@ void CommandEncoder::commit() {
   // Put completion handlers in a batch.
   worker_.commit(stream_);
   node_count_ = 0;
+  bytes_in_graph_ = 0;
 }
 
 void CommandEncoder::synchronize() {
-  cudaStreamSynchronize(stream_);
+  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
   auto p = std::make_shared<std::promise<void>>();
   std::future<void> f = p->get_future();
   add_completed_handler([p = std::move(p)]() { p->set_value(); });

mlx/backend/cuda/device.h

@@ -3,6 +3,7 @@
 #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"
@@ -83,7 +84,7 @@ class CommandEncoder {
   }
 
   void add_completed_handler(std::function<void()> task);
-  int get_num_ops();
+  bool needs_commit();
   void commit();
 
   Device& device() {
@@ -130,6 +131,9 @@ class CommandEncoder {
   std::vector<std::uintptr_t> active_deps_;
   std::vector<std::uintptr_t> active_outputs_;
   std::unordered_map<std::uintptr_t, GraphNode> node_map_;
+  size_t bytes_in_graph_{0};
+  int max_ops_per_graph_;
+  int max_mb_per_graph_;
 };
 
 class Device {
@@ -165,6 +169,7 @@ class Device {
   int device_;
   int compute_capability_major_;
   int compute_capability_minor_;
+  std::string device_name_;
   cublasLtHandle_t lt_;
   cudnnHandle_t cudnn_;
   std::unordered_map<int, CommandEncoder> encoders_;

mlx/backend/cuda/distributed.cu

@@ -15,8 +15,10 @@ void AllReduce::eval_gpu(
   assert(inputs.size() == 1);
   assert(outputs.size() == 1);
 
-  auto set_input_output =
-      [s = stream()](const array& in, array& out) -> std::pair<array, array> {
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto set_input_output = [&](const array& in,
+                              array& out) -> std::pair<array, array> {
     if (!in.flags().row_contiguous) {
       copy_gpu(in, out, CopyType::General, s);
       return {out, out};
@@ -24,19 +26,17 @@ void AllReduce::eval_gpu(
       out.copy_shared_buffer(in);
       return {in, out};
     } else {
-      out.set_data(allocator::malloc(out.nbytes()));
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
       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,4 +53,69 @@ 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));
+
+  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));
+
+  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/eval.cpp

@@ -11,9 +11,6 @@
 
 namespace mlx::core::gpu {
 
-// Can be tuned with MLX_MAX_OPS_PER_BUFFER
-constexpr int default_max_nodes_per_graph = 20;
-
 bool is_available() {
   return true;
 }
@@ -53,8 +50,7 @@ void eval(array& arr) {
     encoder.add_temporary(s);
   }
 
-  if (encoder.get_num_ops() >=
-      env::max_ops_per_buffer(default_max_nodes_per_graph)) {
+  if (encoder.needs_commit()) {
     scheduler::notify_new_task(stream);
     encoder.add_completed_handler(
         [stream]() { scheduler::notify_task_completion(stream); });

mlx/backend/cuda/fence.cpp

@@ -1,6 +1,8 @@
 // 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 {
@@ -20,8 +22,24 @@ void Fence::wait(Stream s, const array&) {
   fence->event.wait(fence->count);
 }
 
-void Fence::update(Stream s, const array&) {
+void Fence::update(Stream s, const array& a, bool cross_device) {
   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 = bias.data<void>();
+  auto* bias_ptr = gpu_ptr<void>(bias);
   CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
       matmul_desc_,
       CUBLASLT_MATMUL_DESC_BIAS_POINTER,
@@ -278,9 +278,9 @@ void CublasGemm::run(
 
   execute(
       encoder,
-      out.data<void>(),
-      a.data<void>(),
-      b.data<void>(),
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
       nullptr,
       alpha);
 }
@@ -321,10 +321,10 @@ void CublasGemm::run(
 
   execute(
       encoder,
-      out.data<void>(),
-      a.data<void>(),
-      b.data<void>(),
-      c.data<void>(),
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(c),
       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(
-        allocator::malloc(nbytes),
+        cu::malloc_async(nbytes, encoder),
         {static_cast<int>(heuristic_.workspaceSize)},
         int8);
     encoder.add_temporary(workspace);
-    workspace_ptr = workspace.data<void>();
+    workspace_ptr = gpu_ptr<void>(workspace);
   }
 
   auto capture = encoder.capture_context();

mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp

@@ -25,9 +25,10 @@ void CublasGemm::run_batched(
   for (size_t i = 0; i < nbatch; ++i) {
     execute(
         encoder,
-        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,
+        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,
         nullptr,
         alpha);
     a_it.step();
@@ -60,10 +61,11 @@ void CublasGemm::run_batched(
   for (size_t i = 0; i < nbatch; ++i) {
     execute(
         encoder,
-        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,
+        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,
         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(
-      allocator::malloc(batch_count * sizeof(void*) * 3),
+      cu::malloc_async(batch_count * sizeof(void*) * 3, encoder),
       {batch_count * 3},
       uint64);
 
@@ -183,10 +183,10 @@ void CublasGemm::run_batched(
           num_blocks,
           block_dims,
           0,
-          pointers.data<int8_t*>(),
-          a.data<int8_t>(),
-          b.data<int8_t>(),
-          out.data<int8_t>(),
+          gpu_ptr<int8_t*>(pointers),
+          gpu_ptr<int8_t>(a),
+          gpu_ptr<int8_t>(b),
+          gpu_ptr<int8_t>(out),
           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,
-        pointers.data<int8_t*>(),
-        a.data<int8_t>(),
-        b.data<int8_t>(),
-        out.data<int8_t>(),
+        gpu_ptr<int8_t*>(pointers),
+        gpu_ptr<int8_t>(a),
+        gpu_ptr<int8_t>(b),
+        gpu_ptr<int8_t>(out),
         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 = pointers.data<int8_t*>();
+  auto a_pointers = gpu_ptr<int8_t*>(pointers);
   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(
-      allocator::malloc(batch_count * sizeof(uint64_t) * 4),
+      cu::malloc_async(batch_count * sizeof(uint64_t) * 4, encoder),
       {batch_count * 4},
       uint64);
 
@@ -271,11 +271,11 @@ void CublasGemm::run_batched(
           num_blocks,
           block_dims,
           0,
-          pointers.data<int8_t*>(),
-          a.data<int8_t>(),
-          b.data<int8_t>(),
-          c.data<int8_t>(),
-          out.data<int8_t>(),
+          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),
           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,
-        pointers.data<int8_t*>(),
-        a.data<int8_t>(),
-        b.data<int8_t>(),
-        c.data<int8_t>(),
-        out.data<int8_t>(),
+        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),
         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 = pointers.data<int8_t*>();
+  auto a_pointers = gpu_ptr<int8_t*>(pointers);
   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 = b.data<DataType>();
-      vec = a.data<DataType>();
+      mat = gpu_ptr<DataType>(b);
+      vec = gpu_ptr<DataType>(a);
       rows = N;
       std::swap(mat_strides, vec_strides);
     } else {
-      mat = a.data<DataType>();
-      vec = b.data<DataType>();
+      mat = gpu_ptr<DataType>(a);
+      vec = gpu_ptr<DataType>(b);
       rows = M;
     }
     uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
@@ -177,7 +177,7 @@ void gemv(
             0,
             mat,
             vec,
-            out.data<DataType>(),
+            gpu_ptr<DataType>(out),
             rows,
             cols);
       } else {
@@ -189,7 +189,7 @@ void gemv(
             0,
             mat,
             vec,
-            out.data<DataType>(),
+            gpu_ptr<DataType>(out),
             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] = inputs[i + 1].data<void>();
+    indices[i] = gpu_ptr<void>(inputs[i + 1]);
   }
   args.append(std::move(indices));
   SmallVector<int32_t> indices_shape(nidx * idx_ndim);
@@ -59,7 +59,9 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() > 0);
   const auto& src = 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));
   if (out.size() == 0) {
     return;
   }
@@ -80,7 +82,6 @@ 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) {
@@ -121,7 +122,6 @@ 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,7 +239,9 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   const auto& src = inputs[0];
   const auto& idx = inputs[1];
 
-  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));
   if (out.size() == 0) {
     return;
   }
@@ -251,7 +253,6 @@ 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) {
@@ -312,7 +313,6 @@ 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.cpp

@@ -279,11 +279,14 @@ void compile(
   // Compile program.
   std::vector<const char*> args;
   bool use_sass = compiler_supports_device_sass(device);
+  auto cc = device.compute_capability_major();
+  std::string arch_tag = (cc == 90 || cc == 100 || cc == 121) ? "a" : "";
   std::string compute = fmt::format(
-      "--gpu-architecture={}_{}{}",
+      "--gpu-architecture={}_{}{}{}",
       use_sass ? "sm" : "compute",
-      device.compute_capability_major(),
-      device.compute_capability_minor());
+      cc,
+      device.compute_capability_minor(),
+      arch_tag);
   args.push_back(compute.c_str());
   std::string cccl_include = cccl_dir();
   if (!cccl_include.empty()) {

mlx/backend/cuda/jit_module.h

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

mlx/backend/cuda/kernel_utils.cuh

@@ -9,6 +9,7 @@
 #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,9 +230,10 @@ 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](const array& x) {
+  auto set_output = [&s, &out, &encoder](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];
@@ -243,7 +244,7 @@ void LayerNorm::eval_gpu(
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            allocator::malloc(x.data_size() * x.itemsize()),
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -265,7 +266,6 @@ 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,
-          x.data<DataType>(),
-          w.data<DataType>(),
-          b.data<DataType>(),
-          out.data<DataType>(),
+          gpu_ptr<DataType>(x),
+          gpu_ptr<DataType>(w),
+          gpu_ptr<DataType>(b),
+          gpu_ptr<DataType>(out),
           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(allocator::malloc(gx.nbytes()));
+    gx.set_data(cu::malloc_async(gx.nbytes(), encoder));
   }
   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(allocator::malloc(gw_temp.nbytes()));
+      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder));
       encoder.add_temporary(gw_temp);
     }
   }
@@ -393,11 +393,11 @@ void LayerNormVJP::eval_gpu(
                 n_rows,
                 block_dim(),
                 0,
-                x.data<DataType>(),
-                w.data<DataType>(),
-                g.data<DataType>(),
-                gx.data<DataType>(),
-                gw_temp.data<DataType>(),
+                gpu_ptr<DataType>(x),
+                gpu_ptr<DataType>(w),
+                gpu_ptr<DataType>(g),
+                gpu_ptr<DataType>(gx),
+                gpu_ptr<DataType>(gw_temp),
                 eps_,
                 axis_size,
                 w_stride);

mlx/backend/cuda/load.cpp

@@ -0,0 +1,60 @@
+// 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));
+  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(allocator::malloc(out.nbytes()));
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
   } 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(
-        allocator::malloc(in.nbytes() / n),
+        cu::malloc_async(in.nbytes() / n, encoder),
         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,
-          in.data<DataType>(),
-          out.data<DataType>(),
+          gpu_ptr<DataType>(in),
+          gpu_ptr<DataType>(out),
           axis_size);
     });
   });

mlx/backend/cuda/lru_cache.h

@@ -135,12 +135,19 @@ class LRUCache {
 };
 
 // Turn a POD struct into a container key by doing bytes compare.
+//
+// Usage:
+//   BytesKey<MyKey> key;
+//   key.pod = { ... };
 template <typename T>
 struct BytesKey {
   T pod;
   static_assert(std::is_standard_layout_v<T>, "T is not POD");
 
-  BytesKey(T pod) : pod(std::move(pod)) {}
+  BytesKey() {
+    // Make sure the paddings between members are filled with 0.
+    memset(&pod, 0, sizeof(T));
+  }
 
   BytesKey(const BytesKey& other) {
     memcpy(&pod, &other.pod, sizeof(T));

mlx/backend/cuda/matmul.cpp

@@ -121,7 +121,7 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
 
   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(allocator::malloc(out.nbytes()));
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
     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(allocator::malloc(out.nbytes()));
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
     } else if (sty == 1 && stx == 0) {
       ldc = 0;
-      out.set_data(allocator::malloc(out.nbytes()));
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
     } else {
       // Copy C into out and set C to out
       ldc = c.shape(-1);

mlx/backend/cuda/primitives.cpp

@@ -28,7 +28,6 @@ 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)
@@ -38,9 +37,9 @@ NO_GPU(Inverse)
 NO_GPU(Cholesky)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
+NO_GPU(MaskedScatter)
 
 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,
-            w.data<T>(),
-            wq.data<uint8_t>(),
-            scales.data<T>(),
-            biases.data<T>(),
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<T>(scales),
+            gpu_ptr<T>(biases),
             w.size());
       });
     });
@@ -318,10 +318,10 @@ void affine_dequantize(
             num_blocks,
             block_dims,
             0,
-            wq.data<uint8_t>(),
-            scales.data<T>(),
-            biases.data<T>(),
-            w.data<T>(),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<T>(scales),
+            gpu_ptr<T>(biases),
+            gpu_ptr<T>(w),
             w.size());
       });
     });

mlx/backend/cuda/quantized/fp_quantize.cu

@@ -156,9 +156,9 @@ void fp_quantize(
           num_blocks,
           block_dims,
           0,
-          w.data<T>(),
-          wq.data<uint8_t>(),
-          scales.data<uint8_t>(),
+          gpu_ptr<T>(w),
+          gpu_ptr<uint8_t>(wq),
+          gpu_ptr<uint8_t>(scales),
           w.size());
     } else {
       throw std::runtime_error(
@@ -202,9 +202,9 @@ void fp_dequantize(
           num_blocks,
           block_dims,
           0,
-          wq.data<uint8_t>(),
-          scales.data<T>(),
-          w.data<T>(),
+          gpu_ptr<uint8_t>(wq),
+          gpu_ptr<uint8_t>(scales),
+          gpu_ptr<T>(w),
           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(allocator::malloc(w.nbytes()));
+    w.set_data(cu::malloc_async(w.nbytes(), enc));
 
     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(allocator::malloc(wq.nbytes()));
-    scales.set_data(allocator::malloc(scales.nbytes()));
+    wq.set_data(cu::malloc_async(wq.nbytes(), enc));
+    scales.set_data(cu::malloc_async(scales.nbytes(), enc));
     if (mode_ == QuantizationMode::Affine) {
       auto& biases = outputs[2];
-      biases.set_data(allocator::malloc(biases.nbytes()));
+      biases.set_data(cu::malloc_async(biases.nbytes(), enc));
       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,7 +143,9 @@ 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;
-  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));
   if (out.size() == 0) {
     return;
   }
@@ -152,8 +154,6 @@ 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,
-        keys.data<uint32_t>(),
-        out.data<uint8_t>(),
+        gpu_ptr<uint32_t>(keys),
+        gpu_ptr<uint8_t>(out),
         grid_dims,
         odd,
         bytes_per_key);
@@ -182,8 +182,8 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
         grid,
         block,
         0,
-        keys.data<uint32_t>(),
-        out.data<uint8_t>(),
+        gpu_ptr<uint32_t>(keys),
+        gpu_ptr<uint8_t>(out),
         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(allocator::malloc(out.nbytes()));
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
 
   auto get_args = [](size_t size, int N) {
     int threads = std::min(512UL, (size + N - 1) / N);
@@ -100,14 +100,14 @@ void all_reduce(
   Dtype dt = in.dtype();
 
   // Cub doesn't like const pointers for load (sigh).
-  void* indata = const_cast<void*>(in.data<void>());
+  void* indata = const_cast<void*>(gpu_ptr<void>(in));
 
   // 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(allocator::malloc(intermediate.nbytes()));
+    intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder));
     encoder.add_temporary(intermediate);
     encoder.set_output_array(intermediate);
     dispatch_all_types(dt, [&](auto type_tag) {
@@ -122,14 +122,14 @@ void all_reduce(
             threads,
             0,
             static_cast<T*>(indata),
-            intermediate.data<U>(),
+            gpu_ptr<U>(intermediate),
             block_step,
             insize);
       });
     });
 
     // Set the input for the next step and recalculate the blocks
-    indata = intermediate.data<void>();
+    indata = gpu_ptr<void>(intermediate);
     dt = intermediate.dtype();
     insize = intermediate.size();
     std::tie(blocks, threads, block_step) = get_args(insize, N_READS);
@@ -149,7 +149,7 @@ void all_reduce(
           threads,
           0,
           static_cast<T*>(indata),
-          out.data<U>(),
+          gpu_ptr<U>(out),
           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);
+  allocate_same_layout(out, in, axes, encoder);
 
   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*>(in.data<T>());
+        T* indata = const_cast<T*>(gpu_ptr<T>(in));
 
         constexpr int N_READS = 4;
         constexpr int BM = 32;
@@ -276,7 +276,7 @@ void col_reduce_looped(
             blocks,
             0,
             indata,
-            out.data<U>(),
+            gpu_ptr<U>(out),
             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);
+  allocate_same_layout(out, in, axes, encoder);
 
   encoder.set_input_array(in);
   encoder.set_output_array(out);
@@ -312,8 +312,8 @@ void col_reduce_small(
           grid,
           block,
           0,
-          in.data<T>(),
-          out.data<U>(),
+          gpu_ptr<T>(in),
+          gpu_ptr<U>(out),
           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(allocator::malloc(out.nbytes()));
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
   }
 
   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, out.data<U>(), out.size());
+          kernel, grid, block, 0, gpu_ptr<U>(out), out.size());
     });
   });
 }

mlx/backend/cuda/reduce/reduce_utils.cuh

@@ -5,6 +5,7 @@
 #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>
@@ -92,9 +93,10 @@ 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) {
+    const std::vector<int>& axes,
+    cu::CommandEncoder& encoder) {
   if (in.flags().row_contiguous) {
-    out.set_data(allocator::malloc(out.nbytes()));
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
     return;
   }
 
@@ -133,7 +135,7 @@ inline void allocate_same_layout(
   fl.col_contiguous = cc;
   fl.contiguous = true;
   out.set_data(
-      allocator::malloc(out.nbytes()),
+      cu::malloc_async(out.nbytes(), encoder),
       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);
+  allocate_same_layout(out, in, axes, encoder);
 
   // 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*>(in.data<T>());
+      T* indata = const_cast<T*>(gpu_ptr<T>(in));
       int size = plan.shape.back();
       encoder.add_kernel_node(
-          kernel, grid, block, 0, indata, out.data<U>(), out.size(), size);
+          kernel, grid, block, 0, indata, gpu_ptr<U>(out), 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);
+  allocate_same_layout(out, in, axes, encoder);
 
   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, in.data<T>(), out.data<U>(), args);
+          kernel, grid, block, 0, gpu_ptr<T>(in), gpu_ptr<U>(out), args);
     });
   });
 }

mlx/backend/cuda/rms_norm.cu

@@ -176,9 +176,10 @@ 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](const array& x) {
+  auto set_output = [&s, &out, &encoder](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];
@@ -189,7 +190,7 @@ void RMSNorm::eval_gpu(
         out.copy_shared_buffer(x);
       } else {
         out.set_data(
-            allocator::malloc(x.data_size() * x.itemsize()),
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -209,7 +210,6 @@ 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,
-          x.data<DataType>(),
-          w.data<DataType>(),
-          out.data<DataType>(),
+          gpu_ptr<DataType>(x),
+          gpu_ptr<DataType>(w),
+          gpu_ptr<DataType>(out),
           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(allocator::malloc(gx.nbytes()));
+    gx.set_data(cu::malloc_async(gx.nbytes(), encoder));
   }
   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(allocator::malloc(gw_temp.nbytes()));
+      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder));
       encoder.add_temporary(gw_temp);
     }
   }
@@ -318,11 +318,11 @@ void RMSNormVJP::eval_gpu(
                 n_rows,
                 block_dim(),
                 0,
-                x.data<DataType>(),
-                w.data<DataType>(),
-                g.data<DataType>(),
-                gx.data<DataType>(),
-                gw_temp.data<DataType>(),
+                gpu_ptr<DataType>(x),
+                gpu_ptr<DataType>(w),
+                gpu_ptr<DataType>(g),
+                gpu_ptr<DataType>(gx),
+                gpu_ptr<DataType>(gw_temp),
                 eps_,
                 axis_size,
                 w_stride);

mlx/backend/cuda/rope.cu

@@ -250,6 +250,7 @@ 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];
@@ -291,14 +292,14 @@ void RoPE::eval_gpu(
       donated = true;
       out.copy_shared_buffer(in);
     } else {
-      out.set_data(allocator::malloc(out.nbytes()));
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
     }
     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(allocator::malloc(out.nbytes()));
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
     strides[0] = in.strides()[ndim - 3];
     strides[1] = in.strides()[ndim - 2];
     strides[2] = in.strides()[ndim - 1];
@@ -319,7 +320,6 @@ 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,
-              (donated ? out : in).data<DataType>(),
-              out.data<DataType>(),
-              offset.data<int32_t>(),
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
               scale_,
               std::log2(base_),
               mat_size,
@@ -357,10 +357,10 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              (donated ? out : in).data<DataType>(),
-              out.data<DataType>(),
-              offset.data<int32_t>(),
-              inputs[2].data<float>(),
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              gpu_ptr<float>(inputs[2]),
               scale_,
               mat_size,
               dims,
@@ -381,10 +381,10 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              (donated ? out : in).data<DataType>(),
-              out.data<DataType>(),
-              offset.data<int32_t>(),
-              inputs[2].data<float>(),
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              gpu_ptr<float>(inputs[2]),
               scale_,
               std::log2(base_),
               strides,
@@ -408,9 +408,9 @@ void RoPE::eval_gpu(
               grid,
               block,
               0,
-              (donated ? out : in).data<DataType>(),
-              out.data<DataType>(),
-              offset.data<int32_t>(),
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
               scale_,
               std::log2(base_),
               strides,

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -0,0 +1,537 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast_primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace fe = cudnn_frontend;
+
+namespace {
+
+#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
+  do {                                                               \
+    auto error = cmd;                                                \
+    if (!error.is_good()) {                                          \
+      throw std::runtime_error(                                      \
+          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
+    }                                                                \
+  } while (0)
+
+std::vector<int64_t> normalized_strides(const array& x) {
+  std::vector<int64_t> strides(x.strides().begin(), x.strides().end());
+  if (std::all_of(
+          strides.begin(), strides.end(), [](int64_t s) { return s == 0; })) {
+    strides.back() = 1;
+    return strides;
+  }
+  if (!x.flags().row_contiguous || x.ndim() < 2) {
+    return strides;
+  }
+  for (int i = x.ndim() - 2; i >= 0; --i) {
+    if (x.shape(i) == 1) {
+      strides[i] = x.shape(i + 1) * strides[i + 1];
+    }
+  }
+  return strides;
+}
+
+void set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  tensor->set_uid(uid)
+      .set_dim({x.shape().begin(), x.shape().end()})
+      .set_stride(normalized_strides(x));
+}
+
+array prepare_sdpa_input(const array& x, Stream s) {
+  // SDPA kernel's requirements on inputs:
+  // 1. last dim's stride be 1;
+  // 2. pointer be aligned.
+  if (x.strides(-1) != 1 || get_alignment(x) < 16) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    auto& encoder = cu::get_command_encoder(s);
+    encoder.add_temporary(x_copy);
+    return x_copy;
+  }
+  return x;
+}
+
+constexpr int QKV_NDIM = 4;
+
+struct SDPACacheKey {
+  int device_id;
+  cudnnDataType_t cudnn_dtype;
+  std::array<int, QKV_NDIM> q_shape;
+  std::array<int, QKV_NDIM> k_shape;
+  std::array<int, QKV_NDIM> v_shape;
+  std::array<int64_t, QKV_NDIM> q_strides;
+  std::array<int64_t, QKV_NDIM> k_strides;
+  std::array<int64_t, QKV_NDIM> v_strides;
+  bool do_causal;
+  bool output_logsumexp;
+};
+
+inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    bool output_logsumexp = true) {
+  BytesKey<SDPACacheKey> cache_key;
+  cache_key.pod = {
+      encoder.device().cuda_device(),
+      dtype_to_cudnn_type(q.dtype()),
+      vector_key<QKV_NDIM>(q.shape()),
+      vector_key<QKV_NDIM>(k.shape()),
+      vector_key<QKV_NDIM>(v.shape()),
+      vector_key<QKV_NDIM>(q.strides()),
+      vector_key<QKV_NDIM>(k.strides()),
+      vector_key<QKV_NDIM>(v.strides()),
+      do_causal,
+      output_logsumexp,
+  };
+  return cache_key;
+}
+
+auto& sdpa_cache() {
+  static LRUBytesKeyCache<SDPACacheKey, fe::graph::Graph> cache(
+      "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 16);
+  return cache;
+}
+
+auto& sdpa_backward_cache() {
+  static LRUBytesKeyCache<SDPACacheKey, fe::graph::Graph> cache(
+      "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 16);
+  return cache;
+}
+
+enum UIDS {
+  Q,
+  K,
+  V,
+  SCALE,
+  O,
+  STATS,
+  // Backward graph:
+  D_Q,
+  D_K,
+  D_V,
+  D_O,
+};
+
+fe::graph::Graph build_sdpa_graph(
+    cudnnHandle_t handle,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    bool output_logsumexp,
+    const array& o,
+    const array& stats) {
+  auto dtype = fe::DataType_t::HALF;
+  if (q.dtype() == bfloat16) {
+    dtype = fe::DataType_t::BFLOAT16;
+  }
+
+  fe::graph::Graph graph;
+  graph.set_io_data_type(dtype)
+      .set_intermediate_data_type(fe::DataType_t::FLOAT)
+      .set_compute_data_type(fe::DataType_t::FLOAT);
+
+  auto q_ = graph.tensor(fe::graph::Tensor_attributes().set_name("Q"));
+  auto k_ = graph.tensor(fe::graph::Tensor_attributes().set_name("K"));
+  auto v_ = graph.tensor(fe::graph::Tensor_attributes().set_name("V"));
+  set_tensor_attrs(q_, Q, q);
+  set_tensor_attrs(k_, K, k);
+  set_tensor_attrs(v_, V, v);
+
+  auto scale = graph.tensor(fe::graph::Tensor_attributes()
+                                .set_name("Scale")
+                                .set_uid(SCALE)
+                                .set_dim({1, 1, 1, 1})
+                                .set_stride({1, 1, 1, 1})
+                                .set_is_pass_by_value(true)
+                                .set_data_type(fe::DataType_t::FLOAT));
+
+  auto options = fe::graph::SDPA_attributes()
+                     .set_name("sdpa_cudnn")
+                     .set_attn_scale(scale)
+                     .set_causal_mask(do_causal)
+                     .set_generate_stats(output_logsumexp);
+
+  auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
+  o_->set_output(true);
+  set_tensor_attrs(o_, O, o);
+  if (output_logsumexp) {
+    stats_->set_output(true).set_data_type(fe::DataType_t::FLOAT);
+    set_tensor_attrs(stats_, STATS, stats);
+  }
+
+  CHECK_CUDNN_FE_ERROR(graph.validate());
+  CHECK_CUDNN_FE_ERROR(graph.build_operation_graph(handle));
+  CHECK_CUDNN_FE_ERROR(graph.create_execution_plans({fe::HeurMode_t::A}));
+  graph.select_behavior_notes(
+      {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+  CHECK_CUDNN_FE_ERROR(graph.check_support(handle));
+  CHECK_CUDNN_FE_ERROR(graph.build_plans(handle));
+
+  return graph;
+}
+
+fe::graph::Graph build_sdpa_backward_graph(
+    cudnnHandle_t handle,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    const array& o,
+    const array& d_o,
+    const array& stats,
+    array& d_q,
+    array& d_k,
+    array& d_v) {
+  auto dtype = fe::DataType_t::HALF;
+  if (q.dtype() == bfloat16) {
+    dtype = fe::DataType_t::BFLOAT16;
+  }
+
+  fe::graph::Graph graph;
+  graph.set_io_data_type(dtype)
+      .set_intermediate_data_type(fe::DataType_t::FLOAT)
+      .set_compute_data_type(fe::DataType_t::FLOAT);
+
+  auto q_ = graph.tensor(fe::graph::Tensor_attributes().set_name("Q"));
+  auto k_ = graph.tensor(fe::graph::Tensor_attributes().set_name("K"));
+  auto v_ = graph.tensor(fe::graph::Tensor_attributes().set_name("V"));
+  auto o_ = graph.tensor(fe::graph::Tensor_attributes().set_name("O"));
+  auto d_o_ = graph.tensor(fe::graph::Tensor_attributes().set_name("D_O"));
+  auto stats_ = graph.tensor(fe::graph::Tensor_attributes().set_name("STATS"));
+  set_tensor_attrs(q_, Q, q);
+  set_tensor_attrs(k_, K, k);
+  set_tensor_attrs(v_, V, v);
+  set_tensor_attrs(o_, O, o);
+  set_tensor_attrs(d_o_, D_O, d_o);
+  set_tensor_attrs(stats_, STATS, stats);
+  stats_->set_data_type(fe::DataType_t::FLOAT);
+
+  auto scale = graph.tensor(fe::graph::Tensor_attributes()
+                                .set_name("Scale")
+                                .set_uid(SCALE)
+                                .set_dim({1, 1, 1, 1})
+                                .set_stride({1, 1, 1, 1})
+                                .set_is_pass_by_value(true)
+                                .set_data_type(fe::DataType_t::FLOAT));
+
+  auto options = fe::graph::SDPA_backward_attributes()
+                     .set_name("sdpa_backward_cudnn")
+                     .set_attn_scale(scale)
+                     .set_causal_mask(do_causal);
+
+  auto [d_q_, d_k_, d_v_] =
+      graph.sdpa_backward(q_, k_, v_, o_, d_o_, stats_, options);
+  d_q_->set_output(true);
+  d_k_->set_output(true);
+  d_v_->set_output(true);
+  set_tensor_attrs(d_q_, D_Q, d_q);
+  set_tensor_attrs(d_k_, D_K, d_k);
+  set_tensor_attrs(d_v_, D_V, d_v);
+
+  CHECK_CUDNN_FE_ERROR(graph.validate());
+  CHECK_CUDNN_FE_ERROR(graph.build_operation_graph(handle));
+  CHECK_CUDNN_FE_ERROR(graph.create_execution_plans({fe::HeurMode_t::A}));
+  graph.select_behavior_notes(
+      {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+  CHECK_CUDNN_FE_ERROR(graph.check_support(handle));
+  CHECK_CUDNN_FE_ERROR(graph.build_plans(handle));
+
+  return graph;
+}
+
+void execute_graph(
+    cu::CommandEncoder& encoder,
+    cudnnHandle_t handle,
+    fe::graph::Graph& graph,
+    std::unordered_map<int64_t, void*>& variant_pack) {
+  int64_t workspace_size = 0;
+  CHECK_CUDNN_FE_ERROR(graph.get_workspace_size(workspace_size));
+  void* workspace_ptr = nullptr;
+  if (workspace_size > 0) {
+    array workspace(
+        cu::malloc_async(workspace_size, encoder),
+        {static_cast<int>(workspace_size)},
+        uint8);
+    encoder.add_temporary(workspace);
+    workspace_ptr = gpu_ptr<void>(workspace);
+  }
+
+  cudnnSetStream(handle, encoder.stream());
+
+  CudaGraph cuda_graph(encoder.device());
+  CHECK_CUDNN_FE_ERROR(graph.populate_cuda_graph(
+      handle, variant_pack, workspace_ptr, cuda_graph));
+  encoder.add_graph_node(cuda_graph);
+}
+
+} // namespace
+
+bool supports_sdpa_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool do_causal,
+    Stream s) {
+  static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1);
+  if (!enabled) {
+    return false;
+  }
+
+  // cuDNN SDPA requires Ampere and later.
+  if (cu::device(s.device).compute_capability_major() < 8) {
+    return false;
+  }
+
+  if (has_mask) {
+    // TODO: Support array masks.
+    if (!do_causal) {
+      return false;
+    }
+    // FIXME: Causal mask generates wrong results when L_Q != L_K.
+    if (q.shape(2) != k.shape(2)) {
+      return false;
+    }
+  }
+
+  // Only use cuDNN for prefilling and training.
+  if (q.shape(2) != k.shape(2)) {
+    return false;
+  }
+
+  // D_qk and D_v must be a multiple of 8 with maximum value 128.
+  if ((q.shape(-1) % 8 != 0) || (q.shape(-1) > 128) || (v.shape(-1) % 8 != 0) ||
+      (v.shape(-1) > 128)) {
+    return false;
+  }
+
+  Dtype dtype = q.dtype();
+  return dtype == float16 || dtype == bfloat16;
+}
+
+void sdpa_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    array& o,
+    array& stats,
+    bool do_causal,
+    bool output_logsumexp,
+    Stream s) {
+  auto& encoder = cu::get_command_encoder(s);
+  auto handle = encoder.device().cudnn_handle();
+
+  // TODO: Handle donation.
+  // TODO: Make O use same memory layout with Q.
+  o.set_data(cu::malloc_async(o.nbytes(), encoder));
+
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  encoder.set_output_array(o);
+
+  if (output_logsumexp) {
+    stats.set_data(cu::malloc_async(stats.nbytes(), encoder));
+    encoder.set_output_array(stats);
+  }
+
+  // Search cache.
+  auto cache_key =
+      build_sdpa_cache_key(encoder, q, k, v, do_causal, output_logsumexp);
+  auto it = sdpa_cache().find(cache_key);
+  if (it == sdpa_cache().end()) {
+    auto graph = build_sdpa_graph(
+        handle, q, k, v, do_causal, output_logsumexp, o, stats);
+    it = sdpa_cache().emplace(cache_key, std::move(graph)).first;
+  }
+  auto& graph = it->second;
+
+  std::unordered_map<int64_t, void*> variant_pack{
+      {Q, const_cast<void*>(gpu_ptr<void>(q))},
+      {K, const_cast<void*>(gpu_ptr<void>(k))},
+      {V, const_cast<void*>(gpu_ptr<void>(v))},
+      {SCALE, &scale},
+      {O, gpu_ptr<void>(o)}};
+  if (output_logsumexp) {
+    variant_pack[STATS] = gpu_ptr<void>(stats);
+  }
+
+  execute_graph(encoder, handle, graph, variant_pack);
+}
+
+void sdpa_backward_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    const array& o,
+    const array& stats,
+    bool do_causal,
+    const array& d_o,
+    array& d_q,
+    array& d_k,
+    array& d_v,
+    Stream s) {
+  auto& encoder = cu::get_command_encoder(s);
+  auto handle = encoder.device().cudnn_handle();
+
+  // TODO: Handle donation.
+  d_q.set_data(cu::malloc_async(d_q.nbytes(), encoder));
+  d_k.set_data(cu::malloc_async(d_k.nbytes(), encoder));
+  d_v.set_data(cu::malloc_async(d_v.nbytes(), encoder));
+
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  encoder.set_input_array(o);
+  encoder.set_input_array(stats);
+  encoder.set_input_array(d_o);
+  encoder.set_output_array(d_q);
+  encoder.set_output_array(d_k);
+  encoder.set_output_array(d_v);
+
+  // Search cache.
+  auto cache_key = build_sdpa_cache_key(encoder, q, k, v, do_causal);
+  auto it = sdpa_backward_cache().find(cache_key);
+  if (it == sdpa_backward_cache().end()) {
+    auto graph = build_sdpa_backward_graph(
+        handle, q, k, v, do_causal, o, d_o, stats, d_q, d_k, d_v);
+    it = sdpa_backward_cache().emplace(cache_key, std::move(graph)).first;
+  }
+  auto& graph = it->second;
+
+  std::unordered_map<int64_t, void*> variant_pack{
+      {Q, const_cast<void*>(gpu_ptr<void>(q))},
+      {K, const_cast<void*>(gpu_ptr<void>(k))},
+      {V, const_cast<void*>(gpu_ptr<void>(v))},
+      {SCALE, &scale},
+      {O, const_cast<void*>(gpu_ptr<void>(o))},
+      {STATS, const_cast<void*>(gpu_ptr<void>(stats))},
+      {D_O, const_cast<void*>(gpu_ptr<void>(d_o))},
+      {D_Q, gpu_ptr<void>(d_q)},
+      {D_K, gpu_ptr<void>(d_k)},
+      {D_V, gpu_ptr<void>(d_v)}};
+
+  execute_graph(encoder, handle, graph, variant_pack);
+}
+
+// Defined in scaled_dot_product_attention.cu file.
+bool supports_sdpa_vector(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp);
+void sdpa_vector(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks,
+    Stream s);
+
+namespace fast {
+
+bool ScaledDotProductAttention::use_fallback(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool is_training,
+    bool output_logsumexp,
+    Stream s) {
+  if (s.device == Device::cpu) {
+    return true;
+  }
+
+  return !supports_sdpa_vector(
+             q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) &&
+      !supports_sdpa_cudnn(q, k, v, has_mask, do_causal, s);
+}
+
+void ScaledDotProductAttention::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("ScaledDotProductAttention::eval_gpu");
+
+  auto& s = stream();
+
+  array q = prepare_sdpa_input(inputs[0], s);
+  array k = prepare_sdpa_input(inputs[1], s);
+  array v = prepare_sdpa_input(inputs[2], s);
+  auto& out = outputs[0];
+  auto& stats = outputs[1];
+  bool has_mask = inputs.size() - has_sinks_ > 3;
+  bool has_arr_mask = has_mask && !do_causal_;
+
+  if (supports_sdpa_vector(
+          q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) {
+    if (has_sinks_) {
+      sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s);
+    } else {
+      sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s);
+    }
+  } else {
+    sdpa_cudnn(q, k, v, scale_, out, stats, do_causal_, output_logsumexp_, s);
+  }
+}
+
+bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) {
+  // The frontend adds a padding mask when sequence length is not a multiple of
+  // tile size.
+  if (q.shape(2) % 128 != 0) {
+    return true;
+  }
+  return s.device == Device::cpu;
+}
+
+void ScaledDotProductAttentionVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("ScaledDotProductAttentionVJP::eval_gpu");
+
+  auto& s = stream();
+
+  assert(inputs.size() == 6);
+  array q = prepare_sdpa_input(inputs[0], s);
+  array k = prepare_sdpa_input(inputs[1], s);
+  array v = prepare_sdpa_input(inputs[2], s);
+  array o = prepare_sdpa_input(inputs[3], s);
+  array stats = prepare_sdpa_input(inputs[4], s);
+  array d_o = prepare_sdpa_input(inputs[5], s);
+
+  assert(outputs.size() == 3);
+  auto& d_q = outputs[0];
+  auto& d_k = outputs[1];
+  auto& d_v = outputs[2];
+
+  sdpa_backward_cudnn(
+      q, k, v, scale_, o, stats, do_causal_, d_o, d_q, d_k, d_v, s);
+}
+
+} // namespace fast
+
+} // namespace mlx::core

mlx/backend/cuda/scaled_dot_product_attention.cu

@@ -6,10 +6,6 @@
 #include "mlx/backend/cuda/kernel_utils.cuh"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/dtype_utils.h"
-#include "mlx/fast_primitives.h"
-#include "mlx/transforms_impl.h"
-
-#include <nvtx3/nvtx3.hpp>
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
@@ -513,11 +509,11 @@ void sdpa_vector_1pass_fallback(
             grid_dim,
             block_dim,
             0,
-            q.data<DataType>(),
-            k.data<DataType>(),
-            v.data<DataType>(),
-            o.data<DataType>(),
-            sinks ? (*sinks).data<DataType>() : nullptr,
+            gpu_ptr<DataType>(q),
+            gpu_ptr<DataType>(k),
+            gpu_ptr<DataType>(v),
+            gpu_ptr<DataType>(o),
+            sinks ? gpu_ptr<DataType>(*sinks) : nullptr,
             params);
       });
     });
@@ -565,9 +561,9 @@ void sdpa_vector_2pass_fallback(
   array sums(intermediate_shape, float32, nullptr, {});
   array maxs(std::move(intermediate_shape), float32, nullptr, {});
 
-  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
-  sums.set_data(allocator::malloc(sums.nbytes()));
-  maxs.set_data(allocator::malloc(maxs.nbytes()));
+  intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder));
+  sums.set_data(cu::malloc_async(sums.nbytes(), encoder));
+  maxs.set_data(cu::malloc_async(maxs.nbytes(), encoder));
 
   encoder.add_temporary(intermediate);
   encoder.add_temporary(sums);
@@ -601,13 +597,13 @@ void sdpa_vector_2pass_fallback(
               grid_dim,
               block_dim,
               0,
-              q.data<DataType>(),
-              k.data<DataType>(),
-              v.data<DataType>(),
-              sinks ? (*sinks).data<DataType>() : nullptr,
-              intermediate.data<float>(),
-              sums.data<float>(),
-              maxs.data<float>(),
+              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),
               params);
         }
 
@@ -628,10 +624,10 @@ void sdpa_vector_2pass_fallback(
               grid_dim,
               block_dim,
               0,
-              intermediate.data<float>(),
-              sums.data<float>(),
-              maxs.data<float>(),
-              o.data<DataType>(),
+              gpu_ptr<float>(intermediate),
+              gpu_ptr<float>(sums),
+              gpu_ptr<float>(maxs),
+              gpu_ptr<DataType>(o),
               params);
         }
       });
@@ -662,21 +658,16 @@ void sdpa_vector_fallback(
 
 } // namespace
 
-namespace fast {
-
-bool ScaledDotProductAttention::use_fallback(
+bool supports_sdpa_vector(
     const array& q,
     const array& k,
     const array& v,
     bool has_mask,
     bool has_arr_mask,
     bool do_causal,
-    Stream s) {
-  if (detail::in_grad_tracing()) {
-    return true;
-  }
-  if (s.device == Device::cpu) {
-    return true;
+    bool output_logsumexp) {
+  if (output_logsumexp) {
+    return false;
   }
 
   const int value_head_dim = v.shape(-1);
@@ -690,29 +681,24 @@ bool ScaledDotProductAttention::use_fallback(
   const bool supported_vector_config =
       sdpa_supported_head_dim && query_sequence_length < 4;
 
-  const bool supported_config = supported_vector_config;
-
-  return has_arr_mask || !supported_config;
+  return supported_vector_config && !has_arr_mask;
 }
 
-void ScaledDotProductAttention::eval_gpu(
-    const std::vector<array>& inputs,
-    array& out) {
-  nvtx3::scoped_range r("ScaledDotProductAttention::eval_gpu");
-
-  auto& s = stream();
+void sdpa_vector(
+    const array& q_pre,
+    const array& k_pre,
+    const array& v_pre,
+    float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks_pre,
+    Stream s) {
   auto& encoder = cu::get_command_encoder(s);
-
-  auto& q_pre = inputs[0];
-  auto& k_pre = inputs[1];
-  auto& v_pre = inputs[2];
-  auto& o = out;
-
   std::vector<array> copies;
 
   // Define some copy functions to ensure the layout of the inputs is as
   // expected.
-  copies.reserve(inputs.size());
+  copies.reserve(4);
   auto copy_unless = [&copies, &s](
                          auto predicate, const array& arr) -> const array& {
     if (!predicate(arr)) {
@@ -730,8 +716,8 @@ void ScaledDotProductAttention::eval_gpu(
   };
 
   std::optional<array> sinks = std::nullopt;
-  if (has_sinks_) {
-    sinks = copy_unless(is_matrix_contiguous, inputs.back());
+  if (sinks_pre) {
+    sinks = copy_unless(is_matrix_contiguous, sinks_pre.value());
   }
 
   // We are in vector mode ie single query
@@ -787,7 +773,7 @@ void ScaledDotProductAttention::eval_gpu(
       };
 
       o.set_data(
-          allocator::malloc(o.nbytes()),
+          cu::malloc_async(o.nbytes(), encoder),
           o.size(),
           {str_oB, str_oH, str_oL, str_oD},
           flags);
@@ -797,8 +783,7 @@ void ScaledDotProductAttention::eval_gpu(
       encoder.add_temporary(cp);
     }
 
-    return sdpa_vector_fallback(
-        s, encoder, q, k, v, scale_, o, do_causal_, sinks);
+    sdpa_vector_fallback(s, encoder, q, k, v, scale, o, do_causal, sinks);
   }
 
   // Full attention mode should never reach here
@@ -807,6 +792,4 @@ void ScaledDotProductAttention::eval_gpu(
   }
 }
 
-} // namespace fast
-
 } // namespace mlx::core

mlx/backend/cuda/scan.cu

@@ -367,13 +367,14 @@ 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(
-          allocator::malloc(in.data_size() * out.itemsize()),
+          cu::malloc_async(in.data_size() * out.itemsize(), encoder),
           in.data_size(),
           in.strides(),
           in.flags());
@@ -387,7 +388,6 @@ 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,
-                  in.data<T>(),
-                  out.data<U>(),
+                  gpu_ptr<T>(in),
+                  gpu_ptr<U>(out),
                   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,
-                  in.data<T>(),
-                  out.data<U>(),
+                  gpu_ptr<T>(in),
+                  gpu_ptr<U>(out),
                   axis_size,
                   stride,
                   stride_blocks);

mlx/backend/cuda/slicing.cpp

@@ -23,14 +23,15 @@ void concatenate_gpu(
   }
   std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
 
-  out.set_data(allocator::malloc(out.nbytes()));
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
 
   auto strides = out.strides();
   auto flags = out.flags();
   flags.row_contiguous = false;
   flags.col_contiguous = false;
   flags.contiguous = false;
-  auto concurrent = cu::get_command_encoder(s).concurrent_context();
+  auto concurrent = encoder.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];
@@ -80,6 +81,7 @@ 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() &&
@@ -87,10 +89,9 @@ array compute_dynamic_offset(
   if (donate) {
     offset.copy_shared_buffer(indices);
   } else {
-    offset.set_data(allocator::malloc(offset.itemsize()));
+    offset.set_data(cu::malloc_async(offset.itemsize(), encoder));
   }
 
-  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,15 +109,16 @@ 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](const array& x) {
+  auto set_output = [&s, &out, &encoder](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(
-            allocator::malloc(x.data_size() * x.itemsize()),
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
             x.data_size(),
             x.strides(),
             x.flags());
@@ -136,7 +137,6 @@ 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,
-          in.data<DataType>(),
-          out.data<DataType>(),
+          gpu_ptr<DataType>(in),
+          gpu_ptr<DataType>(out),
           axis_size);
     });
   });

mlx/backend/cuda/sort.cu

@@ -49,11 +49,12 @@ 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(allocator::malloc(out.nbytes()), in.shape(), out.dtype());
+    out =
+        array(cu::malloc_async(out.nbytes(), encoder), in.shape(), out.dtype());
     encoder.add_temporary(out);
   } else {
     out.set_data(
-        allocator::malloc(in.data_size() * out.itemsize()),
+        cu::malloc_async(in.data_size() * out.itemsize(), encoder),
         in.data_size(),
         in.strides(),
         in.flags());
@@ -70,22 +71,24 @@ 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(allocator::malloc(out.nbytes()), in.shape(), out.dtype());
+        array indices(
+            cu::malloc_async(out.nbytes(), encoder), 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(allocator::malloc(in.nbytes()), in.shape(), in.dtype());
+        array discard(
+            cu::malloc_async(in.nbytes(), encoder), in.shape(), in.dtype());
         encoder.add_temporary(discard);
 
         size_t size;
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortPairs(
             nullptr,
             size,
-            in.data<Type>(),
-            discard.data<Type>(),
-            indices.data<uint32_t>(),
-            out.data<uint32_t>(),
+            gpu_ptr<Type>(in),
+            gpu_ptr<Type>(discard),
+            gpu_ptr<uint32_t>(indices),
+            gpu_ptr<uint32_t>(out),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -94,7 +97,8 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
             sizeof(Type) * 8,
             stream));
 
-        array temp(allocator::malloc(size), {static_cast<int>(size)}, uint8);
+        array temp(
+            cu::malloc_async(size, encoder), {static_cast<int>(size)}, uint8);
         encoder.add_temporary(temp);
 
         // Start capturing after allocations
@@ -103,16 +107,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(indices.data<uint32_t>()),
+            thrust::device_pointer_cast(gpu_ptr<uint32_t>(indices)),
             ModOp<uint32_t>{static_cast<uint32_t>(nsort)});
 
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortPairs(
-            temp.data<void>(),
+            gpu_ptr<void>(temp),
             size,
-            in.data<Type>(),
-            discard.data<Type>(),
-            indices.data<uint32_t>(),
-            out.data<uint32_t>(),
+            gpu_ptr<Type>(in),
+            gpu_ptr<Type>(discard),
+            gpu_ptr<uint32_t>(indices),
+            gpu_ptr<uint32_t>(out),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -125,8 +129,8 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortKeys(
             nullptr,
             size,
-            in.data<Type>(),
-            out.data<Type>(),
+            gpu_ptr<Type>(in),
+            gpu_ptr<Type>(out),
             in.data_size(),
             in.data_size() / nsort,
             offsets,
@@ -135,16 +139,17 @@ void gpu_sort(const Stream& s, array in, array& out_, int axis, bool argsort) {
             sizeof(Type) * 8,
             stream));
 
-        array temp(allocator::malloc(size), {static_cast<int>(size)}, uint8);
+        array temp(
+            cu::malloc_async(size, encoder), {static_cast<int>(size)}, uint8);
         encoder.add_temporary(temp);
 
         // Start capturing after allocations
         auto capture = encoder.capture_context();
         CHECK_CUDA_ERROR(cub::DeviceSegmentedRadixSort::SortKeys(
-            temp.data<void>(),
+            gpu_ptr<void>(temp),
             size,
-            in.data<Type>(),
-            out.data<Type>(),
+            gpu_ptr<Type>(in),
+            gpu_ptr<Type>(out),
             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,
-            a.data<bool>(),
-            b.data<DType>(),
-            c.data<DType>(),
-            out.data<DType>(),
+            gpu_ptr<bool>(a),
+            gpu_ptr<DType>(b),
+            gpu_ptr<DType>(c),
+            gpu_ptr<DType>(out),
             out.data_size());
       });
     } else {
@@ -211,10 +211,10 @@ void ternary_op_gpu_inplace(
                     {num_blocks_x, num_blocks_y},
                     block_dims,
                     0,
-                    a.data<bool>(),
-                    b.data<DType>(),
-                    c.data<DType>(),
-                    out.data<DType>(),
+                    gpu_ptr<bool>(a),
+                    gpu_ptr<DType>(b),
+                    gpu_ptr<DType>(c),
+                    gpu_ptr<DType>(out),
                     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,
-                  a.data<bool>(),
-                  b.data<DType>(),
-                  c.data<DType>(),
-                  out.data<DType>(),
+                  gpu_ptr<bool>(a),
+                  gpu_ptr<DType>(b),
+                  gpu_ptr<DType>(c),
+                  gpu_ptr<DType>(out),
                   rest,
                   const_param(shape),
                   const_param(a_strides),
@@ -256,7 +256,9 @@ void ternary_op_gpu(
   auto& b = inputs[1];
   auto& c = inputs[2];
   auto topt = get_ternary_op_type(a, b, c);
-  set_ternary_op_output_data(a, b, c, out, topt);
+  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); });
   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,
-                in.data<InType>(),
-                out.data<OutType>(),
+                gpu_ptr<InType>(in),
+                gpu_ptr<OutType>(out),
                 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,
-                in.data<InType>(),
-                out.data<OutType>(),
+                gpu_ptr<InType>(in),
+                gpu_ptr<OutType>(out),
                 rest,
                 const_param(shape),
                 const_param(strides),
@@ -207,7 +207,9 @@ void unary_op_gpu(
     array& out,
     const char* op,
     const Stream& s) {
-  set_unary_output_data(inputs[0], out);
+  auto& encoder = cu::get_command_encoder(s);
+  set_unary_output_data(
+      inputs[0], out, [&](auto n) { return cu::malloc_async(n, encoder); });
   unary_op_gpu_inplace<Op>(inputs, out, op, s);
 }
 

mlx/backend/cuda/utils.cpp

@@ -60,7 +60,7 @@ const char* dtype_to_cuda_type(const Dtype& dtype) {
     case float64:
       return "double";
     case complex64:
-      return "complex64_t";
+      return "mlx::core::cu::complex64_t";
     default:
       return "unknown";
   }