Fix edge checking in attention mask

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

@@ -18,7 +18,7 @@ runs:
       env:
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
       run: |
-        pip install auditwheel "build<=1.4.2" patchelf setuptools
+        pip install auditwheel build patchelf setuptools
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
 

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

@@ -25,7 +25,7 @@ runs:
     - name: Build Python package
       shell: bash
       run: |
-        pip install auditwheel patchelf "build<=1.4.2"
+        pip install auditwheel patchelf build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
         auditwheel repair dist/mlx-*.whl \

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

@@ -9,7 +9,6 @@ inputs:
 runs:
   using: "composite"
   steps:
-
     - name: Install Python package
       id: python_build
       shell: sh
@@ -21,7 +20,7 @@ runs:
       run: |
         if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
           # There is no GPU in arm64 runner, use a common arch.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
           # Can not build tests and stubs when the built executables can not run.
           CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
         fi

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

@@ -18,7 +18,6 @@ runs:
     - name: Build Python package
       shell: bash -l {0}
       env:
-        DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
         pip install build
@@ -29,7 +28,6 @@ runs:
       if: ${{ inputs.build-backend }}
       shell: bash -l {0}
       env:
-        DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
         python setup.py clean --all

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

@@ -12,20 +12,20 @@ runs:
       run: |
         pip install --upgrade pip
         pip install cmake setuptools typing_extensions
-        pip install -e ".[dev]" -v
+        pip install -e . -v
 
     - name: Install tests dependencies
       shell: bash -l {0}
       run: |
-        pip install tensorflow
+        pip install numpy torch tensorflow unittest-xml-reporting
 
     - name: Run Python tests
       shell: bash -l {0}
       env:
         LOW_MEMORY: 1
       run: |
-        DEVICE=cpu python -m unittest discover -v python/tests
-        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
+        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
@@ -45,17 +45,15 @@ runs:
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
-
+    
     - name: Run CPP tests
       shell: bash -l {0}
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
-      run: |
-        ./build/tests/tests
-        ./build/tests/test_teardown
-
+      run: ./build/tests/tests
+    
     - name: Build small binary with JIT
       shell: bash -l {0}
       run: |
@@ -79,4 +77,6 @@ runs:
       run: |
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
           pip install -e . -v
-        python -m unittest discover -v python/tests
+        python -m xmlrunner discover \
+            -v python/tests \
+            -o test-results/gpu_jit

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

@@ -1,26 +0,0 @@
-name: 'Build on Windows'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Install Python package
-      id: python-build
-      shell: cmd
-      env:
-        # For MSVC, Ninja/Release is the only config supported by ccache.
-        CMAKE_ARGS: >-
-          -G Ninja
-          -DCMAKE_BUILD_TYPE=Release
-          -DCMAKE_C_COMPILER=cl
-          -DCMAKE_CXX_COMPILER=cl
-          -DCMAKE_RC_COMPILER=rc
-      run: |
-        uv pip install ".[dev]" -v
-        :: Pass the CMAKE_ARGS to following steps.
-        >>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS%
-
-    - name: Build CPP only
-      shell: cmd
-      run: |
-        cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }}
-        cmake --build build -j %NUMBER_OF_PROCESSORS%

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

@@ -14,9 +14,6 @@ inputs:
     description: 'Whether to enable ccache'
     required: false
     default: 'true'
-  ccache-key:
-    required: false
-    default: 'ccache'
 
 runs:
   using: "composite"
@@ -30,13 +27,14 @@ runs:
             zip \
             libblas-dev liblapack-dev liblapacke-dev \
             openmpi-bin openmpi-common libopenmpi-dev
+        sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc || true
         echo "::endgroup::"
 
     - name: Use ccache
       if: ${{ inputs.use-ccache == 'true' }}
       uses: hendrikmuhs/ccache-action@v1.2
       with:
-        key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
+        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
         max-size: 1GB
         # ccache-action bug: running "apt-get update" fails on large arm runner.
         update-package-index: false
@@ -57,12 +55,6 @@ runs:
         echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
         echo "::endgroup::"
 
-    - name: Set swap space
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
-      with:
-        swap-size-gb: 16
-
     - name: Install CUDA toolkit
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       shell: bash

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

@@ -1,42 +0,0 @@
-name: 'Setup Windows environment'
-
-inputs:
-  python-version:
-    description: 'Version of python to set up'
-    required: false
-    default: '3.14'
-  use-ccache:
-    description: 'Whether to enable ccache'
-    required: false
-    default: 'true'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Use ccache
-      if: ${{ inputs.use-ccache == 'true' }}
-      uses: hendrikmuhs/ccache-action@v1.2
-      with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-cpu
-        max-size: 1GB
-
-    - name: Setup Visual Studio cmd
-      shell: cmd
-      run: |
-        :: Find out path to VS.
-        pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\"
-        for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x
-        popd
-        :: Import VS vars.
-        call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" x64
-        :: Export to all steps.
-        >>%GITHUB_ENV% set
-
-    - uses: astral-sh/setup-uv@v7
-
-    - name: Setup Python venv
-      shell: cmd
-      run: |
-        uv venv --python ${{ inputs.python-version }}
-        call ".venv/Scripts/activate.bat"
-        >>%GITHUB_ENV% set

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

@@ -65,5 +65,5 @@ runs:
         DEVICE: gpu
       run: |
         echo "::group::CPP tests - GPU"
-        ./build/tests/tests -sfe="*linalg_tests.cpp"
+        ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
         echo "::endgroup::"

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

@@ -1,21 +0,0 @@
-name: 'Run tests on Windows'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Run Python tests - CPU
-      shell: bash
-      run: |
-        echo "::group::Python tests - CPU"
-        python -m unittest discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run CPP tests - CPU
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::CPP tests - CPU"
-        ./build/tests.exe -tce="*gguf*,test random uniform"
-        ./build/test_teardown.exe
-        echo "::endgroup::"

.github/workflows/build-wheels.yml

@@ -1,94 +0,0 @@
-name: Build macOS arm64 wheels
-
-on:
-  push:
-    branches:
-      - main
-      - 'metal-*'
-      - 'q-*'
-      - attn-mask-fix
-      - fix-rope
-  workflow_dispatch:
-    inputs:
-      branch_to_build:
-        description: 'Branch to build (optional, defaults to current ref)'
-        required: false
-        default: ''
-
-concurrency:
-  group: build-${{ github.ref }}-${{ github.event.inputs.branch_to_build }}
-  cancel-in-progress: true
-
-jobs:
-  build:
-    name: Build wheel (Python ${{ matrix.python }})
-    runs-on: macos-14
-    timeout-minutes: 60
-    strategy:
-      fail-fast: false
-      matrix:
-        python: ['3.11', '3.12']
-    env:
-      CMAKE_BUILD_PARALLEL_LEVEL: '4'
-    steps:
-      - name: Determine target branch
-        id: branch
-        run: |
-          NAME="${{ github.event.inputs.branch_to_build }}"
-          if [ -z "$NAME" ]; then
-            NAME="${{ github.ref_name }}"
-          fi
-          # Sanitize for artifact naming (replace / with -)
-          SAFE_NAME=$(echo "$NAME" | tr '/' '-')
-          echo "name=$NAME" >> $GITHUB_OUTPUT
-          echo "safe_name=$SAFE_NAME" >> $GITHUB_OUTPUT
-          echo "Target branch: $NAME (safe: $SAFE_NAME)"
-
-      - name: Checkout
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ steps.branch.outputs.name }}
-          submodules: recursive
-
-      - name: Set up Python ${{ matrix.python }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python }}
-
-      - name: Cache pip
-        uses: actions/cache@v4
-        with:
-          path: |
-            ~/.cache/pip
-            ~/Library/Caches/pip
-          key: pip-${{ runner.os }}-py${{ matrix.python }}-${{ hashFiles('CMakeLists.txt', 'setup.py', 'pyproject.toml') }}
-          restore-keys: |
-            pip-${{ runner.os }}-py${{ matrix.python }}-
-
-      - name: Install build dependencies
-        run: |
-          python -m pip install -U pip wheel build setuptools cmake nanobind
-
-      - name: Build wheel
-        run: |
-          mkdir -p ./wheels
-          pip wheel --no-deps . -w ./wheels
-
-      - name: List built wheels
-        run: ls -lh ./wheels
-
-      - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: mlx-${{ steps.branch.outputs.safe_name }}-py${{ matrix.python }}-wheels
-          path: ./wheels/*.whl
-          retention-days: 30
-          if-no-files-found: error
-
-      - name: Create GitHub Release (on tag)
-        if: startsWith(github.ref, 'refs/tags/v')
-        uses: softprops/action-gh-release@v2
-        with:
-          files: ./wheels/*.whl
-          fail_on_unmatched_files: false
-          generate_release_notes: true

.github/workflows/build_and_test.yml

@@ -36,7 +36,6 @@ jobs:
       - uses: ./.github/actions/setup-linux
       - uses: ./.github/actions/build-linux
       - uses: ./.github/actions/test-linux
-      - run: df -h
 
   cuda_build_and_test:
     name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }})
@@ -76,16 +75,6 @@ jobs:
       - uses: ./.github/actions/setup-macos
       - uses: ./.github/actions/build-macos
 
-  windows_build_and_test:
-    name: Windows (cpu, x86_64)
-    needs: check_lint
-    runs-on: windows-2025
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-windows
-      - uses: ./.github/actions/build-windows
-      - uses: ./.github/actions/test-windows
-
   build_documentation:
     name: Build Documentation
     if: github.repository == 'ml-explore/mlx'

.github/workflows/documentation.yml

@@ -25,4 +25,4 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v5
+        uses: actions/deploy-pages@v4

.github/workflows/nightly.yml

@@ -23,19 +23,18 @@ jobs:
           build-backend: ${{ matrix.python-version == '3.10' }}
           arch: "x86_64"
       - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           name: linux-wheels-${{ matrix.python_version }}
           path: wheelhouse/mlx-*.whl
           retention-days: 7
       - name: Upload mlx-cpu artifacts
         if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           name: mlx-cpu
           path: wheelhouse/mlx_cpu-*.whl
           retention-days: 7
-      - run: df -h
 
   build_linux_with_tests:
     strategy:
@@ -53,7 +52,6 @@ jobs:
           python-version: ${{ matrix.python_version }}
       - uses: ./.github/actions/build-linux
       - uses: ./.github/actions/test-linux
-      - run: df -h
 
   build_mac_release:
     if: github.repository == 'ml-explore/mlx'
@@ -85,24 +83,20 @@ jobs:
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        arch: ['x86_64', 'aarch64']
-        toolkit: ['cuda-12.9', 'cuda-13.0']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
+    runs-on: ubuntu-22-large
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
-          toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
+          toolkit: 'cuda-12.9'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          arch: ${{ matrix.arch }}
+          toolkit: 'cuda-12.9'
+          arch: 'x86_64'
       - name: Upload artifacts
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
-          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          name: mlx-cuda
           path: wheelhouse/mlx_cuda_*.whl
           retention-days: 7

.github/workflows/release.yml

@@ -4,18 +4,18 @@ on:
   push:
     tags:
       - 'v*'
-    branches:
-      - 'test-publish/*'
   workflow_dispatch:
     inputs:
-      dry_run:
-        description: 'Dry run (do not publish to PyPi)'
+      publish:
+        description: 'Publish to PyPI (uncheck for dry run)'
         required: false
         type: boolean
+        default: true
       dev_release:
         description: 'Development release (DEV_RELEASE=1)'
         required: false
         type: boolean
+        default: false
 
 permissions:
   contents: read
@@ -29,7 +29,7 @@ jobs:
       - uses: ./.github/actions/build-docs
 
   deploy_documentation:
-    if: ${{ !inputs.dry_run }}
+    if: ${{ github.event_name == 'push' || inputs.publish }}
     needs: build_documentation
     permissions:
       pages: write
@@ -41,7 +41,7 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v5
+        uses: actions/deploy-pages@v4
 
   build_linux_release:
     if: github.repository == 'ml-explore/mlx'
@@ -64,7 +64,7 @@ jobs:
           build-backend: ${{ matrix.python_version == '3.10' }}
           arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
           name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
@@ -72,7 +72,7 @@ jobs:
           if-no-files-found: error
       - name: Upload CPU artifacts
         if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
           name: mlx-cpu-${{ matrix.arch }}
@@ -110,13 +110,8 @@ jobs:
         with:
           macos-target: 15.0
           build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 26 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 26.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
           name: mac-wheels-${{ matrix.python-version }}
@@ -124,7 +119,7 @@ jobs:
           if-no-files-found: error
       - name: Upload Metal artifacts
         if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
           name: mlx-metal
@@ -146,13 +141,13 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
+          use-ccache: false
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
           arch: ${{ matrix.arch }}
       - name: Upload artifacts
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
           name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
@@ -166,15 +161,15 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      name: pypi
       url: https://pypi.org/p/mlx
     steps:
-      - uses: actions/download-artifact@v8
+      - uses: actions/download-artifact@v7
         with:
           pattern: linux-wheels-*
           merge-multiple: true
           path: dist
-      - uses: actions/download-artifact@v8
+      - uses: actions/download-artifact@v7
         with:
           pattern: mac-wheels-*
           merge-multiple: true
@@ -182,7 +177,7 @@ jobs:
       - name: Display structure of downloaded files
         run: du -ah dist
       - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
+        if: ${{ github.event_name == 'push' || inputs.publish }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -194,10 +189,10 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      name: pypi
       url: https://pypi.org/p/mlx-cuda
     steps:
-      - uses: actions/download-artifact@v8
+      - uses: actions/download-artifact@v7
         with:
           pattern: mlx-cuda-*
           merge-multiple: true
@@ -205,7 +200,7 @@ jobs:
       - name: Display structure of downloaded files
         run: du -ah dist
       - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
+        if: ${{ github.event_name == 'push' || inputs.publish }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -217,10 +212,10 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      name: pypi
       url: https://pypi.org/p/mlx-cpu
     steps:
-      - uses: actions/download-artifact@v8
+      - uses: actions/download-artifact@v7
         with:
           pattern: mlx-cpu-*
           merge-multiple: true
@@ -228,7 +223,7 @@ jobs:
       - name: Display structure of downloaded files
         run: du -ah dist
       - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
+        if: ${{ github.event_name == 'push' || inputs.publish }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -240,17 +235,17 @@ jobs:
     permissions:
       id-token: write
     environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      name: pypi
       url: https://pypi.org/p/mlx-metal
     steps:
-      - uses: actions/download-artifact@v8
+      - uses: actions/download-artifact@v7
         with:
           name: mlx-metal
           path: dist
       - name: Display structure of downloaded files
         run: du -ah dist
       - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
+        if: ${{ github.event_name == 'push' || inputs.publish }} 
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/

.gitignore

@@ -3,12 +3,16 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+# C extensions
+*.so
+
 # tensor files
 *.safe
 *.safetensors
 
 # Metal libraries
 *.metallib
+venv/
 
 # Distribution / packaging
 python/mlx/core
@@ -26,7 +30,6 @@ lib64/
 parts/
 sdist/
 var/
-venv/
 wheels/
 share/python-wheels/
 *.egg-info/
@@ -34,7 +37,12 @@ share/python-wheels/
 *.egg
 MANIFEST
 uv.lock
-.DS_Store
+
+# vim
+*.swp
+
+# Ignore build dir
+build/
 
 # Prerequisites
 *.d
@@ -44,7 +52,6 @@ uv.lock
 *.lo
 *.o
 *.obj
-*.ilk
 
 # Precompiled Headers
 *.gch
@@ -73,9 +80,9 @@ uv.lock
 # Debug symbols
 *.pdb
 
-# VSCode
+# VSCode 
 .vscode/
+.DS_Store
+
 # Jetbrains
-.cache/
-# vim
-*.swp
+.cache

.pre-commit-config.yaml

@@ -6,17 +6,17 @@ repos:
     # -   id: end-of-file-fixer
     # -   id: trailing-whitespace
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v21.1.8
+    rev: v19.1.7
     hooks:
     -   id: clang-format
 # Using this mirror lets us use mypyc-compiled black, which is about 2x faster
 -   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 26.1.0
+    rev: 25.1.0
     hooks:
     -   id: black
     
 -   repo: https://github.com/pycqa/isort
-    rev: 7.0.0
+    rev: 6.0.0
     hooks:
     -   id: isort
         args:

CMakeLists.txt

@@ -22,7 +22,7 @@ project(
 
 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
-set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
@@ -40,6 +40,7 @@ option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
+option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
@@ -149,17 +150,15 @@ cmake_policy(SET CMP0135 NEW)
 
 add_library(mlx)
 
+# Supress warnings: note: parameter passing for argument of type
+# ‘std::pair<float, float>’ when C++17 is enabled changed to match C++14 in GCC
+# 10.1
+target_compile_options(mlx PRIVATE -Wno-psabi)
 target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
 target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
 
 if(MLX_BUILD_CUDA)
   enable_language(CUDA)
-  find_package(CUDAToolkit REQUIRED)
-  find_package(CUDNN REQUIRED)
-  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.1" AND CUDAToolkit_VERSION
-                                                          VERSION_LESS "13.2")
-    message(FATAL_ERROR "CUDA Toolkit 13.1 is not supported.")
-  endif()
 endif()
 
 if(MLX_BUILD_METAL)
@@ -223,17 +222,14 @@ if(WIN32)
   if(MSVC)
     # GGUF does not build with MSVC.
     set(MLX_BUILD_GGUF OFF)
-  endif()
-  # Generate DLL and EXE in the same dir, otherwise EXE will not be able to run.
-  # This is only done when MLX is built as the top project.
-  if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
-    set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
+    # There is no prebuilt OpenBLAS distribution for MSVC.
+    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
   endif()
   # Windows implementation of dlfcn.h APIs.
   FetchContent_Declare(
     dlfcn-win32
     GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git
-    GIT_TAG v1.4.2
+    GIT_TAG v1.4.1
     EXCLUDE_FROM_ALL)
   block()
   set(BUILD_SHARED_LIBS OFF)
@@ -257,25 +253,20 @@ if(MLX_BUILD_CPU)
     target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
     add_compile_definitions(MLX_USE_ACCELERATE)
     add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(WIN32)
-    # Download and link prebuilt binaries of OpenBLAS. Note that we can only
-    # link with the dynamic library, the prebuilt binaries were built with MinGW
-    # so static-linking would require linking with MinGW's runtime.
+  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
+    # Download and build OpenBLAS from source code.
     FetchContent_Declare(
       openblas
-      URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip"
-    )
+      GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git
+      GIT_TAG v0.3.28
+      EXCLUDE_FROM_ALL)
+    set(BUILD_STATIC_LIBS ON) # link statically
+    set(NOFORTRAN ON) # msvc has no fortran compiler
     FetchContent_MakeAvailable(openblas)
-    target_link_libraries(mlx
-                          PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib")
-    target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include")
-    # Make sure the DLL file is placed in the same dir with executables.
-    set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll")
-    add_custom_command(
-      TARGET mlx
-      POST_BUILD
-      COMMAND ${CMAKE_COMMAND} -E copy_if_different ${OPENBLAS_DLL_FILE}
-              ${CMAKE_BINARY_DIR})
+    target_link_libraries(mlx PRIVATE openblas)
+    target_include_directories(
+      mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include"
+                  "${CMAKE_BINARY_DIR}/generated" "${CMAKE_BINARY_DIR}")
   else()
     if(${CMAKE_HOST_APPLE})
       # The blas shipped in macOS SDK is not supported, search homebrew for
@@ -321,28 +312,22 @@ FetchContent_MakeAvailable(json)
 target_include_directories(
   mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
 
-# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
-if(MLX_BUILD_CPU
-   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
-   AND DEFINED MACOS_SDK_VERSION
-   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
-  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
-                   ${CMAKE_BINARY_DIR}/jaccl)
-endif()
-
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(
   mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
              $<INSTALL_INTERFACE:include>)
 
+# Do not add mlx_EXPORTS define for shared library.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL "")
+
 if(USE_SYSTEM_FMT)
   find_package(fmt REQUIRED)
 else()
   FetchContent_Declare(
     fmt
     GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-    GIT_TAG 12.1.0
+    GIT_TAG 10.2.1
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(fmt)
 endif()
@@ -357,7 +342,7 @@ if(MLX_BUILD_PYTHON_BINDINGS)
   FetchContent_Declare(
     nanobind
     GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v2.12.0
+    GIT_TAG v2.10.2
     GIT_SHALLOW TRUE
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(nanobind)
@@ -380,15 +365,6 @@ endif()
 # ----------------------------- Installation -----------------------------
 include(GNUInstallDirs)
 
-if(WIN32)
-  # Install DLLs to the same dir with extension file (core.pyd) on Windows.
-  set(CMAKE_INSTALL_BINDIR ".")
-  if(MLX_BUILD_CPU)
-    # Install OpenBLAS.
-    install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN)
-  endif()
-endif()
-
 # Install library
 install(
   TARGETS mlx

benchmarks/python/block_masked_mm_bench.py

@@ -1,193 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-import time
-
-import mlx.core as mx
-import numpy as np
-
-MLX_DTYPES = {
-    "float16": mx.float16,
-    "bfloat16": mx.bfloat16,
-    "float32": mx.float32,
-}
-
-
-def parse_cases(cases):
-    parsed = []
-    for spec in cases.split(","):
-        parts = spec.split("x")
-        m, n, k, bs = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
-        sparsity = float(parts[4]) if len(parts) > 4 else 0.5
-        parsed.append((m, n, k, bs, sparsity))
-    return parsed
-
-
-def make_masks(m, n, k, block_size, sparsity, rng):
-    """Create block masks with given sparsity (fraction of blocks zeroed)."""
-    tm = (m + block_size - 1) // block_size
-    tn = (n + block_size - 1) // block_size
-    tk = (k + block_size - 1) // block_size
-
-    lhs_mask = (rng.random((tm, tk)) >= sparsity).astype(np.bool_)
-    rhs_mask = (rng.random((tk, tn)) >= sparsity).astype(np.bool_)
-    out_mask = (rng.random((tm, tn)) >= sparsity).astype(np.bool_)
-    return lhs_mask, rhs_mask, out_mask
-
-
-def mlx_naive_block_masked_mm(a, b, block_size, out_mask, lhs_mask, rhs_mask):
-    """MLX naive: expand masks and use regular matmul."""
-    M, K = a.shape[-2], a.shape[-1]
-    N = b.shape[-1]
-
-    def expand(mask, rows, cols):
-        e = mx.repeat(mx.repeat(mask, block_size, axis=-2), block_size, axis=-1)
-        return e[..., :rows, :cols]
-
-    a_masked = a * expand(lhs_mask, M, K)
-    b_masked = b * expand(rhs_mask, K, N)
-    c = a_masked @ b_masked
-    c = c * expand(out_mask, M, N)
-    return c
-
-
-def bench_mlx(fn, warmup, iters):
-    for _ in range(warmup):
-        y = fn()
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = fn()
-        mx.eval(y)
-    mx.synchronize()
-    return (time.perf_counter() - start) * 1e3 / iters
-
-
-def print_table(headers, rows):
-    widths = [len(h) for h in headers]
-    for row in rows:
-        for i, cell in enumerate(row):
-            widths[i] = max(widths[i], len(cell))
-
-    def fmt_row(row):
-        return (
-            "| "
-            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
-            + " |"
-        )
-
-    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
-    print(fmt_row(headers))
-    print(sep)
-    for row in rows:
-        print(fmt_row(row))
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Benchmark block_masked_mm vs naive expand+matmul"
-    )
-    parser.add_argument(
-        "--cases",
-        default=(
-            "256x256x256x32x0.5,"
-            "512x512x512x32x0.5,"
-            "1024x1024x1024x32x0.5,"
-            "1024x1024x1024x64x0.5,"
-            "2048x2048x2048x64x0.5,"
-            "256x256x256x32x0.0,"
-            "1024x1024x1024x32x0.0,"
-            "1024x1024x1024x32x0.9"
-        ),
-        help="Comma-separated MxNxKxBSxSparsity list. Sparsity=fraction of blocks zeroed.",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="float32",
-        choices=["float16", "bfloat16", "float32"],
-    )
-    parser.add_argument("--warmup", type=int, default=10)
-    parser.add_argument("--iters", type=int, default=50)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--no-check", action="store_true")
-    args = parser.parse_args()
-
-    mlx_dtype = MLX_DTYPES[args.dtype]
-
-    print(f"dtype={args.dtype}  warmup={args.warmup}  iters={args.iters}")
-
-    headers = [
-        "Case (MxNxKxBS)",
-        "Sparsity",
-        "MLX ms",
-        "Naive ms",
-        "Speedup",
-    ]
-    if not args.no_check:
-        headers.append("Max err")
-    rows = []
-
-    cases = parse_cases(args.cases)
-    for idx, (m, n, k, bs, sparsity) in enumerate(cases):
-        rng = np.random.default_rng(args.seed + idx)
-        a_np = rng.standard_normal((m, k)).astype(np.float32)
-        b_np = rng.standard_normal((k, n)).astype(np.float32)
-        lhs_mask_np, rhs_mask_np, out_mask_np = make_masks(m, n, k, bs, sparsity, rng)
-
-        a_mx = mx.array(a_np, dtype=mlx_dtype)
-        b_mx = mx.array(b_np, dtype=mlx_dtype)
-        lhs_mask_mx = mx.array(lhs_mask_np)
-        rhs_mask_mx = mx.array(rhs_mask_np)
-        out_mask_mx = mx.array(out_mask_np)
-        mx.eval(a_mx, b_mx, lhs_mask_mx, rhs_mask_mx, out_mask_mx)
-
-        # Correctness check: block_masked_mm vs naive expand+matmul
-        err_str = ""
-        if not args.no_check:
-            y_op = mx.block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            )
-            y_naive = mlx_naive_block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            )
-            mx.eval(y_op, y_naive)
-            err = float(mx.max(mx.abs(y_op - y_naive)).item())
-            err_str = f"{err:.2e}"
-
-        # Benchmark
-        t_mlx = bench_mlx(
-            lambda: mx.block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            ),
-            args.warmup,
-            args.iters,
-        )
-        t_naive = bench_mlx(
-            lambda: mlx_naive_block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            ),
-            args.warmup,
-            args.iters,
-        )
-        speedup = f"{t_naive / t_mlx:.2f}x" if t_mlx > 0 else "-"
-
-        row = [
-            f"{m}x{n}x{k}x{bs}",
-            f"{sparsity:.0%}",
-            f"{t_mlx:.3f}",
-            f"{t_naive:.3f}",
-            speedup,
-        ]
-        if not args.no_check:
-            row.append(err_str)
-        rows.append(row)
-
-    print_table(headers, rows)
-    if not args.no_check:
-        print("err: max|block_masked_mm - naive_expand_matmul|")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/python/conv3d_bench.py

@@ -1,152 +0,0 @@
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_warmup = 2
-N_iter_bench = 10
-N_iter_func = 10
-
-
-def bench(f, a, b, b_prime):
-    for i in range(N_warmup):
-        f(a, b, b_prime)
-    torch.mps.synchronize()
-
-    s = time.perf_counter_ns()
-    for i in range(N_iter_bench):
-        f(a, b, b_prime)
-    e = time.perf_counter_ns()
-    return (e - s) * 1e-9
-
-
-def make_mx_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
-    def mx_conv_3D(a, b, b_prime):
-        y = a
-        for i in range(N_iter_func):
-            y = mx.conv3d(y, b, stride=strides, padding=padding, groups=groups)
-            y = mx.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
-        mx.eval(y)
-        return y
-
-    return mx_conv_3D
-
-
-def make_pt_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
-    @torch.no_grad()
-    def pt_conv_3D(a, b, b_prime):
-        y = a
-        for i in range(N_iter_func):
-            y = torch.conv3d(y, b, stride=strides, padding=padding, groups=groups)
-            y = torch.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
-        torch.mps.synchronize()
-        return y
-
-    return pt_conv_3D
-
-
-def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype):
-    scale = 1.0 / math.sqrt(kD * kH * kW * C)
-    a_np = np.random.uniform(0, 0.5, (N, D, H, W, C))
-    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups)))
-    b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(O / groups)))
-
-    a_np, b_np, b_prime_np = map(lambda x: x.astype(np_dtype), (a_np, b_np, b_prime_np))
-    a_mx, b_mx, b_prime_mx = map(lambda x: mx.array(x), (a_np, b_np, b_prime_np))
-    a_pt, b_pt, b_prime_pt = map(
-        lambda x: torch.from_numpy(x.transpose(0, 4, 1, 2, 3)).to("mps"),
-        (a_np, b_np, b_prime_np),
-    )
-
-    torch.mps.synchronize()
-
-    f_mx = make_mx_conv_3D(strides, padding, groups)
-    f_pt = make_pt_conv_3D(strides, padding, groups)
-
-    time_torch = bench(f_pt, a_pt, b_pt, b_prime_pt)
-    time_mlx = bench(f_mx, a_mx, b_mx, b_prime_mx)
-
-    # Measure MLX memory
-    mx.clear_cache()
-    mx.reset_peak_memory()
-    y = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
-    mx.eval(y)
-    mlx_peak_mb = mx.get_peak_memory() / 1024**2
-    mlx_active_mb = mx.get_active_memory() / 1024**2
-    del y
-
-    # Measure PyTorch MPS memory
-    torch.mps.synchronize()
-    torch.mps.empty_cache()
-    y = torch.conv3d(a_pt, b_pt, stride=strides, padding=padding, groups=groups)
-    torch.mps.synchronize()
-    pt_current_mb = torch.mps.current_allocated_memory() / 1024**2
-    pt_driver_mb = torch.mps.driver_allocated_memory() / 1024**2
-    del y
-
-    out_mx = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
-    out_pt = torch.conv3d(
-        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
-    )
-    out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1))
-    out_pt = out_pt.numpy(force=True)
-
-    atol = 2e-5 if np_dtype == np.float32 else 5e-4
-
-    if not np.allclose(out_pt, out_mx, atol=atol):
-        print(
-            f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} "
-            f"[strides = {strides}, padding = {padding}, groups = {groups}] "
-            f"with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
-        )
-
-    return time_mlx, time_torch, mlx_peak_mb, mlx_active_mb, pt_current_mb, pt_driver_mb
-
-
-if __name__ == "__main__":
-    dtypes = ("float16", "float32")
-    shapes = (
-        # (C % 16 == 0)
-        (4, 16, 16, 16, 32, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        # Larger spatial dims
-        (2, 64, 64, 64, 32, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (1, 64, 64, 64, 64, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        # Strided
-        (4, 32, 32, 32, 64, 3, 3, 3, 128, (2, 2, 2), (1, 1, 1), 1),
-        # Asymmetric kernels
-        (4, 32, 32, 32, 64, 3, 1, 1, 128, (1, 1, 1), (1, 0, 0), 1),
-        (4, 32, 32, 32, 64, 1, 3, 3, 128, (1, 1, 1), (0, 1, 1), 1),
-        # (C % 16 != 0)
-        (4, 16, 16, 16, 21, 3, 3, 3, 21, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 3, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
-    )
-
-    for dtype in dtypes:
-        print(f"\n{'=' * 120}" f"\n  dtype: {dtype}" f"\n{'=' * 120}")
-        print(
-            f"{'(N,   D,   H,   W,   C)':<26s} {'(  O, kD, kH, kW,   C)':<24s} "
-            f"{'stride':<12s} {'pads':<12s} {'groups':>6s} "
-            f"{'diff%':>7s}  "
-            f"{'MLX peak':>9s} {'MLX act':>8s} {'PT cur':>8s} {'PT drv':>8s}"
-        )
-        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
-            np_dtype = getattr(np, dtype)
-            time_mlx, time_torch, mlx_peak, mlx_act, pt_cur, pt_drv = bench_shape(
-                N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype
-            )
-            diff = time_torch / time_mlx - 1.0
-
-            print(
-                f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), "
-                f"{strides}, {padding}, {groups:6d}, "
-                f"{100. * diff:+6.1f}%  "
-                f"{mlx_peak:8.1f}  {mlx_act:7.1f}  {pt_cur:7.1f}  {pt_drv:7.1f}"
-            )

benchmarks/python/large_gemm_bench.py

@@ -1,119 +0,0 @@
-# Copyright © 2026 Apple Inc.
-
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_WARMUP = 5
-N_BENCH = 20
-
-
-def bench_mlx(a, b):
-    for _ in range(N_WARMUP):
-        mx.eval(a @ b)
-
-    times = []
-    for _ in range(N_BENCH):
-        start = time.perf_counter_ns()
-        mx.eval(a @ b)
-        end = time.perf_counter_ns()
-        times.append((end - start) * 1e-9)
-
-    return np.mean(times), np.std(times)
-
-
-@torch.no_grad()
-def bench_torch(a, b):
-    for _ in range(N_WARMUP):
-        _ = a @ b
-        torch.mps.synchronize()
-
-    times = []
-    for _ in range(N_BENCH):
-        start = time.perf_counter_ns()
-        _ = a @ b
-        torch.mps.synchronize()
-        end = time.perf_counter_ns()
-        times.append((end - start) * 1e-9)
-
-    return np.mean(times), np.std(times)
-
-
-def check_correctness(out_mx, out_pt, rtol, M, N, K):
-    if not np.allclose(out_pt, out_mx, rtol=rtol, atol=0):
-        abs_diff = np.abs(out_pt - out_mx)
-        rel_diff = abs_diff / np.maximum(np.abs(out_pt), 1e-10)
-
-        print(
-            f"  WARNING: Correctness failed at {M}x{N}x{K}: "
-            f"max_abs={np.max(abs_diff):.6e}, max_rel={np.max(rel_diff):.6e}"
-        )
-
-
-def bench_gemm(M, N, K, dtype, rtol):
-    scale = 0.5 / math.sqrt(K)
-    a_np = np.random.uniform(0, scale, (M, K)).astype(np.float32)
-    b_np = np.random.uniform(0, scale, (K, N)).astype(np.float32)
-
-    a_mx = mx.array(a_np).astype(getattr(mx, dtype))
-    b_mx = mx.array(b_np).astype(getattr(mx, dtype))
-
-    a_pt = torch.from_numpy(a_np).to(dtype=getattr(torch, dtype), device="mps")
-    b_pt = torch.from_numpy(b_np).to(dtype=getattr(torch, dtype), device="mps")
-    torch.mps.synchronize()
-
-    torch_mean, torch_std = bench_torch(a_pt, b_pt)
-    mlx_mean, mlx_std = bench_mlx(a_mx, b_mx)
-
-    out_mx = (a_mx @ b_mx).astype(mx.float32)
-    out_pt = (a_pt @ b_pt).to(torch.float32).to("cpu").numpy(force=True)
-    check_correctness(out_mx, out_pt, rtol, M, N, K)
-
-    return mlx_mean, mlx_std, torch_mean, torch_std
-
-
-if __name__ == "__main__":
-    dtypes = ("bfloat16", "float16", "float32")
-
-    rtols = {
-        "float32": 1e-3,
-        "float16": 5e-3,
-        "bfloat16": 1e-2,
-    }
-
-    shapes = (
-        (2048, 2048, 10240),
-        (2048, 3072, 10240),
-        (3072, 3072, 10240),
-        (3072, 3072, 12288),
-        (3072, 4096, 12288),
-        (4096, 4096, 12288),
-        (4096, 4096, 18432),
-        (4096, 4096, 21504),
-        (4096, 6144, 21504),
-        (6144, 6144, 21504),
-    )
-
-    for dtype in dtypes:
-        print(f"\nPerformance ({dtype}):")
-        print(
-            f"{'M':>5s} {'N':>5s} {'K':>6s}  "
-            f"{'MLX (ms)':>15s}  {'Torch (ms)':>15s}  {'Speedup':>10s}"
-        )
-        print("-" * 80)
-
-        for M, N, K in shapes:
-            mlx_mean, mlx_std, torch_mean, torch_std = bench_gemm(
-                M, N, K, dtype, rtols[dtype]
-            )
-            speedup = torch_mean / mlx_mean
-
-            print(
-                f"{M:5d} {N:5d} {K:6d}  "
-                f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f}  "
-                f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f}  "
-                f"{speedup:8.2f}x"
-            )

benchmarks/python/masked_scatter.py

@@ -1,6 +1,5 @@
 import math
 import os
-import platform
 import subprocess
 import time
 from copy import copy
@@ -18,6 +17,9 @@ 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()
@@ -25,36 +27,11 @@ TORCH_DEVICE = torch.device(
 )
 
 
-def get_device_name():
-    if TORCH_DEVICE.type == "cuda":
-        try:
-            out = subprocess.check_output(
-                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
-                stderr=subprocess.DEVNULL,
-            )
-            return out.decode("utf-8").splitlines()[0].strip()
-        except Exception:
-            return "CUDA_GPU"
-    if TORCH_DEVICE.type == "mps":
-        try:
-            out = subprocess.check_output(
-                ["sysctl", "-n", "machdep.cpu.brand_string"],
-                stderr=subprocess.DEVNULL,
-            )
-            return out.decode("utf-8").strip()
-        except Exception:
-            return "Apple_Silicon"
-    return platform.processor() or platform.machine() or "CPU"
-
-
-DEVICE_NAME = get_device_name()
-
-
 N_WARMUP = 5
 N_ITER_BENCH = 50
 N_ITER_FUNC = 20
 
-VECTOR_LENGTHS = [4096 * (2**i) for i in range(12)]
+VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)]
 MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
 D_TYPES = ("float32", "float16")
 
@@ -225,10 +202,9 @@ def main():
         )
         output_path = os.path.join(
             RESULTS_DIR,
-            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.png",
+            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf",
         )
         fig.savefig(output_path)
-        print(f"Saved benchmark image: {output_path}")
         plt.close(fig)
 
 

benchmarks/python/sdpa_bench.py

@@ -176,8 +176,6 @@ if __name__ == "__main__":
           (  1,  1024,  1024,       64,   32,     8),
           (  1,  2048,  2048,       64,   32,     8),
           (  1,  4096,  4096,       64,   32,     8),
-          (  1,  4096,  5000,       64,   32,     8),
-          (  1,  2048,  32121,      64,   32,     8),
     )
 
     shapes_80 = (
@@ -185,8 +183,6 @@ if __name__ == "__main__":
           (  1,  1024,  1024,       80,   32,     8),
           (  1,  2048,  2048,       80,   32,     8),
           (  1,  4096,  4096,       80,   32,     8),
-          (  1,  4096,  5000,       80,   32,     8),
-          (  1,  2048,  32121,      80,   32,     8),
     )
 
     shapes_128 = (
@@ -194,8 +190,6 @@ if __name__ == "__main__":
           (  1,  1024,  1024,      128,   32,     8),
           (  1,  2048,  2048,      128,   32,     8),
           (  1,  4096,  4096,      128,   32,     8),
-          (  1,  4096,  5000,      128,   32,     8),
-          (  1,  2048,  32121,     128,   32,     8),
     )
     # fmt: on
 

benchmarks/python/segmented_mm_bench.py

@@ -1,209 +0,0 @@
-# Copyright © 2026 Apple Inc.
-
-import argparse
-import time
-
-import mlx.core as mx
-import numpy as np
-
-MLX_DTYPES = {
-    "float16": mx.float16,
-    "bfloat16": mx.bfloat16,
-    "float32": mx.float32,
-}
-
-
-def parse_cases(cases):
-    parsed = []
-    for spec in cases.split(","):
-        m, n, k, s = [int(x) for x in spec.split("x")]
-        parsed.append((m, n, k, s))
-    return parsed
-
-
-def make_segments(k, num_segments, pattern, seed):
-    if pattern == "equal":
-        cuts = np.linspace(0, k, num_segments + 1, dtype=np.int64)
-    else:
-        rng = np.random.default_rng(seed)
-        cuts = rng.integers(0, k + 1, size=(num_segments - 1,), dtype=np.int64)
-        cuts = np.sort(cuts)
-        cuts = np.concatenate(([0], cuts, [k]))
-    return np.stack([cuts[:-1], cuts[1:]], axis=1).astype(np.uint32)
-
-
-def numpy_segmented_mm_ref(a, b, segments):
-    """Ground-truth reference in float64."""
-    out = []
-    for start, end in segments:
-        out.append(a[:, start:end] @ b[start:end, :])
-    return np.stack(out, axis=0)
-
-
-def mlx_segmented_mm_loop(a, b, segments):
-    """MLX loop-of-matmuls baseline."""
-    segments_list = segments.tolist()
-    out = []
-    for start, end in segments_list:
-        out.append(a[:, start:end] @ b[start:end, :])
-    return mx.stack(out, axis=0)
-
-
-def bench_mlx(a, b, segments, warmup, iters):
-    for _ in range(warmup):
-        y = mx.segmented_mm(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = mx.segmented_mm(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-    end = time.perf_counter()
-    return (end - start) * 1e3 / iters
-
-
-def bench_mlx_loop(a, b, segments, warmup, iters):
-    for _ in range(warmup):
-        y = mlx_segmented_mm_loop(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = mlx_segmented_mm_loop(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-    end = time.perf_counter()
-    return (end - start) * 1e3 / iters
-
-
-def print_table(headers, rows):
-    widths = [len(h) for h in headers]
-    for row in rows:
-        for i, cell in enumerate(row):
-            widths[i] = max(widths[i], len(cell))
-
-    def fmt_row(row):
-        return (
-            "| "
-            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
-            + " |"
-        )
-
-    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
-    print(fmt_row(headers))
-    print(sep)
-    for row in rows:
-        print(fmt_row(row))
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--cases",
-        default=(
-            "128x128x1024x16,"
-            "128x128x1024x32,"
-            "256x256x2048x16,"
-            "512x512x4096x32,"
-            "1024x1024x4096x32,"
-            "1024x1024x8192x64"
-        ),
-        help="Comma-separated MxNxKxS list.",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="float32",
-        choices=["float16", "bfloat16", "float32"],
-    )
-    parser.add_argument("--warmup", type=int, default=10)
-    parser.add_argument("--iters", type=int, default=50)
-    parser.add_argument(
-        "--segments",
-        choices=["equal", "random"],
-        default="random",
-        help="Segment generation pattern.",
-    )
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--no-check", action="store_true")
-    args = parser.parse_args()
-
-    mlx_dtype = MLX_DTYPES[args.dtype]
-
-    print(
-        f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}"
-    )
-
-    headers = [
-        "Case",
-        "MLX ms",
-        "Loop ms",
-        "Speedup",
-        "MLX err",
-        "Loop err",
-    ]
-    rows = []
-
-    cases = parse_cases(args.cases)
-    for idx, (m, n, k, s) in enumerate(cases):
-        rng = np.random.default_rng(args.seed + idx)
-        a_np = rng.standard_normal((m, k)).astype(np.float32)
-        b_np = rng.standard_normal((k, n)).astype(np.float32)
-        seg_np = make_segments(k, s, args.segments, args.seed + idx)
-
-        a_mx = mx.array(a_np, dtype=mlx_dtype)
-        b_mx = mx.array(b_np, dtype=mlx_dtype)
-        seg_mx = mx.array(seg_np, dtype=mx.uint32)
-        mx.eval(a_mx, b_mx, seg_mx)
-
-        mlx_err_str = ""
-        loop_err_str = ""
-        if not args.no_check:
-            y_mlx = mx.segmented_mm(a_mx, b_mx, seg_mx)
-            y_loop = mlx_segmented_mm_loop(a_mx, b_mx, seg_mx)
-            mx.eval(y_mlx, y_loop)
-
-            if args.dtype == "float32":
-                ref = numpy_segmented_mm_ref(
-                    a_np.astype(np.float64),
-                    b_np.astype(np.float64),
-                    seg_np.tolist(),
-                )
-                mlx_err = np.max(np.abs(np.array(y_mlx, dtype=np.float64) - ref))
-                loop_err = np.max(np.abs(np.array(y_loop, dtype=np.float64) - ref))
-            else:
-                a_mx_f32 = mx.array(a_np, dtype=mx.float32)
-                b_mx_f32 = mx.array(b_np, dtype=mx.float32)
-                ref = mx.segmented_mm(a_mx_f32, b_mx_f32, seg_mx)
-                mx.eval(ref)
-                mlx_err = float(mx.max(mx.abs(ref - y_mlx.astype(mx.float32))).item())
-                loop_err = float(mx.max(mx.abs(ref - y_loop.astype(mx.float32))).item())
-            mlx_err_str = f"{mlx_err:.2e}"
-            loop_err_str = f"{loop_err:.2e}"
-
-        t_mlx = bench_mlx(a_mx, b_mx, seg_mx, args.warmup, args.iters)
-        t_loop = bench_mlx_loop(a_mx, b_mx, seg_mx, args.warmup, args.iters)
-        ratio = t_loop / t_mlx if t_mlx > 0 else float("inf")
-        rows.append(
-            [
-                f"{m}x{n}x{k}x{s}",
-                f"{t_mlx:.3f}",
-                f"{t_loop:.3f}",
-                f"{ratio:.2f}x",
-                mlx_err_str,
-                loop_err_str,
-            ]
-        )
-
-    print_table(headers, rows)
-    if not args.no_check:
-        if args.dtype == "float32":
-            print("err: max|result - numpy_fp64_ref|")
-        else:
-            print("err: max|result - own_fp32_result|")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/python/slice_update_bench.py

@@ -1,109 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import argparse
-
-import mlx.core as mx
-import torch
-from time_utils import measure_runtime
-
-
-def benchmark_slice_update_mlx(dst_shape, slice_shape, slice_range, dtype, iters=10):
-    def slice_update(arguments):
-        for i in range(iters):
-            arguments["dst"] = (
-                arguments["dst"].at[slice_range].add(arguments["updates"])
-            )
-        mx.eval(arguments)
-
-    dtype = getattr(mx, dtype)
-    arguments = {
-        "dst": mx.random.normal(dst_shape).astype(dtype),
-        "updates": mx.random.normal(slice_shape).astype(dtype),
-    }
-
-    runtime = measure_runtime(slice_update, arguments=arguments)
-    bytes_processed = (
-        arguments["dst"][slice_range].nbytes * 2 + arguments["updates"].nbytes
-    ) * iters
-    bandwidth_gb_s = bytes_processed / runtime / 1e6
-    return runtime, bandwidth_gb_s
-
-
-def benchmark_slice_update_torch(
-    dst_shape, slice_shape, slice_range, device, dtype, iters=10
-):
-    def slice_update(dst, updates, slice_range):
-        for i in range(iters):
-            dst[slice_range] = dst[slice_range] + updates
-        if device == torch.device("mps"):
-            torch.mps.synchronize()
-
-    dtype = getattr(torch, dtype)
-    updates = torch.randn(slice_shape, dtype=dtype).to(device)
-    dst = torch.randn(dst_shape, dtype=dtype).to(device)
-
-    runtime = measure_runtime(
-        slice_update, dst=dst, updates=updates, slice_range=slice_range
-    )
-    bytes_processed = (dst[slice_range].nbytes * 2 + updates.nbytes) * iters
-    bandwidth_gb_s = bytes_processed / runtime / 1e6
-    return runtime, bandwidth_gb_s
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser("Slice update benchmarks.")
-    parser.add_argument("--cpu", action="store_true", help="Use the CPU.")
-    args = parser.parse_args()
-
-    if args.cpu:
-        mx.set_default_device(mx.cpu)
-        device = torch.device("cpu")
-    elif torch.mps.is_available():
-        device = torch.device("mps")
-    elif torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        raise ValueError()
-
-    dtypes = ["float32", "bfloat16"]
-
-    test_cases = [
-        ((10_000_000,), slice(0, 1_000_000), (1_000_000,)),
-        ((100_000,), slice(10_000, 20_000), (10_000,)),
-        ((1000, 64), slice(100, 200), (100, 64)),
-        ((100, 100, 64), slice(20, 40), (20, 100, 64)),
-        (
-            (2048, 2048, 128),
-            (slice(500, 1500), slice(200, 1200), slice(32, 96)),
-            (1000, 1000, 64),
-        ),
-        (
-            (2048, 2048, 128),
-            (slice(1800, 1850), slice(100, 200), slice(64, 128)),
-            (50, 100, 64),
-        ),
-        (
-            (2048, 2048, 128),
-            (slice(1000, 1010), slice(1000, 1010), slice(64, 128)),
-            (10, 10, 64),
-        ),
-    ]
-
-    print(
-        f"{'Dtype':<12} {'Dst Shape':<25} {'Update Shape':<20} "
-        f"{'MLX (ms)':<12} {'MLX GB/s':<12} {'Torch (ms)':<12} {'Torch GB/s':<12}"
-    )
-    print("-" * 110)
-
-    for dtype in dtypes:
-        for dst_shape, slice_range, update_shape in test_cases:
-            mlx_time, mlx_bw = benchmark_slice_update_mlx(
-                dst_shape, update_shape, slice_range, dtype
-            )
-            torch_time, torch_bw = benchmark_slice_update_torch(
-                dst_shape, update_shape, slice_range, device, dtype
-            )
-            print(
-                f"{dtype:<12} {str(dst_shape):<25} {str(update_shape):<20} "
-                f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} {torch_bw:<12.2f}"
-            )

cmake/FindCUDNN.cmake

@@ -1,177 +0,0 @@
-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# Modified from
-# https://github.com/NVIDIA/cudnn-frontend/blob/main/cmake/cuDNN.cmake
-
-# Return the last file matching the pattern.
-function(find_file_glob VAR PATTERN)
-  file(GLOB _RESULT "${PATTERN}")
-  if(_RESULT)
-    list(LENGTH ${_RESULT} _RESULT_LENGTH)
-    if(_RESULT_LENGTH GREATER 0)
-      list(GET ${_RESULT} -1 _RESULT)
-    endif()
-    set(${VAR}
-        "${_RESULT}"
-        PARENT_SCOPE)
-  endif()
-endfunction()
-
-# Find the dir including the "cudnn.h" file.
-find_path(
-  CUDNN_INCLUDE_DIR cudnn.h
-  HINTS ${CUDNN_INCLUDE_PATH} ${CUDAToolkit_INCLUDE_DIRS}
-  PATH_SUFFIXES include OPTIONAL)
-
-# Glob searching "cudnn.h" for Windows.
-if(WIN32 AND NOT CUDNN_INCLUDE_DIR)
-  find_file_glob(
-    CUDNN_H_PATH
-    "C:/Program Files/NVIDIA/CUDNN/*/include/${CUDAToolkit_VERSION_MAJOR}.*/cudnn.h"
-  )
-  if(CUDNN_H_PATH)
-    get_filename_component(CUDNN_INCLUDE_DIR "${CUDNN_H_PATH}" DIRECTORY)
-  endif()
-endif()
-
-if(NOT CUDNN_INCLUDE_DIR)
-  message(
-    FATAL_ERROR
-      "Unable to find cudnn.h, please make sure cuDNN is installed and pass CUDNN_INCLUDE_PATH to cmake."
-  )
-endif()
-
-# Get cudnn version.
-file(READ "${CUDNN_INCLUDE_DIR}/cudnn_version.h" cudnn_version_header)
-string(REGEX MATCH "#define CUDNN_MAJOR [1-9]+" macrodef
-             "${cudnn_version_header}")
-string(REGEX MATCH "[1-9]+" CUDNN_MAJOR_VERSION "${macrodef}")
-
-# Function for searching library files.
-function(find_cudnn_library NAME)
-  if(NOT "${ARGV1}" STREQUAL "OPTIONAL")
-    set(_CUDNN_REQUIRED TRUE)
-  else()
-    set(_CUDNN_REQUIRED FALSE)
-  endif()
-
-  find_library(
-    ${NAME}_LIBRARY
-    NAMES ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}" NAMES_PER_DIR
-    HINTS ${CUDNN_LIBRARY_PATH} ${CUDAToolkit_LIBRARY_DIR}
-    PATH_SUFFIXES lib64 lib/x64 lib OPTIONAL)
-
-  if(WIN32 AND NOT ${NAME}_LIBRARY)
-    find_file_glob(
-      ${NAME}_LIBRARY
-      "C:/Program Files/NVIDIA/CUDNN/*/lib/${CUDAToolkit_VERSION_MAJOR}.*/x64/${NAME}.lib"
-    )
-  endif()
-
-  if(NOT ${NAME}_LIBRARY AND ${_CUDNN_REQUIRED})
-    message(
-      FATAL_ERROR
-        "Unable to find ${NAME}, please make sure cuDNN is installed and pass CUDNN_LIBRARY_PATH to cmake."
-    )
-  endif()
-
-  if(${NAME}_LIBRARY)
-    add_library(CUDNN::${NAME} UNKNOWN IMPORTED)
-    set_target_properties(
-      CUDNN::${NAME}
-      PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR}
-                 IMPORTED_LOCATION ${${NAME}_LIBRARY})
-    set(${NAME}_LIBRARY
-        "${${NAME}_LIBRARY}"
-        PARENT_SCOPE)
-  else()
-    message(STATUS "${NAME} not found.")
-  endif()
-endfunction()
-
-# Search for the main cudnn library.
-find_cudnn_library(cudnn)
-
-include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(CUDNN REQUIRED_VARS CUDNN_INCLUDE_DIR
-                                                      cudnn_LIBRARY)
-
-if(CUDNN_INCLUDE_DIR AND cudnn_LIBRARY)
-  set(CUDNN_FOUND
-      ON
-      CACHE INTERNAL "cuDNN Library Found")
-else()
-  set(CUDNN_FOUND
-      OFF
-      CACHE INTERNAL "cuDNN Library Not Found")
-endif()
-
-# Find out all the DLL files for Windows.
-if(WIN32 AND cudnn_LIBRARY)
-  get_filename_component(CUDNN_BIN_DIR "${cudnn_LIBRARY}" DIRECTORY)
-  string(REPLACE "/lib/" "/bin/" CUDNN_BIN_DIR "${CUDNN_BIN_DIR}")
-  file(
-    GLOB CUDNN_DLL_NAMES
-    RELATIVE "${CUDNN_BIN_DIR}"
-    "${CUDNN_BIN_DIR}/*.dll")
-endif()
-
-# Create an interface library that users can link with.
-add_library(CUDNN::cudnn_all INTERFACE IMPORTED)
-target_link_libraries(CUDNN::cudnn_all INTERFACE CUDNN::cudnn)
-target_include_directories(
-  CUDNN::cudnn_all INTERFACE $<INSTALL_INTERFACE:include>
-                             $<BUILD_INTERFACE:${CUDNN_INCLUDE_DIR}>)
-
-# Add other components of cudnn.
-if(CUDNN_MAJOR_VERSION EQUAL 8)
-  find_cudnn_library(cudnn_adv_infer)
-  find_cudnn_library(cudnn_adv_train)
-  find_cudnn_library(cudnn_cnn_infer)
-  find_cudnn_library(cudnn_cnn_train)
-  find_cudnn_library(cudnn_ops_infer)
-  find_cudnn_library(cudnn_ops_train)
-
-  target_link_libraries(
-    CUDNN::cudnn_all
-    INTERFACE CUDNN::cudnn_adv_train CUDNN::cudnn_ops_train
-              CUDNN::cudnn_cnn_train CUDNN::cudnn_adv_infer
-              CUDNN::cudnn_cnn_infer CUDNN::cudnn_ops_infer)
-
-elseif(CUDNN_MAJOR_VERSION EQUAL 9)
-  find_cudnn_library(cudnn_graph)
-  find_cudnn_library(cudnn_engines_runtime_compiled)
-  find_cudnn_library(cudnn_ops OPTIONAL)
-  find_cudnn_library(cudnn_cnn OPTIONAL)
-  find_cudnn_library(cudnn_adv OPTIONAL)
-  find_cudnn_library(cudnn_engines_precompiled OPTIONAL)
-  find_cudnn_library(cudnn_heuristic OPTIONAL)
-
-  target_link_libraries(
-    CUDNN::cudnn_all
-    INTERFACE CUDNN::cudnn_graph
-              CUDNN::cudnn_engines_runtime_compiled
-              CUDNN::cudnn_ops
-              CUDNN::cudnn_cnn
-              CUDNN::cudnn_adv
-              CUDNN::cudnn_engines_precompiled
-              CUDNN::cudnn_heuristic)
-endif()

docs/Doxyfile

@@ -26,7 +26,6 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
-PREDEFINED             = MLX_API=
 
 ################################################################################
 # Compound extraction control.                                                 #

docs/README.md

@@ -38,17 +38,3 @@ the docs. Then force add the `build/html` directory:
 `git add -f build/html`
 
 Commit and push the changes to the `gh-pages` branch.
-
-## Doc Development Setup
-
-To enable live refresh of docs while writing:
-
-Install sphinx autobuild
-```
-pip install sphinx-autobuild
-```
-
-Run auto build on docs/src folder
-```
-sphinx-autobuild ./src ./build/html
-```

docs/logo/mlx_logo.svg

@@ -1,36 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
-<defs>
-<g>
-<g id="glyph-0-0">
-<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
-</g>
-<g id="glyph-0-1">
-<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-2">
-<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-3">
-<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
-</g>
-</g>
-<clipPath id="clip-0">
-<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
-</clipPath>
-<clipPath id="clip-1">
-<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
-</clipPath>
-</defs>
-<g clip-path="url(#clip-0)">
-<g fill="rgb(0%, 0%, 0%)" fill-opacity="1">
-<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
-<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
-</g>
-</g>
-<g clip-path="url(#clip-1)">
-<g fill="rgb(82.998657%, 82.998657%, 82.998657%)" fill-opacity="1">
-<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
-</g>
-</g>
-</svg>

docs/logo/mlx_logo_dark.svg

@@ -1,36 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
-<defs>
-<g>
-<g id="glyph-0-0">
-<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
-</g>
-<g id="glyph-0-1">
-<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-2">
-<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-3">
-<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
-</g>
-</g>
-<clipPath id="clip-0">
-<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
-</clipPath>
-<clipPath id="clip-1">
-<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
-</clipPath>
-</defs>
-<g clip-path="url(#clip-0)">
-<g fill="rgb(100%, 100%, 100%)" fill-opacity="1">
-<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
-<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
-</g>
-</g>
-<g clip-path="url(#clip-1)">
-<g fill="rgb(56.999207%, 56.999207%, 56.999207%)" fill-opacity="1">
-<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
-</g>
-</g>
-</svg>

docs/src/dev/extensions.rst

@@ -404,7 +404,7 @@ below.
         auto kernel = d.get_kernel(kname, lib);
 
         // Prepare to encode kernel
-        auto& compute_encoder = mx::metal::get_command_encoder(s);
+        auto& compute_encoder = d.get_command_encoder(s.index);
         compute_encoder.set_compute_pipeline_state(kernel);
 
         // Kernel parameters are registered with buffer indices corresponding to
@@ -448,7 +448,7 @@ We can now call the :meth:`axpby` operation on both the CPU and the GPU!
 
 A few things to note about MLX and Metal before moving on. MLX keeps track of
 the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is
-associated. We rely on :meth:`metal::get_command_encoder` to give us the active
+associated. We rely on :meth:`d.get_command_encoder` to give us the active
 metal compute command encoder instead of building a new one and calling
 :meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute
 pipelines) to the active command buffer until some specified limit is hit or

docs/src/dev/mlx_in_cpp.rst

@@ -45,7 +45,7 @@ The next step is to setup a CMake file in ``CMakeLists.txt``:
 
   project(example LANGUAGES CXX)
 
-  set(CMAKE_CXX_STANDARD 20)
+  set(CMAKE_CXX_STANDARD 17)
   set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 

docs/src/examples/data_parallelism.rst

@@ -1,91 +0,0 @@
-.. _data_parallelism:
-
-Data Parallelism
-================
-
-MLX enables efficient data parallel distributed training through its
-distributed communication primitives.
-
-.. _training_example:
-
-Training Example
-----------------
-
-In this section we will adapt an MLX training loop to support data parallel
-distributed training. Namely, we will average the gradients across a set of
-hosts before applying them to the model.
-
-Our training loop looks like the following code snippet if we omit the model,
-dataset, and optimizer initialization.
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())
-
-All we have to do to average the gradients across machines is perform an
-:func:`all_sum` and divide by the size of the :class:`Group`. Namely we
-have to :func:`mlx.utils.tree_map` the gradients with following function.
-
-.. code:: python
-
-    def all_avg(x):
-        return mx.distributed.all_sum(x) / mx.distributed.init().size()
-
-Putting everything together our training loop step looks as follows with
-everything else remaining the same.
-
-.. code:: python
-
-    from mlx.utils import tree_map
-
-    def all_reduce_grads(grads):
-        N = mx.distributed.init().size()
-        if N == 1:
-            return grads
-        return tree_map(
-            lambda x: mx.distributed.all_sum(x) / N,
-            grads
-        )
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = all_reduce_grads(grads)  # <--- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-Using ``nn.average_gradients``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Although the code example above works correctly; it performs one communication
-per gradient. It is significantly more efficient to aggregate several gradients
-together and perform fewer communication steps.
-
-This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
-almost identical to the example above:
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = mx.nn.average_gradients(grads)  # <---- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())

docs/src/examples/tensor_parallelism.rst

@@ -1,239 +0,0 @@
-.. _tensor_parallelism:
-
-Tensor Parallelism
-==================
-
-In this example, we will explore how tensor parallelism (TP) works in MLX.  We
-will start with an overview of the distributed layers in ``mlx.nn`` and then
-show how to do tensor parallelism Llama-style transformer models.
-
-Sharded Layers
---------------
-
-:class:`AllToShardedLinear <mlx.nn.AllToShardedLinear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This layer replicates a common input and shards the weight matrix along the
-output dimension across all devices in the :class:`mlx.core.distributed.Group`.
-The layer produces a sharded output.
-
-For example, consider an :class:`mlx.nn.AllToShardedLinear` layer with
-``input_dims=2`` and ``output_dims=2``, a batched input of shape ``(4, 2)``,
-and a device group with 2 devices. The layer shards the weight matrix along the
-output dimension across the two devices, where each device receives the full
-input and computes a partial output.
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/all-to-sharded-linear.png" alt="column-wise tensor parallelism" style="width: 100%">
-    </div>
-
-This layer does not automatically gather all outputs from each device. This is
-an intended and :ref:`useful design choice <useful_design_choices>`.
-
-:class:`QuantizedAllToShardedLinear <mlx.nn.QuantizedAllToShardedLinear>` is
-the quantized equivalent of :class:`mlx.nn.AllToShardedLinear`.  Similar to
-:class:`mlx.nn.QuantizedLinear`, its parameters are frozen and will not be
-included in any gradient computation.
-
-
-:class:`ShardedToAllLinear <mlx.nn.ShardedToAllLinear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This layer expects inputs that are sharded along the feature dimension and
-shards the weight matrix along the input dimension across all devices in the
-:class:`mlx.core.distributed.Group`. The layer automatically aggregates the
-results using :class:`mlx.core.distributed.all_sum`, so all devices in the
-group will have the same result.
-
-For example, consider an :class:`mlx.nn.ShardedToAllLinear` layer with
-``input_dims=2`` and ``output_dims=2``, a batched input of shape ``(4, 2)``,
-and a device group with 2 devices. The layer shards the weight matrix along the
-input dimension across the two devices. Each device computes a ``(4,2)``
-output, which is then aggregated with all other device outputs to get layer
-output.
-
-   .. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/sharded-to-all-linear.png" alt="row-wise tensor parallelism" style="width: 100%">
-    </div>
-
-This layer does not automatically shard the inputs along the feature dimension
-for you. It is necessary to create a "partial" input structure to feed into the
-layer. This is an intended and :ref:`useful design choice
-<useful_design_choices>`.
-
-:class:`QuantizedShardedToAllLinear <mlx.nn.QuantizedShardedToAllLinear>` is
-the quantized equivalent of :class:`mlx.nn.ShardedToAllLinear`.  Similar to
-:class:`mlx.nn.QuantizedLinear`, its parameters are frozen and will not be
-included in any gradient computation.
-
-
-Shard Utility Functions
------------------------
-
-:func:`shard_linear <mlx.nn.layers.distributed.shard_linear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Converts a regular linear layer into a tensor parallel layer that distributes
-computation across multiple devices. Takes an existing :class:`mlx.nn.Linear`
-or :class:`mlx.nn.QuantizedLinear` layer and returns a new distributed layer
-(either :class:`mlx.nn.AllToShardedLinear` or
-:class:`mlx.nn.ShardedToAllLinear`, depending on the sharding type). The
-original layer is not modified.
-
-:func:`shard_inplace <mlx.nn.layers.distributed.shard_inplace>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Splits the parameters of an existing layer across multiple devices by modifying
-the layer in-place. Unlike :func:`shard_linear
-<mlx.nn.layers.distributed.shard_linear>`, this function does not create a new
-layer or add distributed communication. The layer itself must handle
-distributed communication if needed.
-
-
-.. _useful_design_choices:
-
-Useful Design Choices
----------------------
-
-The design choices above regarding when operations are done automatically are intentional and make model training and inference easier.
-
-All-to-sharded and sharded-to-all layers naturally go together because the
-output of the former layer is exactly the input needed needed for the latter.
-This removes the need for an intermediate gather step between the layers,
-reducing communication overhead.
-
-This is why :class:`mlx.nn.AllToShardedLinear` does not aggregate results
-automatically and why :class:`mlx.nn.ShardedToAllLinear` does not shard inputs
-automatically. It is so that they can be placed in successive order and work
-together easily.
-
-We can demonstrate this through a simple model using our two types of
-distributed layers.
-
-.. code-block:: python
-
-  x = ... # some (4, 2) model input: batch size 4, feature size 2
-
-  l1 = nn.AllToShardedLinear(2, 2, bias=False)   # initialize the layer
-  l1_out = l1(x) # (4, 1) output
-
-  l2 = nn.ShardedToAllLinear(2, 2, bias=False)
-  l2_out = l2(l1_out) # (4, 2) output
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/column-row-tp.png" alt="two layer tensor parallelism" style="width: 100%">
-      <p style="font-size: 0.85em; margin-top: 0.5em;"><small>A visualization of the simple MLX model using all-to-sharded then sharded-to-all tensor parallelism across 2 devices.</small></p>
-    </div>
-
-
-LLM Inference with Tensor Parallelism
--------------------------------------
-
-We can apply these TP techniques to LLMs in order to enable inference for much
-larger models by sharding parameters from huge layers across multiple devices.
-
-To demonstrate this, let's apply TP to the Transformer block of our :doc:`Llama
-Inference <llama-inference>` example. In this example, we will use the same
-inference script as the Llama Inference example, which can be found in
-`mlx-examples`_.
-
-Our first edit is to initialize the distributed communication group and get the
-current process rank:
-
-.. code-block:: python
-
-  world = mx.distributed.init()
-  rank = world.rank()
-
-Next, let's look at the current architecture of the transformer block and see how we can apply tensor parallelism:
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/llama-transformer.png" alt="llama transformer example" style="width: 100%">
-    </div>
-
-
-This architecture has two natural places where 
-tensor parallelism can be applied: the attention block and the FFN
-block. Both follow the same pattern: multiple parallel linear layers operating
-on the same input, followed by a single output linear layer. In the attention
-block, the Q, K, and V projections are sharded along the output dimension (all-to-sharded), and the output
-projection is sharded along the input dimension (sharded-to-all). Similarly in the FFN block, the gate and up projections
-become all-to-sharded layers, and the down projection becomes an sharded-to-all layer.
-
-The intermediate operations between the linear layers (RoPE, softmax, scaled
-dot-product attention in the attention block, and element-wise multiplication
-in the FFN block) do not impede the use of our TP paradigm. These operations
-are either:
-
-- **Element-wise operations** (RoPE, element-wise multiplication): These
-  operate independently on each element or position, preserving the sharding
-  pattern without requiring cross-device communication.
-
-- **Operations on non-sharded dimensions** (softmax, scaled dot-product
-  attention): These operate along dimensions that are not sharded (such as the
-  sequence length or head dimensions), so they can be computed independently on
-  each device. The attention computation ``Q @ K^T`` and ``scores @ V`` work
-  correctly with sharded Q, K, V tensors because the matrix multiplications are
-  performed along the sharded feature dimension, and the results remain
-  properly sharded for the subsequent sharded-to-all layer.
-
-To implement sharding in our Llama inference, we use :func:`shard_linear
-<mlx.nn.layers.distributed.shard_linear>` to get sharded linear layers with
-distributed communication. This is easier than using :func:`shard_inplace
-<mlx.nn.layers.distributed.shard_inplace>` and implementing the steps manually
-in the :code:`__call__` function.
-
-The following code shows how to shard the Attention block. The Q, K, and V
-projection layers are converted to all-to-sharded layers, while the output
-projection is converted to a sharded-to-all layer. The number of heads are also
-adjusted to account for the sharding:
-
-.. code-block:: python
-
-  # ... in Attention class
-  def shard(self, group: mx.distributed.Group):
-    self.n_heads = self.n_heads // group.size()
-    self.n_kv_heads = self.n_kv_heads // group.size()
-
-    self.wq = nn.layers.distributed.shard_linear(self.wq, "all-to-sharded", group=group)
-    self.wk = nn.layers.distributed.shard_linear(self.wk, "all-to-sharded", group=group)
-    self.wv = nn.layers.distributed.shard_linear(self.wv, "all-to-sharded", group=group)
-    self.wo = nn.layers.distributed.shard_linear(self.wo, "sharded-to-all", group=group)
-
-Similarly, the FeedForward block is sharded by converting the gate (w1) and up
-(w3) projections to all-to-sharded layers, and the down projection (w2) to
-a sharded-to-all layer:
-
-.. code-block:: python
-
-  # ... in FeedForward class
-  def shard(self, group: mx.distributed.Group):
-    self.w1 = nn.layers.distributed.shard_linear(self.w1, "all-to-sharded", group=group)
-    self.w2 = nn.layers.distributed.shard_linear(self.w2, "sharded-to-all", group=group)
-    self.w3 = nn.layers.distributed.shard_linear(self.w3, "all-to-sharded", group=group)
-
-Finally, in our :code:`load_model` function, we need to apply our sharding
-functions to all transformer layers when using multiple devices:
-
-.. code-block:: python
-
-  # ... in load_model function
-  if world.size() > 1:
-    # convert Linear layers in Transformer/FFN to appropriate Sharded Layers
-    for layer in model.layers:
-        layer.attention.shard(group=world)
-        layer.feed_forward.shard(group=world)
-
-This allows us to use the llama inference file as normal when running
-:code:`python llama.py`, but now we can also run it across two (or more)
-devices via :code:`mlx.launch -n 2 llama.py`.
-
-.. _mlx-examples: https://github.com/ml-explore/mlx-examples/tree/main/llms/llama

docs/src/index.rst

@@ -32,7 +32,7 @@ are the CPU and GPU.
    install
 
 .. toctree::
-   :caption: Usage
+   :caption: Usage 
    :maxdepth: 1
 
    usage/quick_start
@@ -54,8 +54,6 @@ are the CPU and GPU.
    examples/linear_regression
    examples/mlp
    examples/llama-inference
-   examples/data_parallelism
-   examples/tensor_parallelism
 
 .. toctree::
    :caption: Python API Reference
@@ -78,7 +76,6 @@ are the CPU and GPU.
    python/optimizers
    python/distributed
    python/tree_utils
-   python/printoptions
 
 .. toctree::
    :caption: C++ API Reference

docs/src/install.rst

@@ -15,7 +15,7 @@ silicon computer is
 
 To install from PyPI your system must meet the following requirements:
 
-- Using `Apple silicon <https://support.apple.com/en-us/116943>`_
+- Using an M series chip (Apple silicon)
 - Using a native Python >= 3.10
 - macOS >= 14.0
 
@@ -83,8 +83,7 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^
 
-- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux)
-- A C++ compiler with C++20 support (e.g. Clang >= 15.0)
+- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
 - `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0
 

docs/src/python/devices_and_streams.rst

@@ -14,10 +14,6 @@ Devices and Streams
    set_default_device
    default_stream
    new_stream
-   new_thread_local_stream
    set_default_stream
    stream
    synchronize
-   clear_streams
-   device_count
-   device_info

docs/src/python/fft.rst

@@ -20,7 +20,5 @@ FFT
   irfft2
   rfftn
   irfftn
-  fftfreq
-  rfftfreq
   fftshift
   ifftshift

docs/src/python/linalg.rst

@@ -14,7 +14,6 @@ Linear Algebra
     cholesky
     cholesky_inv
     cross
-    det
     qr
     svd
     eigvals
@@ -24,6 +23,5 @@ Linear Algebra
     lu
     lu_factor
     pinv
-    slogdet
     solve
     solve_triangular

docs/src/python/nn.rst

@@ -175,7 +175,6 @@ In detail:
    value_and_grad
    quantize
    average_gradients
-   fsdp_apply_gradients
 
 .. toctree::
 
@@ -184,4 +183,3 @@ In detail:
    nn/functions
    nn/losses
    nn/init
-   nn/distributed

docs/src/python/nn/distributed.rst

@@ -1,30 +0,0 @@
-.. _nn_distributed:
-
-Distributed
------------
-
-Helper Routines
-^^^^^^^^^^^^^^^
-
-The :code:`mlx.nn.layers.distributed` package contains helpful routines to 
-create sharded layers from existing :class:`Modules <mlx.nn.Module>`.
-
-.. currentmodule:: mlx.nn.layers.distributed
-.. autosummary::
-   :toctree: _autosummary
-
-   shard_linear
-   shard_inplace
-
-Layers
-^^^^^^
-
-.. currentmodule:: mlx.nn
-.. autosummary::
-   :toctree: _autosummary
-   :template: nn-module-template.rst
-
-   AllToShardedLinear
-   ShardedToAllLinear
-   QuantizedAllToShardedLinear
-   QuantizedShardedToAllLinear

docs/src/python/nn/layers.rst

@@ -10,7 +10,6 @@ Layers
    :template: nn-module-template.rst
 
    ALiBi
-   AllToShardedLinear
    AvgPool1d
    AvgPool2d
    AvgPool3d
@@ -47,10 +46,8 @@ Layers
    Mish
    MultiHeadAttention
    PReLU
-   QuantizedAllToShardedLinear
    QuantizedEmbedding
    QuantizedLinear
-   QuantizedShardedToAllLinear
    RMSNorm
    ReLU
    ReLU2
@@ -59,7 +56,6 @@ Layers
    RoPE
    SELU
    Sequential
-   ShardedToAllLinear
    Sigmoid
    SiLU
    SinusoidalPositionalEncoding

docs/src/python/printoptions.rst

@@ -1,12 +0,0 @@
-Print Options
-===============
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-  PrintOptions
-  set_printoptions
-  printoptions
-  get_printoptions

docs/src/python/transforms.rst

@@ -11,7 +11,6 @@ Transforms
    eval
    async_eval
    compile
-   checkpoint
    custom_function
    disable_compile
    enable_compile

docs/src/usage/distributed.rst

@@ -117,11 +117,89 @@ The following examples aim to clarify the backend initialization logic in MLX:
     world_ring = mx.distributed.init(backend="ring")
     world_any = mx.distributed.init()  # same as MPI because it was initialized first!
 
-Distributed Program Examples
-----------------------------
+.. _training_example:
 
-- :ref:`Data Parallelism <data_parallelism>`
-- :ref:`Tensor Parallelism <tensor_parallelism>`
+Training Example
+----------------
+
+In this section we will adapt an MLX training loop to support data parallel
+distributed training. Namely, we will average the gradients across a set of
+hosts before applying them to the model.
+
+Our training loop looks like the following code snippet if we omit the model,
+dataset and optimizer initialization.
+
+.. code:: python
+
+    model = ...
+    optimizer = ...
+    dataset = ...
+
+    def step(model, x, y):
+        loss, grads = loss_grad_fn(model, x, y)
+        optimizer.update(model, grads)
+        return loss
+
+    for x, y in dataset:
+        loss = step(model, x, y)
+        mx.eval(loss, model.parameters())
+
+All we have to do to average the gradients across machines is perform an
+:func:`all_sum` and divide by the size of the :class:`Group`. Namely we
+have to :func:`mlx.utils.tree_map` the gradients with following function.
+
+.. code:: python
+
+    def all_avg(x):
+        return mx.distributed.all_sum(x) / mx.distributed.init().size()
+
+Putting everything together our training loop step looks as follows with
+everything else remaining the same.
+
+.. code:: python
+
+    from mlx.utils import tree_map
+
+    def all_reduce_grads(grads):
+        N = mx.distributed.init().size()
+        if N == 1:
+            return grads
+        return tree_map(
+            lambda x: mx.distributed.all_sum(x) / N,
+            grads
+        )
+
+    def step(model, x, y):
+        loss, grads = loss_grad_fn(model, x, y)
+        grads = all_reduce_grads(grads)  # <--- This line was added
+        optimizer.update(model, grads)
+        return loss
+
+Utilizing ``nn.average_gradients``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Although the code example above works correctly; it performs one communication
+per gradient. It is significantly more efficient to aggregate several gradients
+together and perform fewer communication steps.
+
+This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
+almost identical to the example above:
+
+.. code:: python
+
+    model = ...
+    optimizer = ...
+    dataset = ...
+
+    def step(model, x, y):
+        loss, grads = loss_grad_fn(model, x, y)
+        grads = mx.nn.average_gradients(grads)  # <---- This line was added
+        optimizer.update(model, grads)
+        return loss
+
+    for x, y in dataset:
+        loss = step(model, x, y)
+        mx.eval(loss, model.parameters())
 
 .. _ring_section:
 

docs/src/usage/export.rst

@@ -155,34 +155,6 @@ parameters, pass them as inputs to the ``call`` wrapper:
    mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)
 
 
-Exporting with a Callback
--------------------------
-
-To inspect the exported graph, you can pass a callback instead of a file path
-to :func:`export_function`.
-
-.. code-block:: python
-
-  def fun(x):
-    return x.astype(mx.int32)
-
-  def callback(args):
-    print(args)
-
-  mx.export_function(callback, fun, mx.array([1.0, 2.0]))
-
-The argument to the callback (``args``) is a dictionary which includes a
-``type`` field. The possible types are:
-
-* ``"inputs"``: The ordered positional inputs to the exported function
-* ``"keyword_inputs"``: The keyword specified inputs to the exported function
-* ``"outputs"``: The ordered outputs of the exported function
-* ``"constants"``: Any graph constants
-* ``"primitives"``: Inner graph nodes representating the operations
-
-Each type has additional fields in the ``args`` dictionary.
-
-
 Shapeless Exports
 -----------------
 

docs/src/usage/numpy.rst

@@ -90,7 +90,10 @@ PyTorch supports the buffer protocol, but it requires an explicit
 
   a = mx.arange(3)
   b = torch.tensor(memoryview(a))
-  c = mx.array(b)
+  c = mx.array(b.numpy())
+
+Conversion from PyTorch tensors back to arrays must be done via intermediate
+NumPy arrays with ``numpy()``.
 
 JAX
 ---

examples/extensions/axpby/axpby.cpp

@@ -192,7 +192,7 @@ void Axpby::eval_gpu(
   auto kernel = d.get_kernel(kname, lib);
 
   // Prepare to encode kernel
-  auto& compute_encoder = mx::metal::get_command_encoder(s);
+  auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Kernel parameters are registered with buffer indices corresponding to

examples/extensions/pyproject.toml

@@ -3,6 +3,6 @@ requires = [
   "setuptools>=42",
   "cmake>=3.25",
   "mlx>=0.18.0",
-  "nanobind==2.12.0",
+  "nanobind==2.10.2",
 ]
 build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
 mlx>=0.21.0
-nanobind==2.12.0
+nanobind==2.10.2

mlx/CMakeLists.txt

@@ -14,7 +14,6 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/stream.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
@@ -23,57 +22,16 @@ target_sources(
 # Define MLX_VERSION only in the version.cpp file.
 add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp)
 target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}")
-target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR})
 target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:mlx_version>)
 
-# Do not export symbols by default.
-set_target_properties(
-  mlx mlx_version
-  PROPERTIES VISIBILITY_INLINES_HIDDEN ON
-             CXX_VISIBILITY_PRESET hidden
-             CUDA_VISIBILITY_PRESET hidden)
-
-# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
-set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
-if(BUILD_SHARED_LIBS)
-  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
-else()
-  target_compile_definitions(mlx PUBLIC MLX_STATIC)
-  target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
-endif()
-
-if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
-  # Supress warnings: note: parameter passing for argument of type
-  # 'std::pair<float, float>' when C++17 is enabled changed to match C++14 in
-  # GCC 10.1
-  target_compile_options(mlx PRIVATE -Wno-psabi)
-endif()
-
 if(MSVC)
-  # Some of CUDA's headers include windows.h, which defines min/max macros.
-  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
-  # Unicode support in fmt does not compile in .cu files.
-  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
   # Disable some MSVC warnings to speed up compilation.
-  target_compile_options(
-    mlx
-    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
-    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
-            /wd4146
-            /wd4700
-            /wd4804
-            /wd4805>
-            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
-            -Xcompiler=/wd4267>)
-  # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
-  # the default 65,535 section limit in COFF object files.
-  target_compile_options(
-    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>
-                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/bigobj>)
-  # Use modern preprocessor, otherwise CCCL would complain.
-  target_compile_options(
-    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
-                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>)
+  target_compile_options(mlx PUBLIC /wd4068 /wd4244 /wd4267 /wd4804)
+endif()
+
+if(WIN32)
+  # Export symbols by default to behave like macOS/linux.
+  set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE)
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)

mlx/allocator.h

@@ -4,14 +4,12 @@
 
 #include <cstdlib>
 
-#include "mlx/api.h"
-
 namespace mlx::core::allocator {
 
 // Simple wrapper around buffer pointers
 // WARNING: Only Buffer objects constructed from and those that wrap
 //          raw pointers from mlx::allocator are supported.
-class MLX_API Buffer {
+class Buffer {
  private:
   void* ptr_;
 
@@ -30,7 +28,7 @@ class MLX_API Buffer {
   };
 };
 
-class MLX_API Allocator {
+class Allocator {
   /** Abstract base class for a memory allocator. */
  public:
   virtual Buffer malloc(size_t size) = 0;
@@ -49,7 +47,7 @@ class MLX_API Allocator {
   virtual ~Allocator() = default;
 };
 
-MLX_API Allocator& allocator();
+Allocator& allocator();
 
 inline Buffer malloc(size_t size) {
   return allocator().malloc(size);

mlx/api.h

@@ -1,29 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-// MLX_API macro for controlling symbol visibility, must add for public APIs.
-//
-// Usage:
-//   MLX_API void some_function(...);
-//   class MLX_API SomeClass { ... };
-
-#if defined(MLX_STATIC)
-
-// Static library build - no import/export decorations needed
-#define MLX_API
-
-#else
-
-// Shared library build.
-#if defined(_WIN32)
-#if defined(MLX_EXPORT)
-#define MLX_API __declspec(dllexport)
-#else
-#define MLX_API __declspec(dllimport)
-#endif // defined(MLX_EXPORT)
-#else
-#define MLX_API __attribute__((visibility("default")))
-#endif // defined(_WIN32)
-
-#endif // defined(MLX_STATIC)

mlx/array.cpp

@@ -21,12 +21,11 @@ array::array(
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
     std::vector<array> inputs)
-    : array_desc_(
-          std::make_shared<ArrayDesc>(
-              std::move(shape),
-              dtype,
-              std::move(primitive),
-              std::move(inputs))) {
+    : array_desc_(std::make_shared<ArrayDesc>(
+          std::move(shape),
+          dtype,
+          std::move(primitive),
+          std::move(inputs))) {
   if (has_primitive() && this->primitive().stream().device == Device::gpu) {
     for (auto& in : this->inputs()) {
       if (in.dtype() == float64) {
@@ -70,18 +69,16 @@ array array::unsafe_weak_copy(const array& other) {
 }
 
 array::array(std::initializer_list<float> data)
-    : array_desc_(
-          std::make_shared<ArrayDesc>(
-              Shape{static_cast<ShapeElem>(data.size())},
-              float32)) {
+    : array_desc_(std::make_shared<ArrayDesc>(
+          Shape{static_cast<ShapeElem>(data.size())},
+          float32)) {
   init(data.begin());
 }
 
 array::array(std::initializer_list<int> data, Dtype dtype)
-    : array_desc_(
-          std::make_shared<ArrayDesc>(
-              Shape{static_cast<ShapeElem>(data.size())},
-              dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(
+          Shape{static_cast<ShapeElem>(data.size())},
+          dtype)) {
   init(data.begin());
 }
 
@@ -134,7 +131,6 @@ bool array::is_available() const {
   } else if (
       status() == Status::evaluated &&
       (!event().valid() || event().is_signaled())) {
-    detach_event();
     set_status(Status::available);
     return true;
   }

mlx/array.h

@@ -8,7 +8,6 @@
 #include <vector>
 
 #include "mlx/allocator.h"
-#include "mlx/api.h"
 #include "mlx/dtype.h"
 #include "mlx/event.h"
 #include "mlx/small_vector.h"
@@ -23,7 +22,7 @@ using ShapeElem = int32_t;
 using Shape = SmallVector<ShapeElem>;
 using Strides = SmallVector<int64_t>;
 
-class MLX_API array {
+class array {
   /* An array is really a node in a graph. It contains a shared ArrayDesc
    * object */
 
@@ -122,7 +121,7 @@ class MLX_API array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto shape(int dim) const {
-    return shape().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
+    return shape().at(dim < 0 ? dim + ndim() : dim);
   }
 
   /** The strides of the array. */
@@ -136,7 +135,7 @@ class MLX_API array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto strides(int dim) const {
-    return strides().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
+    return strides().at(dim < 0 ? dim + ndim() : dim);
   }
 
   /** Get the arrays data type. */
@@ -154,7 +153,7 @@ class MLX_API array {
   template <typename T>
   T item() const;
 
-  struct MLX_API ArrayIterator {
+  struct ArrayIterator {
     using iterator_category = std::random_access_iterator_tag;
     using difference_type = size_t;
     using value_type = const array;
@@ -465,7 +464,7 @@ class MLX_API array {
   template <typename It>
   void init(const It src);
 
-  struct MLX_API ArrayDesc {
+  struct ArrayDesc {
     Shape shape;
     Strides strides;
     size_t size;
@@ -489,10 +488,10 @@ class MLX_API array {
     int64_t offset{0};
 
     // The size in elements of the data buffer the array accesses
-    size_t data_size{0};
+    size_t data_size;
 
     // Contains useful meta data about the array
-    Flags flags{true, true, true};
+    Flags flags;
 
     std::vector<array> inputs;
     // An array to keep track of the siblings from a multi-output
@@ -542,10 +541,9 @@ template <typename T>
 array::array(
     std::initializer_list<T> data,
     Dtype dtype /* = TypeToDtype<T>() */)
-    : array_desc_(
-          std::make_shared<ArrayDesc>(
-              Shape{static_cast<ShapeElem>(data.size())},
-              dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(
+          Shape{static_cast<ShapeElem>(data.size())},
+          dtype)) {
   init(data.begin());
 }
 

mlx/backend/common/buffer_cache.h

@@ -2,7 +2,6 @@
 
 #pragma once
 
-#include <algorithm>
 #include <cassert>
 #include <functional>
 #include <map>

mlx/backend/common/common.cpp

@@ -19,28 +19,27 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
         "AsStrided must be used with row contiguous arrays only.");
   }
 
-  auto [no_bsx_size, row_contiguous, col_contiguous] =
-      check_contiguity(shape_, strides_);
-
-  int64_t l = 0, h = 0;
-  bool has_negative_stride = false;
-  for (int i = 0; i < strides_.size(); i++) {
-    auto delta = strides_[i] * (shape_[i] - 1);
-    if (strides_[i] >= 0) {
-      h += delta;
-    } else {
-      l += delta;
-      has_negative_stride |= shape_[i] > 1;
-    }
+  // Compute the flags given the shape and strides
+  bool row_contiguous = true, col_contiguous = true;
+  size_t r = 1, c = 1;
+  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
   }
-  size_t data_size = out.size() == 0 ? 0 : (h - l) + 1;
-
   auto flags = in.flags();
-  flags.contiguous =
-      out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size);
+  // TODO: Compute the contiguous flag in a better way cause now we are
+  //       unnecessarily strict.
+  flags.contiguous = row_contiguous || col_contiguous;
   flags.row_contiguous = row_contiguous;
   flags.col_contiguous = col_contiguous;
 
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
+
   return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 

mlx/backend/common/quantized.h

@@ -1,14 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-namespace mlx::core {
-
-inline constexpr short get_pack_factor(int bits, int wsize = 8) {
-  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
-}
-
-inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
-  bool power_of_2_bits = (bits & (bits - 1)) == 0;
-  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
-}
-
-} // namespace mlx::core

mlx/backend/common/utils.h

@@ -116,39 +116,6 @@ struct ContiguousIterator {
     loc += strides_[i];
   }
 
-  void step(int64_t s) {
-    int dims = shape_.size();
-    if (dims == 0) {
-      return;
-    }
-    int i = dims - 1;
-    while (s > 0) {
-      if (shape_[i] - pos_[i] > 1) {
-        int steps = static_cast<int>(
-            std::min(static_cast<int64_t>(shape_[i] - pos_[i] - 1), s));
-        pos_[i] += steps;
-        loc += strides_[i] * steps;
-        s -= steps;
-      } else {
-        while (pos_[i] == (shape_[i] - 1) && i > 0) {
-          pos_[i] = 0;
-          loc -= (shape_[i] - 1) * strides_[i];
-          i--;
-        }
-        pos_[i]++;
-        loc += strides_[i];
-        s--;
-      }
-    }
-  }
-
-  int64_t contiguous_suffix() {
-    if (shape_.size() == 0) {
-      return 0;
-    }
-    return (strides_.back() == 1) ? shape_.back() : 0;
-  }
-
   void seek(int64_t n) {
     loc = 0;
     for (int i = shape_.size() - 1; i >= 0; --i) {

mlx/backend/cpu/CMakeLists.txt

@@ -40,7 +40,7 @@ add_dependencies(mlx cpu_compiled_preamble)
 
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp

mlx/backend/cpu/available.cpp

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

mlx/backend/cpu/available.h

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

mlx/backend/cpu/compiled.cpp

@@ -10,6 +10,7 @@
 #include <fmt/format.h>
 
 #include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cpu/compiled_preamble.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/jit_compiler.h"
 #include "mlx/device.h"
@@ -118,15 +119,13 @@ void* compile(
     source_file.close();
 
     try {
-      JitCompiler::exec(
-          JitCompiler::build_command(
-              output_dir, source_file_name, shared_lib_name));
+      JitCompiler::exec(JitCompiler::build_command(
+          output_dir, source_file_name, shared_lib_name));
     } catch (const std::exception& error) {
-      throw std::runtime_error(
-          fmt::format(
-              "[Compile::eval_cpu] Failed to compile function {0}: {1}",
-              kernel_name,
-              error.what()));
+      throw std::runtime_error(fmt::format(
+          "[Compile::eval_cpu] Failed to compile function {0}: {1}",
+          kernel_name,
+          error.what()));
     }
   }
 
@@ -315,9 +314,7 @@ void Compiled::eval_cpu(
   // Get the function
   auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
     std::ostringstream kernel;
-    kernel << std::get<2>(JitCompiler::get_preamble()) << std::endl;
-    kernel << "using namespace mlx::core;" << std::endl;
-    kernel << "using namespace mlx::core::detail;" << std::endl;
+    kernel << get_kernel_preamble() << std::endl;
     kernel << "extern \"C\"  {" << std::endl;
     build_kernel(
         kernel,

mlx/backend/cpu/compiled_preamble.h

@@ -9,4 +9,4 @@
 #include "mlx/backend/cpu/binary_ops.h"
 // clang-format on
 
-const char* get_prebuilt_preamble();
+const char* get_kernel_preamble();

mlx/backend/cpu/device_info.cpp

@@ -1,113 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cpu/device_info.h"
-
-#ifdef __APPLE__
-#include <sys/sysctl.h>
-#include <sys/utsname.h>
-#elif defined(_WIN32)
-#include <windows.h>
-#else
-#include <sys/utsname.h>
-#include <fstream>
-#endif
-
-namespace mlx::core::cpu {
-
-namespace {
-
-// Get CPU architecture string at runtime
-std::string get_cpu_architecture() {
-#ifdef _WIN32
-  // Use GetNativeSystemInfo to get the actual hardware architecture,
-  // even when running under WoW64 emulation
-  SYSTEM_INFO sysInfo;
-  GetNativeSystemInfo(&sysInfo);
-  switch (sysInfo.wProcessorArchitecture) {
-    case PROCESSOR_ARCHITECTURE_AMD64:
-      return "x86_64";
-    case PROCESSOR_ARCHITECTURE_ARM64:
-      return "arm64";
-    case PROCESSOR_ARCHITECTURE_INTEL:
-      return "x86";
-    case PROCESSOR_ARCHITECTURE_ARM:
-      return "arm";
-    default:
-      return "unknown";
-  }
-#else
-  // Use uname() for runtime detection on Unix-like systems.
-  // This returns the actual hardware architecture (e.g., "arm64" on Apple
-  // Silicon even when running x86_64 binaries via Rosetta 2)
-  struct utsname info;
-  if (uname(&info) == 0) {
-    return std::string(info.machine);
-  }
-  return "unknown";
-#endif
-}
-
-// Get CPU device name (brand string)
-std::string get_cpu_name() {
-#ifdef __APPLE__
-  char model[256];
-  size_t len = sizeof(model);
-  if (sysctlbyname("machdep.cpu.brand_string", &model, &len, NULL, 0) == 0) {
-    return std::string(model);
-  }
-#elif defined(_WIN32)
-  // Read CPU brand string from registry
-  HKEY hKey;
-  if (RegOpenKeyExA(
-          HKEY_LOCAL_MACHINE,
-          "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0",
-          0,
-          KEY_READ,
-          &hKey) == ERROR_SUCCESS) {
-    char brand[256];
-    DWORD size = sizeof(brand);
-    if (RegQueryValueExA(
-            hKey, "ProcessorNameString", NULL, NULL, (LPBYTE)brand, &size) ==
-        ERROR_SUCCESS) {
-      RegCloseKey(hKey);
-      return std::string(brand);
-    }
-    RegCloseKey(hKey);
-  }
-#else
-  // Try reading from /proc/cpuinfo on Linux
-  std::ifstream cpuinfo("/proc/cpuinfo");
-  if (cpuinfo.is_open()) {
-    std::string line;
-    while (std::getline(cpuinfo, line)) {
-      if (line.starts_with("model name")) {
-        if (auto n = line.find(": "); n != std::string::npos) {
-          return line.substr(n + 2);
-        }
-      }
-    }
-  }
-#endif
-  return get_cpu_architecture();
-}
-
-} // anonymous namespace
-
-bool is_available() {
-  return true;
-}
-
-int device_count() {
-  return 1;
-}
-
-const std::unordered_map<std::string, std::variant<std::string, size_t>>&
-device_info(int /* device_index */) {
-  static auto info =
-      std::unordered_map<std::string, std::variant<std::string, size_t>>{
-          {"device_name", get_cpu_name()},
-          {"architecture", get_cpu_architecture()}};
-  return info;
-}
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/device_info.h

@@ -1,28 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#pragma once
-
-#include <string>
-#include <unordered_map>
-#include <variant>
-
-namespace mlx::core::cpu {
-
-bool is_available();
-
-/**
- * Get the number of available CPU devices.
- *
- * For CPU, always returns 1.
- */
-int device_count();
-
-/**
- * Get CPU device information.
- *
- * Returns a map with basic CPU device properties.
- */
-const std::unordered_map<std::string, std::variant<std::string, size_t>>&
-device_info(int device_index = 0);
-
-} // namespace mlx::core::cpu

mlx/backend/cpu/encoder.h

@@ -12,7 +12,7 @@ namespace mlx::core::cpu {
 // Number of dispatches per scheduler task
 constexpr int DISPATCHES_PER_TASK = 10;
 
-struct MLX_API CommandEncoder {
+struct CommandEncoder {
   CommandEncoder(Stream stream) : stream_(stream) {}
 
   CommandEncoder(const CommandEncoder&) = delete;
@@ -62,6 +62,6 @@ struct MLX_API CommandEncoder {
   int num_ops_{0};
 };
 
-MLX_API CommandEncoder& get_command_encoder(Stream stream);
+CommandEncoder& get_command_encoder(Stream stream);
 
 } // namespace mlx::core::cpu

mlx/backend/cpu/indexing.cpp

@@ -4,14 +4,11 @@
 #include <cmath>
 
 #include "mlx/allocator.h"
+#include "mlx/primitives.h"
+
 #include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/binary.h"
-#include "mlx/backend/cpu/binary_ops.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
-#include "mlx/backend/cpu/slicing.h"
-#include "mlx/dtype_utils.h"
-#include "mlx/primitives.h"
 
 namespace mlx::core {
 
@@ -764,7 +761,7 @@ void masked_scatter_impl(const array& mask, const array& src, array& out) {
   const size_t mask_batch_size = mask.size() / batch_count;
   const size_t src_batch_size = src.size() / batch_count;
 
-  for (size_t b = 0; b < batch_count; ++b) {
+  for (uint b = 0; b < batch_count; ++b) {
     size_t src_consumed = 0;
     src_it.seek(b * src_batch_size);
 
@@ -791,7 +788,7 @@ void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& mask = inputs[1];
   auto& src = inputs[2];
 
-  // Copy dst into out (copy allocates memory for out)
+  // 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());
@@ -854,128 +851,4 @@ void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
   });
 }
 
-template <typename T, typename Op>
-void slice_update_impl(
-    array& out,
-    const array& upd,
-    int64_t data_offset,
-    const Strides& out_strides) {
-  ContiguousIterator out_it(upd.shape(), out_strides, upd.ndim());
-  ContiguousIterator upd_it(upd);
-  Op op;
-
-  constexpr int SIMD_START = 32;
-
-  T* out_ptr = out.data<T>() + data_offset;
-  const T* upd_ptr = upd.data<T>();
-  int64_t size = upd.size();
-  int64_t suffix = out_it.contiguous_suffix();
-
-  if (upd.data_size() == 1) {
-    if (suffix >= SIMD_START) {
-      for (int64_t i = 0; i < size; i += suffix) {
-        VectorScalar<Op>{}(
-            out_ptr + out_it.loc, upd_ptr, out_ptr + out_it.loc, suffix);
-        out_it.step(suffix);
-      }
-    } else {
-      T update = upd_ptr[0];
-      for (int64_t i = 0; i < size; i++) {
-        out_ptr[out_it.loc] = op(out_ptr[out_it.loc], update);
-        out_it.step();
-      }
-    }
-  } else if (suffix == upd_it.contiguous_suffix() && suffix >= SIMD_START) {
-    for (int64_t i = 0; i < size; i += suffix) {
-      VectorVector<Op>{}(
-          out_ptr + out_it.loc,
-          upd_ptr + upd_it.loc,
-          out_ptr + out_it.loc,
-          suffix);
-      out_it.step(suffix);
-      upd_it.step(suffix);
-    }
-  } else {
-    for (int64_t i = 0; i < size; i++) {
-      out_ptr[out_it.loc] = op(out_ptr[out_it.loc], upd_ptr[upd_it.loc]);
-      out_it.step();
-      upd_it.step();
-    }
-  }
-}
-
-void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
-    return;
-  }
-
-  auto& in = inputs[0];
-  auto& upd = inputs[1];
-
-  if (upd.size() == 0) {
-    out.copy_shared_buffer(in);
-    return;
-  }
-
-  // Check if materialization is needed
-  auto ctype = in.flags().contiguous && in.size() == in.data_size()
-      ? CopyType::Vector
-      : CopyType::General;
-  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
-
-  // Calculate out strides, initial offset and if copy needs to be made
-  auto [data_offset, out_strides] =
-      prepare_slice(out, start_indices_, strides_);
-
-  // Do copy
-  if (reduce_type_ == SliceUpdate::None) {
-    copy_cpu_inplace(
-        /* const array& src = */ upd,
-        /* array& dst = */ out,
-        /* const std::vector<int>& data_shape = */ upd.shape(),
-        /* const std::vector<stride_t>& i_strides = */ upd.strides(),
-        /* const std::vector<stride_t>& o_strides = */ out_strides,
-        /* int64_t i_offset = */ 0,
-        /* int64_t o_offset = */ data_offset,
-        /* CopyType ctype = */ CopyType::GeneralGeneral,
-        stream());
-    return;
-  }
-
-  auto& encoder = cpu::get_command_encoder(stream());
-  encoder.set_input_array(upd);
-  encoder.set_output_array(out);
-  encoder.dispatch([upd = array::unsafe_weak_copy(upd),
-                    out = array::unsafe_weak_copy(out),
-                    data_offset = data_offset,
-                    out_strides = std::move(out_strides),
-                    reduce_type = reduce_type_]() mutable {
-    dispatch_all_types(out.dtype(), [&](auto type_tag) {
-      using T = MLX_GET_TYPE(type_tag);
-      switch (reduce_type) {
-        case SliceUpdate::Sum:
-          slice_update_impl<T, detail::Add>(out, upd, data_offset, out_strides);
-          break;
-        case SliceUpdate::Prod:
-          slice_update_impl<T, detail::Multiply>(
-              out, upd, data_offset, out_strides);
-          break;
-        case SliceUpdate::Max:
-          slice_update_impl<T, detail::Maximum>(
-              out, upd, data_offset, out_strides);
-          break;
-        case SliceUpdate::Min:
-          slice_update_impl<T, detail::Minimum>(
-              out, upd, data_offset, out_strides);
-          break;
-        case SliceUpdate::None:
-          // Should never be here
-          break;
-      }
-    });
-  });
-}
-
 } // namespace mlx::core

mlx/backend/cpu/jit_compiler.cpp

@@ -1,8 +1,6 @@
 // Copyright © 2024 Apple Inc.
 
 #include "mlx/backend/cpu/jit_compiler.h"
-#include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/compiled_preamble.h"
 
 #include <algorithm>
 #include <sstream>
@@ -36,30 +34,18 @@ struct VisualStudioInfo {
     arch = "x64";
 #endif
     // Get path of Visual Studio.
-    // Use -latest to get only the most recent installation when multiple
-    // versions are installed, avoiding path concatenation issues.
-    std::string vs_path = JitCompiler::exec(
-        fmt::format(
-            "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
-            " -latest -property installationPath",
-            std::getenv("ProgramFiles(x86)")));
+    std::string vs_path = JitCompiler::exec(fmt::format(
+        "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
+        " -property installationPath",
+        std::getenv("ProgramFiles(x86)")));
     if (vs_path.empty()) {
       throw std::runtime_error("Can not find Visual Studio.");
     }
-    // Trim any trailing whitespace/newlines from the path
-    vs_path.erase(
-        std::find_if(
-            vs_path.rbegin(),
-            vs_path.rend(),
-            [](unsigned char ch) { return !std::isspace(ch); })
-            .base(),
-        vs_path.end());
     // Read the envs from vcvarsall.
-    std::string envs = JitCompiler::exec(
-        fmt::format(
-            "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
-            vs_path,
-            arch));
+    std::string envs = JitCompiler::exec(fmt::format(
+        "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
+        vs_path,
+        arch));
     for (const std::string& line : str_split(envs, '\n')) {
       // Each line is in the format "ENV_NAME=values".
       auto pos = line.find_first_of('=');
@@ -88,61 +74,30 @@ const VisualStudioInfo& GetVisualStudioInfo() {
 
 #endif // _MSC_VER
 
-const std::tuple<bool, std::string, std::string>& JitCompiler::get_preamble() {
-  static auto preamble = []() -> std::tuple<bool, std::string, std::string> {
-    // Check whether the headers are shipped with the binary, if so use the
-    // preamble from the headers, otherwise use the prebuilt one embeded in
-    // binary, which may not work with all compilers.
-    auto root_dir = current_binary_dir();
-#if !defined(_WIN32)
-    root_dir = root_dir.parent_path();
-#endif
-    auto include_dir = root_dir / "include";
-    if (std::filesystem::exists(include_dir / "mlx")) {
-      return std::make_tuple(
-          true,
-          include_dir.string(),
-          "#include \"mlx/backend/cpu/compiled_preamble.h\"\n");
-    } else {
-      return std::make_tuple(false, "", get_prebuilt_preamble());
-    }
-  }();
-  return preamble;
-}
-
 std::string JitCompiler::build_command(
     const std::filesystem::path& dir,
     const std::string& source_file_name,
     const std::string& shared_lib_name) {
-  auto& [use_include, include_dir, preamble] = get_preamble();
 #ifdef _MSC_VER
-  std::string extra_flags;
-  if (use_include) {
-    extra_flags += fmt::format("/I \"{}\"", include_dir);
-  }
   const VisualStudioInfo& info = GetVisualStudioInfo();
+  std::string libpaths;
   for (const std::string& lib : info.libpaths) {
-    extra_flags += fmt::format(" /libpath:\"{}\"", lib);
+    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
   }
   return fmt::format(
       "\""
-      "cd /D \"{}\" && "
-      "\"{}\" /LD /EHsc /MD /Ox /nologo /std:c++17 {} \"{}\" "
-      "/link /out:\"{}\" 2>&1"
+      "cd /D \"{0}\" && "
+      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
+      "/link /out:\"{3}\" {4} 2>&1"
       "\"",
       dir.string(),
       info.cl_exe,
-      extra_flags,
       source_file_name,
-      shared_lib_name);
+      shared_lib_name,
+      libpaths);
 #else
-  std::string extra_flags;
-  if (use_include) {
-    extra_flags = fmt::format("-I \"{}\"", include_dir);
-  }
   return fmt::format(
-      "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1",
-      extra_flags,
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1",
       (dir / source_file_name).string(),
       (dir / shared_lib_name).string());
 #endif
@@ -185,13 +140,12 @@ std::string JitCompiler::exec(const std::string& cmd) {
   int code = WEXITSTATUS(status);
 #endif
   if (code != 0) {
-    throw std::runtime_error(
-        fmt::format(
-            "Failed to execute command with return code {0}: \"{1}\", "
-            "the output is: {2}",
-            code,
-            cmd,
-            ret));
+    throw std::runtime_error(fmt::format(
+        "Failed to execute command with return code {0}: \"{1}\", "
+        "the output is: {2}",
+        code,
+        cmd,
+        ret));
   }
   return ret;
 }

mlx/backend/cpu/jit_compiler.h

@@ -7,9 +7,6 @@ namespace mlx::core {
 
 class JitCompiler {
  public:
-  // Return the includes that should be prepended to the source code.
-  static const std::tuple<bool, std::string, std::string>& get_preamble();
-
   // Build a shell command that compiles a source code file to a shared library.
   static std::string build_command(
       const std::filesystem::path& dir,

mlx/backend/cpu/luf.cpp

@@ -67,10 +67,11 @@ void luf_impl(
               /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
               /* info */ &info);
 
-          if (info < 0) {
+          if (info != 0) {
             std::stringstream ss;
             ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-               << " because argument had an illegal value";
+               << ((info > 0) ? " because matrix is singular"
+                              : " because argument had an illegal value");
             throw std::runtime_error(ss.str());
           }
 

mlx/backend/cpu/make_compiled_preamble.ps1

@@ -15,6 +15,13 @@ $CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
 # Concatenate to string.
 $CONTENT = $CONTENT -join "`n"
 
+# Append extra content.
+$CONTENT = @"
+$($CONTENT)
+using namespace mlx::core;
+using namespace mlx::core::detail;
+"@
+
 # Convert each char to ASCII code.
 # Unlike the unix script that outputs string literal directly, the output from
 # MSVC is way too large to be embedded as string and compilation will fail, so
@@ -22,7 +29,7 @@ $CONTENT = $CONTENT -join "`n"
 $CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
 
 $OUTPUT = @"
-const char* get_prebuilt_preamble() {
+const char* get_kernel_preamble() {
   static char preamble[] = { $CHARCODES };
   return preamble;
 }

mlx/backend/cpu/make_compiled_preamble.sh

@@ -30,10 +30,12 @@ fi
 CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
-const char* get_prebuilt_preamble() {
+const char* get_kernel_preamble() {
 return R"preamble(
 $INCLUDES
 $CONTENT
+using namespace mlx::core;
+using namespace mlx::core::detail;
 )preamble";
 }
 EOF

mlx/backend/cpu/primitives.cpp

@@ -398,6 +398,44 @@ void DynamicSliceUpdate::eval_cpu(
   }
 }
 
+void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] =
+      prepare_slice(out, start_indices_, strides_);
+
+  // Do copy
+  copy_cpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+      /* const std::vector<stride_t>& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream());
+}
+
 void View::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];

mlx/backend/cpu/quantized.cpp

@@ -1,6 +1,5 @@
 // Copyright © 2023 Apple Inc.
 
-#include "mlx/backend/common/quantized.h"
 #include "mlx/backend/common/unary.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
@@ -15,19 +14,6 @@ namespace mlx::core {
 
 namespace {
 
-array ensure_row_contiguous(
-    const array& arr,
-    cpu::CommandEncoder& encoder,
-    Stream s) {
-  if (arr.flags().row_contiguous) {
-    return arr;
-  } else {
-    auto arr_cpy = contiguous_copy_cpu(arr, s);
-    encoder.add_temporary(arr_cpy);
-    return arr_cpy;
-  }
-};
-
 const static float FP4_LUT[16] = {
     +0.0f,
     +0.5f,
@@ -61,6 +47,15 @@ static inline T dequantize_scale(uint8_t s) {
   }
 }
 
+inline constexpr short get_pack_factor(int bits, int wsize = 8) {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
+  auto power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
 template <typename T, int bits>
 void extract_bits(const uint8_t* w_in, T* w_out) {
   static_assert(bits == 3 || bits == 5 || bits == 6);
@@ -927,9 +922,20 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& scales_pre = inputs[2];
 
   auto& encoder = cpu::get_command_encoder(stream());
-  auto x = ensure_row_contiguous(x_pre, encoder, stream());
-  auto w = ensure_row_contiguous(w_pre, encoder, stream());
-  auto scales = ensure_row_contiguous(scales_pre, encoder, stream());
+  auto ensure_row_contiguous = [s = stream(), &encoder](const array& arr) {
+    if (arr.flags().row_contiguous) {
+      return arr;
+    } else {
+      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
+      copy_cpu(arr, arr_cpy, CopyType::General, s);
+      encoder.add_temporary(arr_cpy);
+      return arr_cpy;
+    }
+  };
+
+  auto x = ensure_row_contiguous(x_pre);
+  auto w = ensure_row_contiguous(w_pre);
+  auto scales = ensure_row_contiguous(scales_pre);
 
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -938,7 +944,7 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   encoder.set_input_array(scales);
   encoder.set_output_array(out);
   if (mode_ == QuantizationMode::Affine) {
-    auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
+    auto biases = ensure_row_contiguous(inputs[3]);
     encoder.set_input_array(biases);
     encoder.dispatch([out = array::unsafe_weak_copy(out),
                       x = array::unsafe_weak_copy(x),
@@ -1046,105 +1052,6 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
-uint8_t to_fp8_e8m0(float x) {
-  if (!std::isfinite(x)) {
-    return 0xFF;
-  }
-  if (x < 0.0f) {
-    return 0x00;
-  }
-  float le = std::log2(x);
-  int n = int(std::round(le));
-
-  n = n < -127 ? -127 : n;
-  n = n > 127 ? 127 : n;
-  return static_cast<uint8_t>(n + 127);
-}
-
-uint8_t to_fp4_e2m1(float x) {
-  if (std::isnan(x)) {
-    return 0x7;
-  }
-
-  const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
-  x = std::abs(x);
-
-  uint8_t bits;
-  if (x > 5.0f) {
-    bits = 0x7;
-  } else if (x >= 3.5f) {
-    bits = 0x6;
-  } else if (x > 2.5f) {
-    bits = 0x5;
-  } else if (x >= 1.75f) {
-    bits = 0x4;
-  } else if (x > 1.25f) {
-    bits = 0x3;
-  } else if (x >= 0.75f) {
-    bits = 0x2;
-  } else if (x > 0.25f) {
-    bits = 0x1;
-  } else {
-    bits = 0x0;
-  }
-  return bits | sign_bit;
-}
-
-template <typename T>
-void fp_quantize_dequantize(
-    const array& w_arr,
-    array& out_arr,
-    int bits,
-    int group_size,
-    size_t w_size) {
-  auto w = w_arr.data<T>();
-  auto out = out_arr.data<T>();
-
-  size_t n_groups = w_size / group_size;
-
-  for (size_t i = 0; i < n_groups; ++i) {
-    size_t idx = i * group_size;
-    float scale = -std::numeric_limits<float>::infinity();
-    for (int j = 0; j < group_size; ++j) {
-      scale = std::max(scale, std::abs(w[idx + j]));
-    }
-    scale /= bits == 4 ? 6.0f : 448.0f;
-    if (group_size == 16) {
-      scale = dequantize_scale<float, 16>(detail::ToFP8()(scale));
-    } else {
-      scale = dequantize_scale<float, 32>(to_fp8_e8m0(scale));
-    }
-
-    for (int j = 0; j < group_size; ++j) {
-      float w_el = scale == 0 ? 0.0f : w[idx + j] / scale;
-      float output;
-      if (bits == 8) {
-        output = detail::FromFP8()(detail::ToFP8()(w_el));
-      } else {
-        output = FP4_LUT[to_fp4_e2m1(w_el)];
-      }
-      out[idx + j] = static_cast<T>(scale * output);
-    }
-  }
-}
-
-void dispatch_quantize_dequantize(
-    const array& w,
-    array& out,
-    int bits,
-    int group_size) {
-  if (w.dtype() == float16) {
-    fp_quantize_dequantize<float16_t>(w, out, bits, group_size, w.size());
-  } else if (w.dtype() == bfloat16) {
-    fp_quantize_dequantize<bfloat16_t>(w, out, bits, group_size, w.size());
-  } else if (w.dtype() == float32) {
-    fp_quantize_dequantize<float>(w, out, bits, group_size, w.size());
-  } else {
-    throw std::runtime_error(
-        "[quantize_dequantize] Only supports floating point inputs");
-  }
-}
-
 template <typename T, typename U>
 void quantize(
     const T* w,
@@ -1229,8 +1136,15 @@ void dispatch_quantize(
 void fast::Quantize::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  auto& encoder = cpu::get_command_encoder(stream());
-  auto w = ensure_row_contiguous(inputs[0], encoder, stream());
+  auto ensure_row_contiguous = [s = stream()](const array& arr) {
+    if (arr.flags().row_contiguous) {
+      return std::make_pair(arr, false);
+    } else {
+      return std::make_pair(contiguous_copy_cpu(arr, s), true);
+    }
+  };
+
+  auto [w, copied] = ensure_row_contiguous(inputs[0]);
   auto& out = outputs[0];
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -1238,6 +1152,10 @@ void fast::Quantize::eval_cpu(
   auto& biases = outputs[2];
   scales.set_data(allocator::malloc(scales.nbytes()));
   biases.set_data(allocator::malloc(biases.nbytes()));
+  auto& encoder = cpu::get_command_encoder(stream());
+  if (copied) {
+    encoder.add_temporary(w);
+  }
   encoder.set_input_array(w);
   encoder.set_input_array(scales);
   encoder.set_input_array(biases);
@@ -1320,43 +1238,6 @@ void fast::ConvertFP8::eval_cpu(
 }
 
 void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
-  auto& encoder = cpu::get_command_encoder(stream());
-
-  bool w_quantized = (inputs[1].dtype() == uint32);
-  if (w_quantized && inputs[0].shape(-2) == 1) {
-    bool donate_x = inputs[0].is_donatable();
-    auto x = ensure_row_contiguous(inputs[0], encoder, stream());
-    auto w = ensure_row_contiguous(inputs[1], encoder, stream());
-    auto scales = ensure_row_contiguous(inputs[2], encoder, stream());
-
-    out.set_data(allocator::malloc(out.nbytes()));
-
-    // If x is a copy it should be donatable
-    donate_x |= x.is_donatable();
-    auto xhat = donate_x
-        ? x
-        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
-    if (!donate_x) {
-      encoder.add_temporary(xhat);
-    }
-    encoder.set_input_array(x);
-    encoder.set_input_array(w);
-    encoder.set_input_array(scales);
-    encoder.set_output_array(out);
-    encoder.dispatch([out = array::unsafe_weak_copy(out),
-                      x = array::unsafe_weak_copy(x),
-                      xhat = array::unsafe_weak_copy(xhat),
-                      w = array::unsafe_weak_copy(w),
-                      scales = array::unsafe_weak_copy(scales),
-                      group_size_ = group_size_,
-                      bits_ = bits_]() mutable {
-      dispatch_quantize_dequantize(x, xhat, bits_, group_size_);
-      fp_qmm_dispatch(out, xhat, w, scales, group_size_, bits_, true);
-    });
-    return;
-  } else {
-    throw std::runtime_error("[QQMatmul] NYI for the general case");
-  }
+  throw std::runtime_error("QQMatmul not implemented on CPU.");
 }
-
 } // namespace mlx::core

mlx/backend/cpu/simd/base_simd.h

@@ -1,18 +1,11 @@
 #pragma once
 
-// Required for using M_LN2 in MSVC.
-#define _USE_MATH_DEFINES
-
-#include <math.h>
 #include <stdint.h>
 #include <algorithm>
+#include <cmath>
 #include <complex>
 #include <functional>
 
-#ifdef _MSC_VER
-#include <intrin.h> // For _BitScanReverse
-#endif
-
 namespace mlx::core::simd {
 template <typename T, int N>
 struct Simd;
@@ -29,14 +22,6 @@ struct Simd<T, 1> {
   Simd(Simd<U, 1> v) : value(v.value) {}
   template <typename U>
   Simd(U v) : value(v) {}
-
-  T operator[](int) const {
-    return value;
-  }
-
-  T& operator[](int) {
-    return value;
-  }
 };
 
 template <typename T, int N>
@@ -120,7 +105,7 @@ Simd<T, 1> log1p(Simd<T, 1> in) {
       if (r == 0) { // handle underflow
         return Simd<T, 1>{T{x, theta}};
       }
-      return Simd<T, 1>{T{((decltype(x))(0.5)) * std::log1p(r), theta}};
+      return Simd<T, 1>{T{((typeof(x))(0.5)) * std::log1p(r), theta}};
     } else {
       auto z0 = std::hypot(x + 1, y);
       return Simd<T, 1>{T{std::log(z0), theta}};
@@ -188,16 +173,7 @@ DEFAULT_BINARY(||)
 
 template <typename T>
 Simd<T, 1> clz(Simd<T, 1> x_) {
-#ifdef _MSC_VER
-  // MSVC doesn't have __builtin_clz, use _BitScanReverse instead
-  unsigned long index;
-  if (_BitScanReverse(&index, static_cast<unsigned long>(x_.value))) {
-    return static_cast<T>(31 - index);
-  }
-  return static_cast<T>(32); // All zeros case
-#else
   return __builtin_clz(x_.value);
-#endif
 }
 
 template <typename T>

mlx/backend/cpu/sort.cpp

@@ -15,14 +15,10 @@ namespace mlx::core {
 
 namespace {
 
-template <typename T>
-inline constexpr bool is_floating_v = std::is_floating_point_v<T> ||
-    std::is_same_v<T, float16_t> || std::is_same_v<T, bfloat16_t>;
-
 // NaN-aware comparator that places NaNs at the end
 template <typename T>
 bool nan_aware_less(T a, T b) {
-  if constexpr (is_floating_v<T> || std::is_same_v<T, complex64_t>) {
+  if constexpr (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
     if (std::isnan(a))
       return false;
     if (std::isnan(b))
@@ -107,11 +103,11 @@ struct StridedIterator {
     return *this;
   }
 
-  StridedIterator operator+(difference_type diff) const {
+  StridedIterator operator+(difference_type diff) {
     return StridedIterator(ptr_, stride_, diff);
   }
 
-  StridedIterator operator-(difference_type diff) const {
+  StridedIterator operator-(difference_type diff) {
     return StridedIterator(ptr_, stride_, -diff);
   }
 
@@ -202,7 +198,7 @@ void argsort(const array& in, array& out, int axis) {
       auto v2 = data_ptr[b * in_stride];
 
       // Handle NaNs (place them at the end)
-      if constexpr (is_floating_v<T>) {
+      if (std::is_floating_point<T>::value) {
         if (std::isnan(v1))
           return false;
         if (std::isnan(v2))
@@ -303,7 +299,7 @@ void argpartition(const array& in, array& out, int axis, int kth) {
       auto v2 = data_ptr[b * in_stride];
 
       // Handle NaNs (place them at the end)
-      if constexpr (is_floating_v<T>) {
+      if (std::is_floating_point<T>::value) {
         if (std::isnan(v1))
           return false;
         if (std::isnan(v2))

mlx/backend/cpu/unary_ops.h

@@ -155,18 +155,11 @@ struct FromFP8 {
   template <int N>
   Simd<float, N> operator()(Simd<uint8_t, N> x) {
     auto v = Simd<uint16_t, N>(x & 127) << 7;
-    Simd<float, N> out;
-    if constexpr (simd::max_size<float16_t> >= N) {
-      auto converted = *(Simd<float16_t, N>*)(&v);
-      out = converted * 256.0;
-    } else {
-      for (int i = 0; i < N; ++i) {
-        auto converted = *(float16_t*)(&v[i]);
-        out[i] = converted * 256.0;
-      }
-    }
+    auto converted = *(Simd<float16_t, N>*)(&v);
+    converted = converted * 256.0;
     auto sign = Simd<bool, N>(x & 128);
-    return select(sign, -out, out);
+    Simd<float, N> out = select(sign, -converted, converted);
+    return out;
   }
   float operator()(uint8_t x) {
     return (*this)(Simd<uint8_t, 1>(x)).value;

mlx/backend/cuda/CMakeLists.txt

@@ -19,21 +19,17 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/event.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/block_mask.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gather_gemm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
@@ -61,24 +57,21 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary)
-add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
 
 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
   target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
-  target_sources(mlx
-                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp)
 else()
   target_sources(
-    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
-                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
+    mlx
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu)
 endif()
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
@@ -120,26 +113,22 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
-# Ignore warnings from CUTLASS.
-target_compile_options(
-  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=2908,2361">)
-
-if(NOT MSVC)
-  # Required for generating optimized CUTLASS code.
+# CUDA 12.8 emits warning #20280-D for copy kernels which is a false positive.
+# Explicitly pass this flag to suppress the warning, it is safe to set it to
+# true but the warning wouldn't be suppressed.
+if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
   target_compile_options(
-    mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
+    mlx
+    PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--static-global-template-stub=false>")
 endif()
 
-# Suppress nvcc warnings on C++ headers.
+# Suppress warning when building for compute capability 7 used by V100.
 target_compile_options(
-  mlx
-  PRIVATE
-    $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=27,997,1394,20011,20208">
-)
+  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
 
-# Ignore some valid nvcc warnings, we might want to fix them in future.
-target_compile_options(
-  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=177,550">)
+# Suppress nvcc warnings on MLX headers.
+target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
+                                   --diag_suppress=997>)
 
 # Use stronger binaries compression. This feature was introduced in CUDA 12.8
 # and requires drivers released after CUDA 12.4.
@@ -154,11 +143,12 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
     COMMAND __nvcc_device_query
     OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
     OUTPUT_STRIP_TRAILING_WHITESPACE)
+  set(UPGRADABLE_ARCHITECTURES "90;100;121")
   if(MLX_CUDA_ARCHITECTURES STREQUAL "")
     message(
       FATAL_ERROR
         "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES")
-  elseif(MLX_CUDA_ARCHITECTURES GREATER_EQUAL 90)
+  elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES)
     # Use arch-specific compute capability whenever possible.
     set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
   endif()
@@ -167,52 +157,16 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
-# Skip Hopper-only kernels when not building for sm90a.
-if(("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
-                                              MLX_CUDA_ARCHITECTURES))
-  target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED)
-endif()
-
-# Search CUDA libs from installed python packages.
-if(WIN32)
-  # Resolve paths of unfound DLL at runtime.
-  if(BUILD_SHARED_LIBS)
-    target_link_libraries(mlx PRIVATE "delayimp.lib")
-    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp)
-  else()
-    # For static library the delayload must be compiled into final executables.
-    target_link_libraries(mlx PUBLIC "delayimp.lib")
-    target_sources(
-      mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp>)
-  endif()
-  # Get all the CUDA DLLs we could link with.
-  file(
-    GLOB CUDA_DLL_NAMES
-    RELATIVE "${CUDAToolkit_BIN_DIR}/x64"
-    "${CUDAToolkit_BIN_DIR}/x64/*.dll")
-  # Delay load CUDA and cuDNN libs.
-  foreach(CUDA_DLL ${CUDA_DLL_NAMES} ${CUDNN_DLL_NAMES})
-    target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}")
-  endforeach()
-  # Pass the locations where CUDA DLLs are placed.
-  if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
-    target_compile_definitions(
-      mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64"
-                 MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}")
-  endif()
-else()
-  # For POSIX we rely on RPATH to search for CUDA libs.
-  if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
-    set_property(
-      TARGET mlx
-      APPEND
-      PROPERTY INSTALL_RPATH
-               # The paths here should match the install_requires in setup.py.
-               "$ORIGIN/../../nvidia/cublas/lib"
-               "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
-               "$ORIGIN/../../nvidia/cudnn/lib"
-               "$ORIGIN/../../nvidia/nccl/lib")
-  endif()
+if(MLX_BUILD_PYTHON_BINDINGS)
+  set_property(
+    TARGET mlx
+    APPEND
+    PROPERTY INSTALL_RPATH
+             # The paths here should match the install_requires in setup.py.
+             "$ORIGIN/../../nvidia/cublas/lib"
+             "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
+             "$ORIGIN/../../nvidia/cudnn/lib"
+             "$ORIGIN/../../nvidia/nccl/lib")
 endif()
 
 # ------------------------ Dependencies ------------------------
@@ -220,7 +174,7 @@ endif()
 # Use fixed version of CCCL.
 FetchContent_Declare(
   cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v3.1.3/cccl-v3.1.3.zip")
+  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")
 
@@ -248,14 +202,12 @@ FetchContent_MakeAvailable(nvtx3)
 target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)
 
 # Make cuda runtime APIs available in non-cuda files.
+find_package(CUDAToolkit REQUIRED)
 target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
 
 # Use cublasLt.
 target_link_libraries(mlx PRIVATE CUDA::cublasLt)
 
-# Use cuFFT.
-target_link_libraries(mlx PRIVATE CUDA::cufft)
-
 # Use NVRTC and driver APIs.
 target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver)
 
@@ -273,15 +225,16 @@ set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF)
 FetchContent_MakeAvailable(cudnn)
 target_link_libraries(mlx PRIVATE cudnn_frontend)
 # Link with the actual cuDNN libraries.
+include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 
 # Use header-only CUTLASS.
 FetchContent_Declare(
   cutlass
   GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-  GIT_TAG v4.4.2
+  GIT_TAG v4.3.2
   GIT_SHALLOW TRUE
   SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(cutlass)
 target_include_directories(
-  mlx SYSTEM PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
+  mlx PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)

mlx/backend/cuda/allocator.cpp

@@ -3,17 +3,13 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/utils.h"
-#include "mlx/backend/gpu/device_info.h"
-#include "mlx/memory.h"
-#include "mlx/scheduler.h"
 #include "mlx/utils.h"
 
 #include <cuda_runtime.h>
 #include <fmt/format.h>
+#include <unistd.h>
 
 #include <cassert>
-#include <fstream>
-#include <string>
 
 namespace mlx::core {
 
@@ -24,70 +20,6 @@ constexpr int page_size = 16384;
 // Any allocations smaller than this will try to use the small pool
 constexpr int small_block_size = 8;
 
-// The small pool size in bytes. This should be a multiple of the host page
-// size and small_block_size.
-constexpr int small_pool_size = 4 * page_size;
-
-// Check if running on Windows or Windows Subsystem for Linux
-bool is_windows() {
-#if defined(_WIN32)
-  return true;
-#elif defined(__linux__)
-  // WSL kernels contain "microsoft" or "WSL" in /proc/version
-  static bool is_wsl = []() {
-    std::ifstream version("/proc/version");
-    if (version.is_open()) {
-      std::string line;
-      std::getline(version, line);
-      return line.find("microsoft") != std::string::npos ||
-          line.find("Microsoft") != std::string::npos ||
-          line.find("WSL") != std::string::npos;
-    }
-    return false;
-  }();
-  return is_wsl;
-#else
-  return false;
-#endif
-}
-
-bool supports_managed_memory() {
-  static bool managed_memory = []() {
-    int device_count = gpu::device_count();
-    for (int i = 0; i < device_count; ++i) {
-      auto& d = cu::device(i);
-      if (!d.managed_memory()) {
-        return false;
-      }
-      // Empirically on Windows (and WSL) if there is no concurrentManagedAccess
-      // the managed memory also does not work.
-      if (is_windows() && !d.concurrent_managed_access()) {
-        return false;
-      }
-    }
-    return true;
-  }();
-  return managed_memory;
-}
-
-inline void* unified_malloc(size_t size) {
-  void* data = nullptr;
-  if (supports_managed_memory()) {
-    CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
-  } else {
-    CHECK_CUDA_ERROR(cudaMallocHost(&data, size));
-  }
-  return data;
-}
-
-inline void unified_free(void* data) {
-  if (supports_managed_memory()) {
-    CHECK_CUDA_ERROR(cudaFree(data));
-  } else {
-    CHECK_CUDA_ERROR(cudaFreeHost(data));
-  }
-}
-
 #if CUDART_VERSION >= 13000
 inline cudaMemLocation cuda_mem_loc(int i) {
   cudaMemLocation loc;
@@ -101,21 +33,24 @@ inline int cuda_mem_loc(int i) {
 }
 #endif // CUDART_VERSION >= 13000
 
+// The small pool size in bytes. This should be a multiple of the host page
+// size and small_block_size.
+constexpr int small_pool_size = 4 * page_size;
+
 SmallSizePool::SmallSizePool() {
   auto num_blocks = small_pool_size / small_block_size;
   buffer_ = new Block[num_blocks];
+
   next_free_ = buffer_;
 
-  data_ = unified_malloc(small_pool_size);
-  if (supports_managed_memory()) {
-    int device_count = gpu::device_count();
-    for (int i = 0; i < device_count; ++i) {
-      if (device(i).concurrent_managed_access()) {
-        auto loc = cuda_mem_loc(i);
-        CHECK_CUDA_ERROR(cudaMemAdvise(
-            data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
-      }
-    }
+  CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));
+
+  int device_count = 0;
+  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
+  for (int i = 0; i < device_count; ++i) {
+    auto loc = cuda_mem_loc(i);
+    CHECK_CUDA_ERROR(
+        cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
   }
 
   auto curr = next_free_;
@@ -127,7 +62,7 @@ SmallSizePool::SmallSizePool() {
 }
 
 SmallSizePool::~SmallSizePool() {
-  unified_free(data_);
+  CHECK_CUDA_ERROR(cudaFree(data_));
   delete[] buffer_;
 }
 
@@ -161,23 +96,39 @@ CudaAllocator::CudaAllocator()
     : buffer_cache_(
           page_size,
           [](CudaBuffer* buf) { return buf->size; },
-          [this](CudaBuffer* buf) { free_cuda_buffer(buf); }) {
+          [this](CudaBuffer* buf) { cuda_free(buf); }) {
   size_t free;
   CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total_memory_));
   memory_limit_ = total_memory_ * 0.95;
   free_limit_ = total_memory_ - memory_limit_;
   max_pool_size_ = memory_limit_;
 
-  int device_count = gpu::device_count();
-  free_streams_.resize(device_count);
-  mem_pools_.resize(device_count);
+  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) {
-    auto& d = device(i);
-    if (d.memory_pools()) {
-      free_streams_[i] = CudaStream(d);
-      CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i));
-    }
+    CHECK_CUDA_ERROR(cudaSetDevice(i));
+    cudaStream_t s;
+    CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));
+    free_streams_.push_back(s);
+
+    cudaMemPool_t mem_pool;
+    CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pool, i));
+    mem_pools_.push_back(mem_pool);
   }
+  CHECK_CUDA_ERROR(cudaSetDevice(curr));
+}
+
+void copy_to_managed(CudaBuffer& buf) {
+  // TODO maybe make this async on a i/o stream to avoid synchronizing the
+  // device on malloc/and free
+  void* new_data;
+  CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, buf.size));
+  buf.device = -1;
+  CHECK_CUDA_ERROR(cudaMemcpy(new_data, buf.data, buf.size, cudaMemcpyDefault));
+  CHECK_CUDA_ERROR(cudaFree(buf.data));
+  buf.data = new_data;
 }
 
 Buffer
@@ -186,6 +137,8 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     return Buffer{new CudaBuffer{nullptr, 0, -1}};
   }
 
+  // Find available buffer from cache.
+  std::unique_lock lock(mutex_);
   if (size <= small_block_size) {
     size = 8;
   } else if (size < page_size) {
@@ -198,8 +151,6 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     device = -1;
   }
 
-  // Find available buffer from cache.
-  std::unique_lock lock(mutex_);
   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.
@@ -217,14 +168,9 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     if (!buf) {
       void* data = nullptr;
       if (device == -1) {
-        data = unified_malloc(size);
+        CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
       } else {
-        cu::device(device).make_current();
-        if (mem_pools_[device]) { // supports memory pools
-          CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
-        } else {
-          CHECK_CUDA_ERROR(cudaMalloc(&data, size));
-        }
+        CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
       }
       if (!data) {
         std::ostringstream msg;
@@ -240,14 +186,12 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     // from OOM
     if (get_cache_memory() > 0) {
       for (auto p : mem_pools_) {
-        if (p) {
-          size_t used = 0;
-          CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
-              p, cudaMemPoolAttrReservedMemCurrent, &used));
-          if (used > (total_memory_ - free_limit_)) {
-            buffer_cache_.release_cached_buffers(free_limit_);
-            break;
-          }
+        size_t used = 0;
+        CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
+            p, cudaMemPoolAttrReservedMemCurrent, &used));
+        if (used > (total_memory_ - free_limit_)) {
+          buffer_cache_.release_cached_buffers(free_limit_);
+          break;
         }
       }
     }
@@ -259,10 +203,9 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
   if (get_cache_memory() > max_pool_size_) {
     buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
   }
-  lock.unlock();
-  // Copy to unified memory here if the buffer is not on the right device.
+  // Copy to managed here if the buffer is not on the right device
   if (buf->device >= 0 && buf->device != device) {
-    move_to_unified_memory(*buf, stream);
+    copy_to_managed(*buf);
   }
   return Buffer{buf};
 }
@@ -286,7 +229,7 @@ void CudaAllocator::free(Buffer buffer) {
   if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
-    free_cuda_buffer(buf);
+    cuda_free(buf);
   }
 }
 
@@ -298,49 +241,17 @@ size_t CudaAllocator::size(Buffer buffer) const {
   return buf->size;
 }
 
-void CudaAllocator::move_to_unified_memory(
-    CudaBuffer& buf,
-    cudaStream_t stream) {
-  if (buf.device == -1) {
-    return;
-  }
-  void* data = unified_malloc(buf.size);
-  cudaMemcpyKind kind =
-      supports_managed_memory() ? cudaMemcpyDefault : cudaMemcpyDeviceToHost;
-  if (stream && mem_pools_[buf.device]) {
-    CHECK_CUDA_ERROR(cudaMemcpyAsync(data, buf.data, buf.size, kind, stream));
-    free_async(buf, stream);
-  } else {
-    CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
-    free_async(buf);
-  }
-  buf.data = data;
-  buf.device = -1;
-}
-
 // This must be called with mutex_ aquired
-void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) {
+void CudaAllocator::cuda_free(CudaBuffer* buf) {
   if (scalar_pool_.in_pool(buf)) {
     scalar_pool_.free(buf);
   } else {
-    free_async(*buf);
-    delete buf;
-  }
-}
-
-void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) {
-  if (buf.device == -1) {
-    unified_free(buf.data);
-  } else {
-    // Free asynchronously when memory pools is supported.
-    if (mem_pools_[buf.device]) {
-      if (!stream) {
-        stream = free_streams_[buf.device];
-      }
-      CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream));
+    if (buf->device >= 0) {
+      CHECK_CUDA_ERROR(cudaFreeAsync(buf->data, free_streams_[buf->device]));
     } else {
-      CHECK_CUDA_ERROR(cudaFree(buf.data));
+      CHECK_CUDA_ERROR(cudaFree(buf->data));
     }
+    delete buf;
   }
 }
 
@@ -383,20 +294,22 @@ void CudaAllocator::clear_cache() {
 }
 
 CudaAllocator& allocator() {
-  static auto* allocator_ = []() {
-    // Ensure scheduler is created before allocator.
-    scheduler::scheduler();
-    // By creating the |allocator_| on heap, the destructor of CudaAllocator
-    // will not be called on exit and buffers in the cache will be leaked. This
-    // can save some time at program exit.
-    return new CudaAllocator();
-  }();
+  // By creating the |allocator_| on heap, the destructor of CudaAllocator
+  // will not be called on exit and buffers in the cache will be leaked. This
+  // can save some time at program exit.
+  static CudaAllocator* allocator_ = new CudaAllocator;
   return *allocator_;
 }
 
 Buffer malloc_async(size_t size, CommandEncoder& encoder) {
-  return allocator().malloc_async(
+  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
@@ -412,7 +325,9 @@ void* Buffer::raw_ptr() {
     return nullptr;
   }
   auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
-  cu::allocator().move_to_unified_memory(cbuf);
+  if (cbuf.device != -1) {
+    copy_to_managed(cbuf);
+  }
   return cbuf.data;
 }
 

mlx/backend/cuda/allocator.h

@@ -54,10 +54,6 @@ class CudaAllocator : public allocator::Allocator {
   void free(Buffer buffer) override;
   size_t size(Buffer buffer) const override;
 
-  // Replace the memory of |buf| with unified memory (managed memory or pinned
-  // host memory), and copy the data over. Pass |stream| to copy asynchronously.
-  void move_to_unified_memory(CudaBuffer& buf, cudaStream_t stream = nullptr);
-
   size_t get_active_memory() const;
   size_t get_peak_memory() const;
   void reset_peak_memory();
@@ -68,8 +64,7 @@ class CudaAllocator : public allocator::Allocator {
   void clear_cache();
 
  private:
-  void free_cuda_buffer(CudaBuffer* buf);
-  void free_async(CudaBuffer& buf, cudaStream_t stream = nullptr);
+  void cuda_free(CudaBuffer* buf);
 
   CudaAllocator();
   friend CudaAllocator& allocator();
@@ -82,7 +77,7 @@ class CudaAllocator : public allocator::Allocator {
   BufferCache<CudaBuffer> buffer_cache_;
   size_t active_memory_{0};
   size_t peak_memory_{0};
-  std::vector<CudaStream> free_streams_;
+  std::vector<cudaStream_t> free_streams_;
   std::vector<cudaMemPool_t> mem_pools_;
   SmallSizePool scalar_pool_;
 };

mlx/backend/cuda/arange.cu

@@ -56,6 +56,7 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
           cu::arange<OutType, IdxT, N_WRITES>,
           num_blocks,
           block_dims,
+          0,
           gpu_ptr<OutType>(out),
           out.data_size(),
           static_cast<CTYPE>(start_),

mlx/backend/cuda/arg_reduce.cu

@@ -172,6 +172,7 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
           kernel,
           num_blocks,
           block_dim(),
+          0,
           gpu_ptr<T>(in),
           gpu_ptr<uint32_t>(out),
           out.size(),

mlx/backend/cuda/binary/binary.cuh

@@ -16,14 +16,8 @@ namespace cu {
 
 namespace cg = cooperative_groups;
 
-constexpr int BINARY_MAX_BLOCK_DIM = 1024;
-
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_ss(
-    const In* a,
-    const In* b,
-    Out* out,
-    IdxT size) {
+__global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -42,11 +36,7 @@ __global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_ss(
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_sv(
-    const In* a,
-    const In* b,
-    Out* out,
-    IdxT size) {
+__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -67,11 +57,7 @@ __global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_sv(
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vs(
-    const In* a,
-    const In* b,
-    Out* out,
-    IdxT size) {
+__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -92,11 +78,7 @@ __global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vs(
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vv(
-    const In* a,
-    const In* b,
-    Out* out,
-    IdxT size) {
+__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -309,6 +291,7 @@ void binary_op_gpu_inplace(
                         kernel,
                         {num_blocks_x, num_blocks_y},
                         block_dims,
+                        0,
                         gpu_ptr<InType>(a),
                         gpu_ptr<InType>(b),
                         gpu_ptr<OutType>(out),
@@ -326,6 +309,7 @@ void binary_op_gpu_inplace(
                       kernel,
                       {num_blocks_x, num_blocks_y},
                       block_dims,
+                      0,
                       gpu_ptr<InType>(a),
                       gpu_ptr<InType>(b),
                       gpu_ptr<OutType>(out),
@@ -349,16 +333,12 @@ void binary_op_gpu_inplace(
               kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
             }
             auto [num_blocks, block_dims] = get_launch_args(
-                out.data_size(),
-                out.shape(),
-                out.strides(),
-                large(),
-                N_READS,
-                cu::BINARY_MAX_BLOCK_DIM);
+                out.data_size(), out.shape(), out.strides(), large(), N_READS);
             encoder.add_kernel_node(
                 kernel,
                 num_blocks,
                 block_dims,
+                0,
                 gpu_ptr<InType>(a),
                 gpu_ptr<InType>(b),
                 gpu_ptr<OutType>(out),
@@ -366,12 +346,11 @@ void binary_op_gpu_inplace(
           });
         }
       } else {
-        throw std::runtime_error(
-            fmt::format(
-                "Can not do binary op {} on inputs of {} with result of {}.",
-                op,
-                dtype_to_string(a.dtype()),
-                dtype_to_string(out.dtype())));
+        throw std::runtime_error(fmt::format(
+            "Can not do binary op {} on inputs of {} with result of {}.",
+            op,
+            dtype_to_string(a.dtype()),
+            dtype_to_string(out.dtype())));
       }
     });
   });

mlx/backend/cuda/binary_two.cu

@@ -314,6 +314,7 @@ void binary_two_op_gpu_inplace(
                         kernel,
                         {num_blocks_x, num_blocks_y},
                         block_dims,
+                        0,
                         gpu_ptr<InType>(a),
                         gpu_ptr<InType>(b),
                         gpu_ptr<OutType>(out_a),
@@ -332,6 +333,7 @@ void binary_two_op_gpu_inplace(
                       kernel,
                       {num_blocks_x, num_blocks_y},
                       block_dims,
+                      0,
                       gpu_ptr<InType>(a),
                       gpu_ptr<InType>(b),
                       gpu_ptr<OutType>(out_a),
@@ -365,6 +367,7 @@ void binary_two_op_gpu_inplace(
                 kernel,
                 num_blocks,
                 block_dims,
+                0,
                 gpu_ptr<InType>(a),
                 gpu_ptr<InType>(b),
                 gpu_ptr<OutType>(out_a),
@@ -373,12 +376,11 @@ void binary_two_op_gpu_inplace(
           });
         }
       } else {
-        throw std::runtime_error(
-            fmt::format(
-                "Can not do binary op {} on inputs of {} with result of {}.",
-                op,
-                dtype_to_string(a.dtype()),
-                dtype_to_string(out_a.dtype())));
+        throw std::runtime_error(fmt::format(
+            "Can not do binary op {} on inputs of {} with result of {}.",
+            op,
+            dtype_to_string(a.dtype()),
+            dtype_to_string(out_a.dtype())));
       }
     });
   });

mlx/backend/cuda/compiled.cpp

@@ -36,16 +36,14 @@ struct FusedKernelBuilder {
       params.push_back(
           fmt::format("const {}* {}", dtype_to_cuda_type(x.dtype()), xname));
       if (!is_scalar(x) && !contiguous) {
-        params.push_back(
-            fmt::format(
-                "const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
-                xname));
+        params.push_back(fmt::format(
+            "const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
+            xname));
       }
     }
     for (const auto& x : outputs) {
-      params.push_back(
-          fmt::format(
-              "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
+      params.push_back(fmt::format(
+          "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
     }
     if (!contiguous) {
       params.push_back(
@@ -252,30 +250,20 @@ void Compiled::eval_gpu(
     builder.os += "\n} // namespace mlx::core::cu\n";
     // Build kernel names.
     std::vector<std::string> kernel_names;
-    kernel_names.push_back(
-        fmt::format(
-            "mlx::core::cu::{}_contiguous<uint32_t, {}>",
-            lib_name(),
-            work_per_thread));
-    kernel_names.push_back(
-        fmt::format(
-            "mlx::core::cu::{}_contiguous<int64_t, {}>",
-            lib_name(),
-            work_per_thread));
-    for (int wpt : {1, work_per_thread}) {
+    kernel_names.push_back(fmt::format(
+        "mlx::core::cu::{}_contiguous<uint32_t, {}>",
+        lib_name(),
+        work_per_thread));
+    kernel_names.push_back(fmt::format(
+        "mlx::core::cu::{}_contiguous<int64_t, {}>",
+        lib_name(),
+        work_per_thread));
+    for (auto wpt : std::array<int, 2>{1, work_per_thread}) {
       for (int i = 1; i <= MAX_NDIM; ++i) {
-        kernel_names.push_back(
-            fmt::format(
-                "mlx::core::cu::{}_strided<{}, uint32_t, {}>",
-                lib_name(),
-                i,
-                wpt));
-        kernel_names.push_back(
-            fmt::format(
-                "mlx::core::cu::{}_strided<{}, int64_t, {}>",
-                lib_name(),
-                i,
-                wpt));
+        kernel_names.push_back(fmt::format(
+            "mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt));
+        kernel_names.push_back(fmt::format(
+            "mlx::core::cu::{}_strided<{}, int64_t, {}>", lib_name(), i, wpt));
       }
     }
 
@@ -351,8 +339,7 @@ void Compiled::eval_gpu(
   auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name);
   auto [num_blocks, block_dims] =
       get_launch_args(outputs[0], large, work_per_thread, max_block_dims);
-  encoder.add_kernel_node_raw(
-      kernel, num_blocks, block_dims, {}, 0, args.args());
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
 }
 
 } // namespace mlx::core

mlx/backend/cuda/conv.cpp

@@ -39,7 +39,7 @@ struct ConvCacheKey {
 };
 
 auto& conv_cache() {
-  static thread_local LRUBytesKeyCache<
+  static LRUBytesKeyCache<
       ConvCacheKey,
       std::pair<ConvBackendType, std::optional<DnnGraph>>>
       cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
@@ -103,7 +103,7 @@ std::optional<DnnGraph> build_conv_graph(
     const std::vector<int64_t>& dilation) {
   auto compute_dtype =
       (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
-  DnnGraph graph(get_cudnn_handle(encoder.device()), dtype, compute_dtype);
+  DnnGraph graph(encoder.device().cudnn_handle(), dtype, compute_dtype);
   auto x_ = graph.tensor_nchw("X", 'x', x);
   auto w_ = graph.tensor_nchw("W", 'w', w);
 
@@ -139,7 +139,7 @@ std::optional<DnnGraph> build_conv_graph(
   if (dtype == float32 && !env::enable_tf32()) {
     graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
   }
-  CHECK_CUDNN_ERROR(graph.build());
+  CHECK_CUDNN_FE_ERROR(graph.build());
   return graph;
 }
 
@@ -252,10 +252,6 @@ void register_args(
 
 } // namespace
 
-void init_cudnn_conv_cache() {
-  conv_cache();
-}
-
 void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
   nvtx3::scoped_range r("Convolution::eval_gpu");
   if (out_.size() == 0) {
@@ -273,19 +269,20 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
 
   // Search cache.
   BytesKey<ConvCacheKey> cache_key;
-  cache_key.pod.device_id = encoder.device().cuda_device();
-  cache_key.pod.cudnn_dtype = dtype_to_cudnn_type(dtype);
-  cache_key.pod.input_shape = vector_key(in.shape());
-  cache_key.pod.weight_shape = vector_key(wt.shape());
-  cache_key.pod.stride = vector_key(kernel_strides_);
-  cache_key.pod.padding_lo = vector_key(padding_lo_);
-  cache_key.pod.padding_hi = vector_key(padding_hi_);
-  cache_key.pod.dilation = vector_key(kernel_dilation_);
-  cache_key.pod.groups = groups_;
-  cache_key.pod.flip = flip_;
-  cache_key.pod.input_alignment = get_alignment(in);
-  cache_key.pod.weight_alignment = get_alignment(wt);
-  cache_key.pod.output_alignment = get_alignment(out);
+  cache_key.pod = {
+      encoder.device().cuda_device(),
+      dtype_to_cudnn_type(dtype),
+      vector_key(in.shape()),
+      vector_key(wt.shape()),
+      vector_key(kernel_strides_),
+      vector_key(padding_lo_),
+      vector_key(padding_hi_),
+      vector_key(kernel_dilation_),
+      groups_,
+      flip_,
+      get_alignment(in),
+      get_alignment(wt),
+      get_alignment(out)};
   if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
     auto& [backend_type, graph] = it->second;
     if (graph) {
@@ -293,7 +290,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
       std::tie(in, wt, out) =
           prepare_args(encoder, backend_type, in, wt, out, groups_, s);
       register_args(encoder, backend_type, in, wt, out, out_);
-      CHECK_CUDNN_ERROR(graph->encode_capturing(
+      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
           encoder,
           {
               {'x', gpu_ptr<void>(in)},
@@ -375,7 +372,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
 
   if (graph) {
     register_args(encoder, backend_type, in, wt, out, out_);
-    CHECK_CUDNN_ERROR(graph->encode_capturing(
+    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
         encoder,
         {
             {'x', gpu_ptr<void>(in)},

mlx/backend/cuda/conv/gemm_conv.cu

@@ -117,6 +117,7 @@ array unfold_inputs_nd(
         cu::naive_unfold_nd<DataType, NDIM>,
         num_blocks,
         block_dims,
+        0,
         gpu_ptr<DataType>(in),
         gpu_ptr<DataType>(unfolded),
         filter_size,

mlx/backend/cuda/conv/gemm_grouped_conv.cu

@@ -120,6 +120,7 @@ array grouped_unfold_transpose_inputs_nd(
         cu::naive_grouped_unfold_transpose_nd<DataType, NDIM>,
         num_blocks,
         block_dims,
+        0,
         gpu_ptr<DataType>(in),
         gpu_ptr<DataType>(unfolded),
         filter_size,

mlx/backend/cuda/copy/copy_contiguous.cu

@@ -76,6 +76,7 @@ void copy_contiguous(
             kernel,
             num_blocks,
             block_dims,
+            0,
             gpu_ptr<InType>(in) + in_offset,
             gpu_ptr<OutType>(out) + out_offset,
             out.data_size());

mlx/backend/cuda/copy/copy_general.cu

@@ -137,6 +137,7 @@ void copy_general(
                     kernel,
                     {num_blocks_x, num_blocks_y},
                     block_dims,
+                    0,
                     in_ptr,
                     out_ptr,
                     rest,
@@ -153,6 +154,7 @@ void copy_general(
                   kernel,
                   {num_blocks_x, num_blocks_y},
                   block_dims,
+                  0,
                   in_ptr,
                   out_ptr,
                   rest,

mlx/backend/cuda/copy/copy_general_dynamic.cu

@@ -83,6 +83,7 @@ void copy_general_dynamic(
                         dims_constant()>,
                     num_blocks,
                     block_dims,
+                    0,
                     in_ptr,
                     out_ptr,
                     out.size(),
@@ -98,6 +99,7 @@ void copy_general_dynamic(
                   cu::copy_gg_dynamic<InType, OutType, IdxT>,
                   num_blocks,
                   block_dims,
+                  0,
                   in_ptr,
                   out_ptr,
                   out.size(),