Use uv in macOS CI

Co-authored-by: Kanishk <kanishk.chores@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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

@@ -18,7 +18,14 @@ runs:
       env:
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
       run: |
-        pip install auditwheel build patchelf setuptools
+        pip install auditwheel "build<=1.4.2" patchelf setuptools
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
-        bash python/scripts/repair_cuda.sh ${{ inputs.arch }}
+
+        auditwheel repair dist/mlx_cuda*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libcublas* \
+          --exclude libcuda* \
+          --exclude libcudnn* \
+          --exclude libnccl* \
+          --exclude libnvrtc*

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

@@ -18,19 +18,21 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Generate package stubs
+    - name: Build MLX
       shell: bash
-      run: |
-        pip install -e ".[dev]" -v
-        pip install typing_extensions
-        python setup.py generate_stubs
+      run: pip install -e . -v
+
     - name: Build Python package
       shell: bash
       run: |
-        pip install auditwheel patchelf build
+        pip install auditwheel patchelf "build<=1.4.2"
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
-        bash python/scripts/repair_linux.sh ${{ inputs.arch }}
+        auditwheel repair dist/mlx-*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libmlx.so* \
+          --only-plat
+
     - name: Build backend package
       if: ${{ inputs.build-backend }}
       shell: bash

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

@@ -9,6 +9,7 @@ inputs:
 runs:
   using: "composite"
   steps:
+
     - name: Install Python package
       id: python_build
       shell: sh
@@ -20,20 +21,16 @@ 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=90a"
-          # Can not build tests when the built executables can not run.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
+          # 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
+        # Install cpu-only torch to save space
+        pip install torch --index-url https://download.pytorch.org/whl/cpu
         pip install --no-build-isolation -e ".[dev]" -v
         # Pass the CMAKE_ARGS to following steps.
         echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
 
-    - name: Generate package stubs
-      shell: sh
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Build CPP only
       shell: bash
       run: |

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

@@ -18,9 +18,10 @@ 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
+        uv pip install build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
 
@@ -28,6 +29,7 @@ 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

@@ -4,65 +4,72 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Install Python package
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools nanobind==2.10.2
-        pip install -e . -v
-
-    - name: Generate package stubs
-      shell: bash -l {0}
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
 
     - name: Install tests dependencies
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install numpy torch tensorflow unittest-xml-reporting
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
 
     - name: Run Python tests
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
       run: |
-        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
+        echo "::group::Run Python tests"
+        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
         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
-    
+        echo "::endgroup::"
+
     - name: Build example extension
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build example extension"
         cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
     - name: Build CPP only
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build CPP only"
         mkdir -p build
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
     - name: Run CPP tests
-      shell: bash -l {0}
+      shell: bash
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
-    - name: Build small binary with JIT
-      shell: bash -l {0}
       run: |
+        echo "::group::Run CPP tests"
+        ./build/tests/tests
+        ./build/tests/test_teardown
+        echo "::endgroup::"
+
+    - name: Build small binary with JIT
+      shell: bash
+      run: |
+        echo "::group::Build small binary with JIT"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -72,17 +79,18 @@ runs:
           -DMLX_BUILD_GGUF=OFF \
           -DMLX_METAL_JIT=ON
         make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
     - name: Run Python tests with JIT
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run Python tests with JIT"
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
-        python -m xmlrunner discover \
-            -v python/tests \
-            -o test-results/gpu_jit
+          uv pip install -e . -v
+        python -m unittest discover -v python/tests
+        echo "::endgroup::"

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

@@ -0,0 +1,26 @@
+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

@@ -9,11 +9,14 @@ inputs:
   python-version:
     description: 'Version of python to set up'
     required: false
-    default: '3.10'
+    default: '3.14'
   use-ccache:
     description: 'Whether to enable ccache'
     required: false
     default: 'true'
+  ccache-key:
+    required: false
+    default: 'ccache'
 
 runs:
   using: "composite"
@@ -21,14 +24,19 @@ runs:
     - name: Install common dependencies
       shell: bash
       run: |
+        echo "::group::Install common dependencies"
         sudo apt-get update
-        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip
+        sudo apt-get install -y --no-install-recommends \
+            zip \
+            libblas-dev liblapack-dev liblapacke-dev \
+            openmpi-bin openmpi-common libopenmpi-dev
+        echo "::endgroup::"
 
     - name: Use ccache
       if: ${{ inputs.use-ccache == 'true' }}
       uses: hendrikmuhs/ccache-action@v1.2
       with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
+        key: ${{ inputs.ccache-key }}-${{ 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
@@ -40,16 +48,20 @@ runs:
     - name: Setup Python venv
       shell: bash
       run: |
+        echo "::group::Setup Python venv"
         python -m venv .venv
         source .venv/bin/activate
-        pip install setuptools cmake nanobind==2.10.2
+        pip install setuptools cmake typing_extensions
         echo PATH=$PATH >> $GITHUB_ENV
-        # Make cmake search .venv for nanobind
+        # Search python packages in .venv
         echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"
 
-    - name: Install MPI
-      shell: bash
-      run: sudo apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
+    - 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') }}
@@ -60,34 +72,31 @@ runs:
         # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
         PACKAGES: |
           {
-            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-toolkit-12-6",
-            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-toolkit-12-9",
-            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
+            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6",
+            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9",
+            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0"
           }
       run: |
+        echo "::group::Install CUDA toolkit"
         # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is
         # Jetson specific. SBSA means Arm Server Base System Architecture.
         ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
         wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
         sudo dpkg -i cuda-keyring_1.1-1_all.deb
         sudo apt-get update
-        sudo apt-get install -y \
+        sudo apt-get install -y --no-install-recommends \
             libnccl2 libnccl-dev \
             ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
         echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH
+        echo "::endgroup::"
 
     - name: CUDA packages and driver report
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       shell: bash
       run: |
-        sudo apt-get install -y ubuntu-drivers-common dkms
-        echo "NVIDIA Driver Packages Available:"
-        sudo ubuntu-drivers list --gpgpu
-        echo "NVIDIA Driver Version:"
-        cat /proc/driver/nvidia/version || echo "nvidia driver not found"
-        echo "Installed NVIDIA and CUDA packages:"
+        echo "::group::Installed NVIDIA and CUDA packages"
         dpkg -l | egrep "cuda|nvidia" -i
-        echo "DKMS Status:"
-        dkms status || echo "dkms not found"
-        echo "NVIDIA-SMI Status:"
-        nvidia-smi || echo "nvidia-smi not found"
+        echo "::endgroup::"
+        echo "::group::NVIDIA-SMI Status"
+        nvidia-smi || true
+        echo "::endgroup::"

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

@@ -13,12 +13,20 @@ runs:
     - name: Install Homebrew packages
       shell: sh
       run: /opt/homebrew/bin/brew install openmpi
-    
+
     - name: Verify MetalToolchain installed
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        uv venv --python ${{ inputs.python-version }}
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"

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

@@ -0,0 +1,42 @@
+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="*fft_tests.cpp,*linalg_tests.cpp"
+        ./build/tests/tests -sfe="*linalg_tests.cpp"
         echo "::endgroup::"

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

@@ -0,0 +1,21 @@
+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/scripts/build-sanitizer-tests.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+set -ex
+
+export CMAKE_C_COMPILER=/usr/bin/clang
+export CMAKE_CXX_COMPILER=/usr/bin/clang++
+BASE_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=DEBUG -DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+if [[ "$(uname -s)" != "Darwin" ]]; then
+  BASE_CMAKE_ARGS+=" -DMLX_BUILD_METAL=OFF"
+fi
+
+run_test() {
+  local sanitizer_name=$1
+  local cmake_sanitizer_flag="-DUSE_${sanitizer_name}=ON"
+  echo "  Running tests with: ${sanitizer_name}"
+
+  case "$sanitizer_name" in
+    ASAN)
+      export ASAN_OPTIONS="detect_leaks=0"
+      ;;
+    UBSAN)
+      export UBSAN_OPTIONS="halt_on_error=0:print_stacktrace=1"
+      ;;
+    TSAN)
+      export TSAN_OPTIONS=""
+      ;;
+  esac
+
+  rm -rf build
+  mkdir -p build
+  pushd build > /dev/null
+
+  cmake .. ${BASE_CMAKE_ARGS} ${cmake_sanitizer_flag}
+  make -j $(nproc)
+  ./tests/tests
+
+  popd > /dev/null
+  unset ${sanitizer_name}_OPTIONS
+}
+
+sanitizer_arg=$(echo "$1" | tr '[:lower:]' '[:upper:]')
+
+if [[ "$sanitizer_arg" == "ASAN" || "$sanitizer_arg" == "UBSAN" || "$sanitizer_arg" == "TSAN" ]]; then
+  run_test "$sanitizer_arg"
+  echo "  ${sanitizer_arg} test run completed successfully."
+else
+  echo "Error: Invalid sanitizer '$1'. Please use one of: ASAN, UBSAN, TSAN."
+  exit 1
+fi

.github/workflows/build_and_test.yml

@@ -36,6 +36,7 @@ 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 }})
@@ -65,7 +66,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        macos-target: ["14.0", "15.0"]
+        macos-target: ["14.0", "15.0", "26.0"]
     runs-on: [self-hosted, macos]
     env:
       MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }}
@@ -75,6 +76,16 @@ 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'
@@ -84,6 +95,39 @@ jobs:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/build-docs
 
+  linux_sanitizer_build_and_test:
+    name: Linux Sanitizer Tests (${{ matrix.sanitizer }})
+    needs: check_lint
+    strategy:
+      fail-fast: false
+      matrix:
+        sanitizer: [ASAN, UBSAN]
+        # todo 12/16/2025: enable TSAN later + consider enabling ASAN for GPU backend tests.
+        # sanitizer: [ASAN, UBSAN, TSAN]
+    runs-on: ubuntu-22.04-arm
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+
+      - name: Install Dependencies
+        run: |
+          export DEBIAN_FRONTEND=noninteractive
+          sudo apt-get update -y
+          sudo apt-get install -y \
+            build-essential \
+            libblas-dev \
+            liblapacke-dev \
+            libopenblas-dev \
+            cmake \
+            clang \
+            git
+          sudo apt-get clean
+          sudo rm -rf /var/lib/apt/lists/*
+
+      - name: Linux Build and Test with ${{ matrix.sanitizer }}
+        run: |
+          bash .github/scripts/build-sanitizer-tests.sh ${{ matrix.sanitizer }}
+
   linux_fedora_build_cpp:
     name: Linux Fedora (${{ matrix.arch }})
     needs: check_lint

.github/workflows/documentation.yml

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

.github/workflows/nightly.yml

@@ -23,18 +23,19 @@ jobs:
           build-backend: ${{ matrix.python-version == '3.10' }}
           arch: "x86_64"
       - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         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@v6
+        uses: actions/upload-artifact@v7
         with:
           name: mlx-cpu
           path: wheelhouse/mlx_cpu-*.whl
           retention-days: 7
+      - run: df -h
 
   build_linux_with_tests:
     strategy:
@@ -52,6 +53,7 @@ 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'
@@ -65,6 +67,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - uses: ./.github/actions/build-macos
+      - name: Build macOS 26 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 26.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Build macOS 15 package
         uses: ./.github/actions/build-macos-release
         with:
@@ -78,19 +85,24 @@ jobs:
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
+    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' }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
-          toolkit: 'cuda-12.9'
+          toolkit: ${{ matrix.toolkit }}
+          ccache-key: 'ccache-release'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          toolkit: 'cuda-12.9'
+          arch: ${{ matrix.arch }}
       - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          path: wheelhouse/mlx_cuda_*.whl
           retention-days: 7

.github/workflows/release.yml

@@ -4,31 +4,32 @@ on:
   push:
     tags:
       - 'v*'
+    branches:
+      - 'test-publish/*'
   workflow_dispatch:
     inputs:
+      dry_run:
+        description: 'Dry run (do not publish to PyPi)'
+        required: false
+        type: boolean
       dev_release:
-        description: "Do a dev release or regular release"
-        required: true
-        default: "false"
+        description: 'Development release (DEV_RELEASE=1)'
+        required: false
+        type: boolean
 
 permissions:
   contents: read
 
 jobs:
-  setup:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Set publishing variables
-        run: echo "Publishing setup complete"
-
   build_documentation:
     if: github.repository == 'ml-explore/mlx'
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/build-docs
-    
+
   deploy_documentation:
+    if: ${{ !inputs.dry_run }}
     needs: build_documentation
     permissions:
       pages: write
@@ -40,7 +41,7 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5
 
   build_linux_release:
     if: github.repository == 'ml-explore/mlx'
@@ -51,7 +52,7 @@ jobs:
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
@@ -60,22 +61,24 @@ jobs:
           use-ccache: false
       - uses: ./.github/actions/build-linux-release
         with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
+          build-backend: ${{ matrix.python_version == '3.10' }}
           arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
           name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
           path: wheelhouse/mlx-*.whl
+          if-no-files-found: error
       - name: Upload CPU artifacts
         if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
           name: mlx-cpu-${{ matrix.arch }}
           path: wheelhouse/mlx_cpu-*.whl
-  
+          if-no-files-found: error
+
   build_mac_release:
     if: github.repository == 'ml-explore/mlx'
     strategy:
@@ -84,25 +87,14 @@ jobs:
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools nanobind==2.10.2
-          pip install -e . -v
-      - name: Generate package stubs
-        shell: bash -l {0}
-        run: |
-          pip install typing_extensions
-          python setup.py generate_stubs
+      - name: Install Python package
+        run: uv pip install -e . -v
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:
@@ -113,19 +105,26 @@ 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@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
           name: mac-wheels-${{ matrix.python-version }}
           path: dist/mlx-*.whl
+          if-no-files-found: error
       - name: Upload Metal artifacts
         if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
           name: mlx-metal
           path: dist/mlx_metal-*.whl
+          if-no-files-found: error
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
@@ -136,68 +135,72 @@ jobs:
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-          use-ccache: false
+          ccache-key: 'ccache-release'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
           arch: ${{ matrix.arch }}
       - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          path: wheelhouse/mlx_cuda_*.whl
+          if-no-files-found: error
 
   pypi-publish:
     name: Upload release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_linux_release, build_mac_release]
+    needs: [build_linux_release, build_mac_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
       url: https://pypi.org/p/mlx
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           pattern: linux-wheels-*
           merge-multiple: true
           path: dist
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           pattern: mac-wheels-*
           merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: ${{ !inputs.dry_run }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
-  
+
   pypi-publish-cuda:
     name: Upload CUDA release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_cuda_release]
+    needs: [build_cuda_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
       url: https://pypi.org/p/mlx-cuda
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
-          name: mlx-cuda
+          pattern: mlx-cuda-*
+          merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: ${{ !inputs.dry_run }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -205,21 +208,22 @@ jobs:
   pypi-publish-cpu:
     name: Upload CPU release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_linux_release]
+    needs: [build_linux_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
       url: https://pypi.org/p/mlx-cpu
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           pattern: mlx-cpu-*
           merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: ${{ !inputs.dry_run }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -227,20 +231,21 @@ jobs:
   pypi-publish-metal:
     name: Upload Metal release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_mac_release]
+    needs: [build_mac_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
       url: https://pypi.org/p/mlx-metal
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           name: mlx-metal
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: ${{ !inputs.dry_run }}
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/

.gitignore

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

.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: v19.1.7
+    rev: v21.1.8
     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: 25.1.0
+    rev: 26.1.0
     hooks:
     -   id: black
     
 -   repo: https://github.com/pycqa/isort
-    rev: 6.0.0
+    rev: 7.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 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
@@ -40,11 +40,14 @@ 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)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF)
+option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF)
+option(USE_UBSAN "Enable UndefinedBehaviorSanitizer (UBSan)" OFF)
+option(USE_TSAN "Enable ThreadSanitizer (TSan)" OFF)
 
 # --------------------- Processor tests -------------------------
 message(
@@ -81,6 +84,63 @@ if(MLX_USE_CCACHE)
   endif()
 endif()
 
+if(USE_ASAN AND USE_TSAN)
+  message(
+    FATAL_ERROR
+      "AddressSanitizer (ASan) and ThreadSanitizer (TSan) are mutually exclusive and cannot be enabled at the same time."
+  )
+endif()
+
+set(SANITIZER_COMPILE_FLAGS "")
+set(SANITIZER_LINK_FLAGS "")
+
+if(USE_ASAN)
+  if(WIN32 AND MSVC)
+    list(APPEND SANITIZER_COMPILE_FLAGS /fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS /fsanitize=address)
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=address)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
+if(USE_UBSAN)
+  if(WIN32 AND MSVC)
+    if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+      list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+      list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+    else()
+      message(
+        WARNING
+          "UndefinedBehaviorSanitizer (UBSan) is not directly supported via a simple flag in MSVC."
+      )
+    endif()
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+  endif()
+endif()
+
+if(USE_TSAN)
+  if(WIN32 AND MSVC)
+    message(
+      FATAL_ERROR
+        "ThreadSanitizer (TSan) is not supported by the MSVC compiler. Please use Clang or GCC."
+    )
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
+    message(FATAL_ERROR "ThreadSanitizer (TSan) is not supported on macOS.")
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=thread)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=thread)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
 # ----------------------------- Lib -----------------------------
 
 include(FetchContent)
@@ -89,13 +149,17 @@ 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)
@@ -159,14 +223,17 @@ if(WIN32)
   if(MSVC)
     # GGUF does not build with MSVC.
     set(MLX_BUILD_GGUF OFF)
-    # There is no prebuilt OpenBLAS distribution for MSVC.
-    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
+  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})
   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.1
+    GIT_TAG v1.4.2
     EXCLUDE_FROM_ALL)
   block()
   set(BUILD_SHARED_LIBS OFF)
@@ -190,20 +257,25 @@ if(MLX_BUILD_CPU)
     target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
     add_compile_definitions(MLX_USE_ACCELERATE)
     add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
-    # Download and build OpenBLAS from source code.
+  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.
     FetchContent_Declare(
       openblas
-      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
+      URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip"
+    )
     FetchContent_MakeAvailable(openblas)
-    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}")
+    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})
   else()
     if(${CMAKE_HOST_APPLE})
       # The blas shipped in macOS SDK is not supported, search homebrew for
@@ -249,22 +321,28 @@ 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 10.2.1
+    GIT_TAG 12.1.0
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(fmt)
 endif()
@@ -276,11 +354,13 @@ if(MLX_BUILD_PYTHON_BINDINGS)
     Python 3.10
     COMPONENTS Interpreter Development.Module
     REQUIRED)
-  execute_process(
-    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE nanobind_ROOT)
-  find_package(nanobind CONFIG REQUIRED)
+  FetchContent_Declare(
+    nanobind
+    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+    GIT_TAG v2.12.0
+    GIT_SHALLOW TRUE
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(nanobind)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 
@@ -300,6 +380,15 @@ 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

@@ -0,0 +1,193 @@
+# 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/comparative/bench_mlx.py

@@ -38,10 +38,10 @@ def bench(f, *args):
     for i in range(10):
         f(*args)
 
-    s = time.time()
+    s = time.perf_counter()
     for i in range(100):
         f(*args)
-    e = time.time()
+    e = time.perf_counter()
     return e - s
 
 

benchmarks/python/comparative/bench_torch.py

@@ -37,10 +37,10 @@ def bench(f, *args):
     for i in range(10):
         f(*args)
 
-    s = time.time()
+    s = time.perf_counter()
     for i in range(100):
         f(*args)
-    e = time.time()
+    e = time.perf_counter()
     return e - s
 
 

benchmarks/python/conv3d_bench.py

@@ -0,0 +1,152 @@
+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

@@ -0,0 +1,119 @@
+# 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,5 +1,6 @@
 import math
 import os
+import platform
 import subprocess
 import time
 from copy import copy
@@ -17,9 +18,6 @@ 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()
@@ -27,11 +25,36 @@ 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(10)]
+VECTOR_LENGTHS = [4096 * (2**i) for i in range(12)]
 MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
 D_TYPES = ("float32", "float16")
 
@@ -202,9 +225,10 @@ def main():
         )
         output_path = os.path.join(
             RESULTS_DIR,
-            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf",
+            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.png",
         )
         fig.savefig(output_path)
+        print(f"Saved benchmark image: {output_path}")
         plt.close(fig)
 
 

benchmarks/python/sdpa_bench.py

@@ -176,6 +176,8 @@ 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 = (
@@ -183,6 +185,8 @@ 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 = (
@@ -190,6 +194,8 @@ 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

@@ -0,0 +1,209 @@
+# 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

@@ -0,0 +1,109 @@
+# 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}"
+            )

benchmarks/python/time_utils.py

@@ -31,8 +31,8 @@ def measure_runtime(fn, **kwargs):
     for _ in range(5):
         fn(**kwargs)
 
-    tic = time.time()
+    tic = time.perf_counter()
     iters = 100
     for _ in range(iters):
         fn(**kwargs)
-    return (time.time() - tic) * 1000 / iters
+    return (time.perf_counter() - tic) * 1000 / iters

cmake/FindCUDNN.cmake

@@ -0,0 +1,177 @@
+# 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,6 +26,7 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
+PREDEFINED             = MLX_API=
 
 ################################################################################
 # Compound extraction control.                                                 #

docs/README.md

@@ -38,3 +38,17 @@ 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

@@ -0,0 +1,36 @@
+<?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

@@ -0,0 +1,36 @@
+<?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 = d.get_command_encoder(s.index);
+        auto& compute_encoder = mx::metal::get_command_encoder(s);
         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:`d.get_command_encoder` to give us the active
+associated. We rely on :meth:`metal::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
@@ -777,11 +777,11 @@ with the naive :meth:`simple_axpby` we first defined.
             mx.eval(z)
 
         # Timed run
-        s = time.time()
+        s = time.perf_counter()
         for i in range(100):
             z = f(x, y, alpha, beta)
             mx.eval(z)
-        e = time.time()
+        e = time.perf_counter()
         return 1000 * (e - s) / 100
 
     simple_time = bench(simple_axpby)

docs/src/dev/metal_logging.rst

@@ -0,0 +1,40 @@
+Metal Logging
+=============
+
+In debug builds, MLX compiles Metal kernels with ``os_log`` enabled so shader
+warnings and debug messages are visible during development.
+
+.. note::
+    Metal logging is only available with Metal 3.2 or higher (macOS 15 and up,
+    iOS 18 and up).
+
+To enable logging from kernels, first make sure to build in debug mode:
+
+.. code-block:: bash
+
+    DEBUG=1 python -m pip install -e .
+
+Then, in the kernel source code include MLX's logging shim and use
+``mlx::os_log``:
+
+.. code-block::
+
+    #include "mlx/backend/metal/kernels/logging.h"
+
+    constant mlx::os_log logger("mlx", "my_kernel");
+
+    kernel void my_kernel(/* ... */) {
+    // ...
+      logger.log_debug("unexpected state: idx=%u", idx);
+    }
+
+When you run the program, set the Metal log level to your desired level and
+forward logs to ``stderr``:
+
+.. code-block:: bash
+
+    MTL_LOG_LEVEL=MTLLogLevelDebug MTL_LOG_TO_STDERR=1 python script.py
+
+See the `Metal logging guide`_ for more details.
+
+.. _`Metal logging guide`: https://developer.apple.com/documentation/metal/logging-shader-debug-messages

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 17)
+  set(CMAKE_CXX_STANDARD 20)
   set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 

docs/src/examples/data_parallelism.rst

@@ -0,0 +1,91 @@
+.. _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

@@ -0,0 +1,239 @@
+.. _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,6 +54,8 @@ are the CPU and GPU.
    examples/linear_regression
    examples/mlp
    examples/llama-inference
+   examples/data_parallelism
+   examples/tensor_parallelism
 
 .. toctree::
    :caption: Python API Reference
@@ -76,6 +78,7 @@ are the CPU and GPU.
    python/optimizers
    python/distributed
    python/tree_utils
+   python/printoptions
 
 .. toctree::
    :caption: C++ API Reference
@@ -89,5 +92,6 @@ are the CPU and GPU.
 
    dev/extensions
    dev/metal_debugger
+   dev/metal_logging
    dev/custom_metal_kernels
    dev/mlx_in_cpp

docs/src/install.rst

@@ -15,7 +15,7 @@ silicon computer is
 
 To install from PyPI your system must meet the following requirements:
 
-- Using an M series chip (Apple silicon)
+- Using `Apple silicon <https://support.apple.com/en-us/116943>`_
 - Using a native Python >= 3.10
 - macOS >= 14.0
 
@@ -83,7 +83,8 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^
 
-- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
+- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux)
+- A C++ compiler with C++20 support (e.g. Clang >= 15.0)
 - `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0
 
@@ -128,13 +129,6 @@ Run the tests with:
 
   python -m unittest discover python/tests
 
-Optional: Install stubs to enable auto completions and type checking from your
-IDE:
-
-.. code-block:: shell
-
-  python setup.py generate_stubs
-
 C++ API
 ^^^^^^^
 

docs/src/python/data_types.rst

@@ -52,7 +52,7 @@ The default floating point type is ``float32`` and the default integer type is
      - 4 
      - 32-bit float
    * - ``float64``
-     - 4 
+     - 8
      - 64-bit double
    * - ``complex64``
      - 8 

docs/src/python/devices_and_streams.rst

@@ -14,6 +14,10 @@ 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,5 +20,7 @@ FFT
   irfft2
   rfftn
   irfftn
+  fftfreq
+  rfftfreq
   fftshift
   ifftshift

docs/src/python/linalg.rst

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

docs/src/python/nn.rst

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

docs/src/python/nn/distributed.rst

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

docs/src/python/printoptions.rst

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

docs/src/python/transforms.rst

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

docs/src/usage/distributed.rst

@@ -117,89 +117,11 @@ 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!
 
-.. _training_example:
+Distributed Program Examples
+----------------------------
 
-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())
+- :ref:`Data Parallelism <data_parallelism>`
+- :ref:`Tensor Parallelism <tensor_parallelism>`
 
 .. _ring_section:
 

docs/src/usage/export.rst

@@ -155,6 +155,34 @@ 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,10 +90,7 @@ PyTorch supports the buffer protocol, but it requires an explicit
 
   a = mx.arange(3)
   b = torch.tensor(memoryview(a))
-  c = mx.array(b.numpy())
-
-Conversion from PyTorch tensors back to arrays must be done via intermediate
-NumPy arrays with ``numpy()``.
+  c = mx.array(b)
 
 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 = d.get_command_encoder(s.index);
+  auto& compute_encoder = mx::metal::get_command_encoder(s);
   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.10.2",
+  "nanobind==2.12.0",
 ]
 build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

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

examples/python/linear_regression.py

@@ -29,12 +29,12 @@ def loss_fn(w):
 
 grad_fn = mx.grad(loss_fn)
 
-tic = time.time()
+tic = time.perf_counter()
 for _ in range(num_iters):
     grad = grad_fn(w)
     w = w - lr * grad
     mx.eval(w)
-toc = time.time()
+toc = time.perf_counter()
 
 loss = loss_fn(w)
 error_norm = mx.sum(mx.square(w - w_star)).item() ** 0.5

examples/python/logistic_regression.py

@@ -30,13 +30,13 @@ def loss_fn(w):
 
 grad_fn = mx.grad(loss_fn)
 
-tic = time.time()
+tic = time.perf_counter()
 for _ in range(num_iters):
     grad = grad_fn(w)
     w = w - lr * grad
     mx.eval(w)
 
-toc = time.time()
+toc = time.perf_counter()
 
 loss = loss_fn(w)
 final_preds = (X @ w) > 0

mlx/CMakeLists.txt

@@ -14,6 +14,7 @@ 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
@@ -22,16 +23,57 @@ 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>)
 
-if(MSVC)
-  # Disable some MSVC warnings to speed up compilation.
-  target_compile_options(mlx PUBLIC /wd4068 /wd4244 /wd4267 /wd4804)
+# 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(WIN32)
-  # Export symbols by default to behave like macOS/linux.
-  set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE)
+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>)
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)

mlx/allocator.h

@@ -4,12 +4,14 @@
 
 #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 Buffer {
+class MLX_API Buffer {
  private:
   void* ptr_;
 
@@ -28,7 +30,7 @@ class Buffer {
   };
 };
 
-class Allocator {
+class MLX_API Allocator {
   /** Abstract base class for a memory allocator. */
  public:
   virtual Buffer malloc(size_t size) = 0;
@@ -47,7 +49,7 @@ class Allocator {
   virtual ~Allocator() = default;
 };
 
-Allocator& allocator();
+MLX_API Allocator& allocator();
 
 inline Buffer malloc(size_t size) {
   return allocator().malloc(size);

mlx/api.h

@@ -0,0 +1,29 @@
+// 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,11 +21,12 @@ 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) {
@@ -69,16 +70,18 @@ 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());
 }
 
@@ -96,7 +99,7 @@ array::array(
     deleter(data);
   } else {
     auto wrapped_deleter = [deleter](allocator::Buffer buffer) {
-      auto ptr = buffer.ptr();
+      auto ptr = buffer.raw_ptr();
       allocator::release(buffer);
       return deleter(ptr);
     };
@@ -131,6 +134,7 @@ 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,6 +8,7 @@
 #include <vector>
 
 #include "mlx/allocator.h"
+#include "mlx/api.h"
 #include "mlx/dtype.h"
 #include "mlx/event.h"
 #include "mlx/small_vector.h"
@@ -22,7 +23,7 @@ using ShapeElem = int32_t;
 using Shape = SmallVector<ShapeElem>;
 using Strides = SmallVector<int64_t>;
 
-class array {
+class MLX_API array {
   /* An array is really a node in a graph. It contains a shared ArrayDesc
    * object */
 
@@ -121,7 +122,7 @@ class array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto shape(int dim) const {
-    return shape().at(dim < 0 ? dim + ndim() : dim);
+    return shape().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
   }
 
   /** The strides of the array. */
@@ -135,7 +136,7 @@ class array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto strides(int dim) const {
-    return strides().at(dim < 0 ? dim + ndim() : dim);
+    return strides().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
   }
 
   /** Get the arrays data type. */
@@ -153,7 +154,7 @@ class array {
   template <typename T>
   T item() const;
 
-  struct ArrayIterator {
+  struct MLX_API ArrayIterator {
     using iterator_category = std::random_access_iterator_tag;
     using difference_type = size_t;
     using value_type = const array;
@@ -464,7 +465,7 @@ class array {
   template <typename It>
   void init(const It src);
 
-  struct ArrayDesc {
+  struct MLX_API ArrayDesc {
     Shape shape;
     Strides strides;
     size_t size;
@@ -488,10 +489,10 @@ class array {
     int64_t offset{0};
 
     // The size in elements of the data buffer the array accesses
-    size_t data_size;
+    size_t data_size{0};
 
     // Contains useful meta data about the array
-    Flags flags;
+    Flags flags{true, true, true};
 
     std::vector<array> inputs;
     // An array to keep track of the siblings from a multi-output
@@ -541,9 +542,10 @@ 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,6 +2,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <cassert>
 #include <functional>
 #include <map>

mlx/backend/common/quantized.h

@@ -0,0 +1,14 @@
+// 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,6 +116,39 @@ 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}/available.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.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

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

mlx/backend/cpu/available.h

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

mlx/backend/cpu/compiled.cpp

@@ -119,13 +119,15 @@ 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()));
     }
   }
 

mlx/backend/cpu/device_info.cpp

@@ -0,0 +1,113 @@
+// 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

@@ -0,0 +1,28 @@
+// 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 CommandEncoder {
+struct MLX_API CommandEncoder {
   CommandEncoder(Stream stream) : stream_(stream) {}
 
   CommandEncoder(const CommandEncoder&) = delete;
@@ -62,6 +62,6 @@ struct CommandEncoder {
   int num_ops_{0};
 };
 
-CommandEncoder& get_command_encoder(Stream stream);
+MLX_API CommandEncoder& get_command_encoder(Stream stream);
 
 } // namespace mlx::core::cpu

mlx/backend/cpu/indexing.cpp

@@ -4,11 +4,14 @@
 #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 {
 
@@ -761,7 +764,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 (uint b = 0; b < batch_count; ++b) {
+  for (size_t b = 0; b < batch_count; ++b) {
     size_t src_consumed = 0;
     src_it.seek(b * src_batch_size);
 
@@ -788,7 +791,7 @@ void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& mask = inputs[1];
   auto& src = inputs[2];
 
-  // Copy src into out (copy allocates memory for out)
+  // Copy dst into out (copy allocates memory for out)
   auto ctype =
       dst.flags().row_contiguous ? CopyType::Vector : CopyType::General;
   copy_cpu(dst, out, ctype, stream());
@@ -851,4 +854,128 @@ 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

@@ -34,18 +34,30 @@ struct VisualStudioInfo {
     arch = "x64";
 #endif
     // Get path of Visual Studio.
-    std::string vs_path = JitCompiler::exec(fmt::format(
-        "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
-        " -property installationPath",
-        std::getenv("ProgramFiles(x86)")));
+    // 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)")));
     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('=');
@@ -140,12 +152,13 @@ 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/luf.cpp

@@ -67,11 +67,10 @@ 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
-               << ((info > 0) ? " because matrix is singular"
-                              : " because argument had an illegal value");
+               << " because argument had an illegal value";
             throw std::runtime_error(ss.str());
           }
 

mlx/backend/cpu/primitives.cpp

@@ -398,44 +398,6 @@ 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,5 +1,6 @@
 // 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"
@@ -14,7 +15,20 @@ namespace mlx::core {
 
 namespace {
 
-const static float MXFP4_LUT[16] = {
+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,
     +1.0f,
@@ -32,24 +46,19 @@ const static float MXFP4_LUT[16] = {
     -4.0f,
     -6.0f};
 
-template <typename T>
+template <typename T, int group_size>
 static inline T dequantize_scale(uint8_t s) {
-  using FOrI = union {
-    bfloat16_t f;
-    uint16_t i;
-  };
-  FOrI out;
-  out.i = (s == 0 ? 0x40 : (static_cast<uint16_t>(s) << 7));
-  return static_cast<T>(out.f);
-}
-
-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);
+  if constexpr (group_size == 16) {
+    return static_cast<T>(detail::FromFP8{}(s));
+  } else {
+    using FOrI = union {
+      bfloat16_t f;
+      uint16_t i;
+    };
+    FOrI out;
+    out.i = (s == 0 ? 0x40 : (static_cast<uint16_t>(s) << 7));
+    return static_cast<T>(out.f);
+  }
 }
 
 template <typename T, int bits>
@@ -437,8 +446,8 @@ void _qmm_dispatch(
   }
 }
 
-template <typename T>
-void mxfp4_qmm(
+template <typename T, int group_size, int bits>
+void fp_qmm(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -446,8 +455,7 @@ void mxfp4_qmm(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = get_pack_factor(4, 8);
+  constexpr int pack_factor = get_pack_factor(bits, 8);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {
@@ -461,25 +469,27 @@ void mxfp4_qmm(
       T xi = *x++;
 
       for (int n = 0; n < N; n += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
         for (int ng = 0; ng < packs_in_group; ng++) {
-          uint8_t wi = *w_local++;
-#pragma clang loop unroll(full)
-          for (int p = 0; p < pack_factor; p++) {
+          if constexpr (bits == 4) {
             (*result_local++) +=
-                xi * scale * static_cast<T>(MXFP4_LUT[wi & 0xf]);
-            wi >>= 4;
+                xi * scale * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            (*result_local++) +=
+                xi * scale * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            (*result_local++) +=
+                xi * scale * static_cast<T>(detail::FromFP8{}(w_local[0]));
           }
+          w_local++;
         }
       }
     }
-
     result += N;
   }
 }
 
-template <typename T>
-void mxfp4_qmm_t(
+template <typename T, int group_size, int bits>
+void fp_qmm_t(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -487,8 +497,7 @@ void mxfp4_qmm_t(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = get_pack_factor(4, 8);
+  constexpr int pack_factor = get_pack_factor(bits, 8);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {
@@ -499,16 +508,19 @@ void mxfp4_qmm_t(
       const T* x_local = x;
       T sum = 0;
       for (int k = 0; k < K; k += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
 
         T gsum = 0;
         for (int kw = 0; kw < packs_in_group; kw++) {
-          uint8_t wi = *w_local++;
-#pragma clang loop unroll(full)
-          for (int p = 0; p < pack_factor; p++) {
-            gsum += (*x_local++) * static_cast<T>(MXFP4_LUT[wi & 0xf]);
-            wi >>= 4;
+          if constexpr (bits == 4) {
+            gsum += (*x_local++) * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            gsum +=
+                (*x_local++) * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            gsum +=
+                (*x_local++) * static_cast<T>(detail::FromFP8{}(w_local[0]));
           }
+          w_local++;
         }
         sum += scale * gsum;
       }
@@ -520,9 +532,9 @@ void mxfp4_qmm_t(
   }
 }
 
-template <int S>
-simd::Simd<float, S> mxfp4_extract_bits_simd(const uint32_t* w) {
-  if constexpr (S == 8) {
+template <int S, int bits>
+simd::Simd<float, S> fp_extract_bits_simd(const uint32_t* w) {
+  if constexpr (S == 8 && bits == 4) {
     constexpr std::array<uint32_t, 8> shifts_ = {{0, 4, 8, 12, 16, 20, 24, 28}};
     auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
     auto wi = simd::Simd<uint32_t, S>(*w);
@@ -530,17 +542,20 @@ simd::Simd<float, S> mxfp4_extract_bits_simd(const uint32_t* w) {
     wi = wi & 0xf;
     simd::Simd<float, S> w_out;
     for (int i = 0; i < S; ++i) {
-      w_out[i] = MXFP4_LUT[wi[i]];
+      w_out[i] = FP4_LUT[wi[i]];
     }
     return w_out;
+  } else if constexpr (S == 8 && bits == 8) {
+    auto w_out = simd::load<uint8_t, S>(reinterpret_cast<const uint8_t*>(w));
+    return detail::FromFP8{}(w_out);
   } else {
     // Appease compiler.. but should never get here
     throw std::runtime_error("Unsupported combination for simd qmm.");
   }
 }
 
-template <typename T>
-void mxfp4_qmm_t_simd(
+template <typename T, int group_size, int bits>
+void fp_qmm_t_simd(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -548,8 +563,7 @@ void mxfp4_qmm_t_simd(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = 32 / 4;
+  constexpr int pack_factor = get_pack_factor(bits, 32);
   constexpr int packs_in_group = group_size / pack_factor;
   constexpr int S = simd::max_size<T>;
   static_assert(
@@ -564,12 +578,12 @@ void mxfp4_qmm_t_simd(
       simd::Simd<float, S> acc(0);
       auto x_local = x;
       for (int k = 0; k < K; k += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
 
         simd::Simd<float, S> g_acc(0);
         for (int kw = 0; kw < packs_in_group; kw += packs_per_simd) {
           // Extract bits
-          auto wf = mxfp4_extract_bits_simd<S>(w_local);
+          auto wf = fp_extract_bits_simd<S, bits>(w_local);
           w_local += packs_per_simd;
           simd::Simd<float, S> x_simd = simd::load<T, S>(x_local);
           g_acc = g_acc + x_simd * wf;
@@ -585,8 +599,8 @@ void mxfp4_qmm_t_simd(
   }
 }
 
-template <typename T>
-void mxfp4_qmm_dispatch_transpose(
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_transpose(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -598,17 +612,17 @@ void mxfp4_qmm_dispatch_transpose(
   if (transposed_w) {
     // the simd size must be a multiple of the number of elements per word
     if constexpr (simd::max_size<T> % 8 == 0) {
-      mxfp4_qmm_t_simd<T>(result, x, w, scales, M, N, K);
+      fp_qmm_t_simd<T, group_size, bits>(result, x, w, scales, M, N, K);
     } else {
-      mxfp4_qmm_t<T>(result, x, w, scales, M, N, K);
+      fp_qmm_t<T, group_size, bits>(result, x, w, scales, M, N, K);
     }
   } else {
-    mxfp4_qmm<T>(result, x, w, scales, M, N, K);
+    fp_qmm<T, group_size, bits>(result, x, w, scales, M, N, K);
   }
 }
 
-template <typename T>
-void mxfp4_qmm_dispatch_typed(
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_mode(
     array& out,
     const array& x,
     const array& w,
@@ -626,7 +640,7 @@ void mxfp4_qmm_dispatch_typed(
   auto w_ptr = w.data<uint32_t>();
   auto scales_ptr = scales.data<uint8_t>();
   for (int i = 0; i < batch_size; i++) {
-    mxfp4_qmm_dispatch_transpose<T>(
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
         out_ptr + i * M * N,
         x_ptr + elem_to_loc(i * M * K, x.shape(), x.strides()),
         w_ptr + elem_to_loc(i * w_els, w.shape(), w.strides()),
@@ -638,21 +652,44 @@ void mxfp4_qmm_dispatch_typed(
   }
 }
 
-void mxfp4_qmm_dispatch(
+template <typename T>
+void fp_qmm_dispatch_typed(
     array& out,
     const array& x,
     const array& w,
     const array& scales,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_qmm_dispatch_mode<T, 32, 8>(out, x, w, scales, transposed_w);
+  } else if (group_size == 32) {
+    fp_qmm_dispatch_mode<T, 32, 4>(out, x, w, scales, transposed_w);
+  } else {
+    fp_qmm_dispatch_mode<T, 16, 4>(out, x, w, scales, transposed_w);
+  }
+}
+
+void fp_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
     bool transposed_w) {
   switch (x.dtype()) {
     case bfloat16:
-      mxfp4_qmm_dispatch_typed<bfloat16_t>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<bfloat16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     case float16:
-      mxfp4_qmm_dispatch_typed<float16_t>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<float16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     case float32:
-      mxfp4_qmm_dispatch_typed<float>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<float>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     default:
       throw std::invalid_argument(
@@ -765,9 +802,8 @@ void _bs_qmm_dispatch(
           "[quantized_matmul] only floating types are supported");
   }
 }
-
-template <typename T>
-void mxfp4_bs_qmm_dispatch_typed(
+template <typename T, int group_size, int bits>
+void fp_bs_qmm_dispatch_mode(
     array& out,
     const array& x,
     const array& w,
@@ -794,7 +830,7 @@ void mxfp4_bs_qmm_dispatch_typed(
         i, lhs_indices.shape(), lhs_indices.strides())];
     int w_idx = rhs_indices_ptr[elem_to_loc(
         i, rhs_indices.shape(), rhs_indices.strides())];
-    mxfp4_qmm_dispatch_transpose<T>(
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
         out_ptr + i * M * N,
         x_ptr + elem_to_loc(x_idx * M * K, x.shape(), x.strides()),
         w_ptr + elem_to_loc(w_idx * w_els, w.shape(), w.strides()),
@@ -807,26 +843,75 @@ void mxfp4_bs_qmm_dispatch_typed(
   }
 }
 
-void mxfp4_bs_qmm_dispatch(
+template <typename T>
+void fp_bs_qmm_dispatch_typed(
     array& out,
     const array& x,
     const array& w,
     const array& scales,
     const array& lhs_indices,
     const array& rhs_indices,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_bs_qmm_dispatch_mode<T, 32, 8>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else if (group_size == 32) {
+    fp_bs_qmm_dispatch_mode<T, 32, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else {
+    fp_bs_qmm_dispatch_mode<T, 16, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  }
+}
+
+void fp_bs_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int group_size,
+    int bits,
     bool transposed_w) {
   switch (x.dtype()) {
     case float32:
-      mxfp4_bs_qmm_dispatch_typed<float>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<float>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     case float16:
-      mxfp4_bs_qmm_dispatch_typed<float16_t>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<float16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     case bfloat16:
-      mxfp4_bs_qmm_dispatch_typed<bfloat16_t>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<bfloat16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     default:
       throw std::invalid_argument(
@@ -842,20 +927,9 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& scales_pre = inputs[2];
 
   auto& encoder = cpu::get_command_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);
+  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());
 
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -864,7 +938,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]);
+    auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
     encoder.set_input_array(biases);
     encoder.dispatch([out = array::unsafe_weak_copy(out),
                       x = array::unsafe_weak_copy(x),
@@ -881,8 +955,10 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
                       x = array::unsafe_weak_copy(x),
                       w = array::unsafe_weak_copy(w),
                       scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
                       transpose_ = transpose_]() mutable {
-      mxfp4_qmm_dispatch(out, x, w, scales, transpose_);
+      fp_qmm_dispatch(out, x, w, scales, group_size_, bits_, transpose_);
     });
   }
 }
@@ -953,13 +1029,122 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
                       scales = array::unsafe_weak_copy(scales),
                       lhs_indices = array::unsafe_weak_copy(lhs_indices),
                       rhs_indices = array::unsafe_weak_copy(rhs_indices),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
                       transpose_ = transpose_]() mutable {
-      mxfp4_bs_qmm_dispatch(
-          out, x, w, scales, lhs_indices, rhs_indices, transpose_);
+      fp_bs_qmm_dispatch(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size_,
+          bits_,
+          transpose_);
     });
   }
 }
 
+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,
@@ -1044,15 +1229,8 @@ void dispatch_quantize(
 void fast::Quantize::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  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& encoder = cpu::get_command_encoder(stream());
+  auto w = ensure_row_contiguous(inputs[0], encoder, stream());
   auto& out = outputs[0];
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -1060,10 +1238,6 @@ 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);
@@ -1146,6 +1320,43 @@ void fast::ConvertFP8::eval_cpu(
 }
 
 void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
-  throw std::runtime_error("QQMatmul not implemented on CPU.");
+  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");
+  }
 }
+
 } // namespace mlx::core

mlx/backend/cpu/simd/base_simd.h

@@ -1,11 +1,18 @@
 #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;
@@ -22,6 +29,14 @@ 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>
@@ -105,7 +120,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{((typeof(x))(0.5)) * std::log1p(r), theta}};
+      return Simd<T, 1>{T{((decltype(x))(0.5)) * std::log1p(r), theta}};
     } else {
       auto z0 = std::hypot(x + 1, y);
       return Simd<T, 1>{T{std::log(z0), theta}};
@@ -173,7 +188,16 @@ 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,10 +15,14 @@ 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 (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
+  if constexpr (is_floating_v<T> || std::is_same_v<T, complex64_t>) {
     if (std::isnan(a))
       return false;
     if (std::isnan(b))
@@ -103,11 +107,11 @@ struct StridedIterator {
     return *this;
   }
 
-  StridedIterator operator+(difference_type diff) {
+  StridedIterator operator+(difference_type diff) const {
     return StridedIterator(ptr_, stride_, diff);
   }
 
-  StridedIterator operator-(difference_type diff) {
+  StridedIterator operator-(difference_type diff) const {
     return StridedIterator(ptr_, stride_, -diff);
   }
 
@@ -198,7 +202,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 (std::is_floating_point<T>::value) {
+      if constexpr (is_floating_v<T>) {
         if (std::isnan(v1))
           return false;
         if (std::isnan(v2))
@@ -299,7 +303,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 (std::is_floating_point<T>::value) {
+      if constexpr (is_floating_v<T>) {
         if (std::isnan(v1))
           return false;
         if (std::isnan(v2))

mlx/backend/cpu/unary_ops.h

@@ -154,24 +154,19 @@ struct ToFP8 {
 struct FromFP8 {
   template <int N>
   Simd<float, N> operator()(Simd<uint8_t, N> x) {
-    auto w = Simd<uint32_t, N>(x) << 24;
-    auto sign = w & 0x80000000;
-    auto nonsign = w & 0x7FFFFFFF;
-
-    auto renorm_shift = clz(nonsign);
-    renorm_shift = simd::select(
-        renorm_shift > Simd<uint32_t, N>{4},
-        renorm_shift - Simd<uint32_t, N>{4},
-        Simd<uint32_t, N>{0});
-
-    Simd<int32_t, N> inf_nan_mask =
-        (Simd<int32_t, N>(nonsign + 0x01000000) >> 8) & 0x7F800000;
-    auto zero_mask = Simd<int32_t, N>(nonsign - 1) >> 31;
-    auto result = sign |
-        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
-          inf_nan_mask) &
-         ~zero_mask);
-    return fp32_from_bits(result);
+    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 sign = Simd<bool, N>(x & 128);
+    return select(sign, -out, out);
   }
   float operator()(uint8_t x) {
     return (*this)(Simd<uint8_t, 1>(x)).value;

mlx/backend/cuda/CMakeLists.txt

@@ -19,16 +19,21 @@ 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
@@ -56,21 +61,24 @@ 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.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu)
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
+                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
 endif()
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
@@ -112,22 +120,26 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
-# 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)
+# 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.
   target_compile_options(
-    mlx
-    PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--static-global-template-stub=false>")
+    mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
 endif()
 
-# Suppress warning when building for compute capability 7 used by V100.
+# Suppress nvcc warnings on C++ headers.
 target_compile_options(
-  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
+  mlx
+  PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=27,997,1394,20011,20208">
+)
 
-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)
+# 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">)
 
 # Use stronger binaries compression. This feature was introduced in CUDA 12.8
 # and requires drivers released after CUDA 12.4.
@@ -142,12 +154,11 @@ 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 IN_LIST UPGRADABLE_ARCHITECTURES)
+  elseif(MLX_CUDA_ARCHITECTURES GREATER_EQUAL 90)
     # Use arch-specific compute capability whenever possible.
     set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
   endif()
@@ -156,12 +167,60 @@ 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()
+endif()
+
 # ------------------------ Dependencies ------------------------
 
 # Use fixed version of CCCL.
 FetchContent_Declare(
   cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
+  URL "https://github.com/NVIDIA/cccl/releases/download/v3.1.3/cccl-v3.1.3.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
 
@@ -189,12 +248,14 @@ 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)
 
@@ -212,5 +273,15 @@ 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_SHALLOW TRUE
+  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(cutlass)
+target_include_directories(
+  mlx SYSTEM PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)

mlx/backend/cuda/allocator.cpp

@@ -3,13 +3,17 @@
 #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 {
 
@@ -20,6 +24,70 @@ 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;
@@ -33,24 +101,21 @@ 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_;
 
-  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));
+  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));
+      }
+    }
   }
 
   auto curr = next_free_;
@@ -62,7 +127,7 @@ SmallSizePool::SmallSizePool() {
 }
 
 SmallSizePool::~SmallSizePool() {
-  CHECK_CUDA_ERROR(cudaFree(data_));
+  unified_free(data_);
   delete[] buffer_;
 }
 
@@ -96,39 +161,23 @@ CudaAllocator::CudaAllocator()
     : buffer_cache_(
           page_size,
           [](CudaBuffer* buf) { return buf->size; },
-          [this](CudaBuffer* buf) { cuda_free(buf); }) {
+          [this](CudaBuffer* buf) { free_cuda_buffer(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 = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
-  int curr;
-  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
+  int device_count = gpu::device_count();
+  free_streams_.resize(device_count);
+  mem_pools_.resize(device_count);
   for (int i = 0; i < device_count; ++i) {
-    CHECK_CUDA_ERROR(cudaSetDevice(i));
-    cudaStream_t s;
-    CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));
-    free_streams_.push_back(s);
-
-    cudaMemPool_t mem_pool;
-    CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pool, i));
-    mem_pools_.push_back(mem_pool);
+    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(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
@@ -137,8 +186,6 @@ 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) {
@@ -151,6 +198,8 @@ 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.
@@ -168,9 +217,14 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     if (!buf) {
       void* data = nullptr;
       if (device == -1) {
-        CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
+        data = unified_malloc(size);
       } else {
-        CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
+        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));
+        }
       }
       if (!data) {
         std::ostringstream msg;
@@ -186,12 +240,14 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     // from OOM
     if (get_cache_memory() > 0) {
       for (auto p : mem_pools_) {
-        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;
+        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;
+          }
         }
       }
     }
@@ -203,9 +259,10 @@ 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_);
   }
-  // Copy to managed here if the buffer is not on the right device
+  lock.unlock();
+  // Copy to unified memory here if the buffer is not on the right device.
   if (buf->device >= 0 && buf->device != device) {
-    copy_to_managed(*buf);
+    move_to_unified_memory(*buf, stream);
   }
   return Buffer{buf};
 }
@@ -229,7 +286,7 @@ void CudaAllocator::free(Buffer buffer) {
   if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
-    cuda_free(buf);
+    free_cuda_buffer(buf);
   }
 }
 
@@ -241,20 +298,52 @@ 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::cuda_free(CudaBuffer* buf) {
+void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) {
   if (scalar_pool_.in_pool(buf)) {
     scalar_pool_.free(buf);
   } else {
-    if (buf->device >= 0) {
-      CHECK_CUDA_ERROR(cudaFreeAsync(buf->data, free_streams_[buf->device]));
-    } else {
-      CHECK_CUDA_ERROR(cudaFree(buf->data));
-    }
+    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));
+    } else {
+      CHECK_CUDA_ERROR(cudaFree(buf.data));
+    }
+  }
+}
+
 size_t CudaAllocator::get_active_memory() const {
   return active_memory_;
 }
@@ -294,22 +383,20 @@ void CudaAllocator::clear_cache() {
 }
 
 CudaAllocator& allocator() {
-  // By creating the |allocator_| on heap, the destructor of CudaAllocator
-  // will not be called on exit and buffers in the cache will be leaked. This
-  // can save some time at program exit.
-  static CudaAllocator* allocator_ = new CudaAllocator;
+  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();
+  }();
   return *allocator_;
 }
 
 Buffer malloc_async(size_t size, CommandEncoder& encoder) {
-  auto buffer = allocator().malloc_async(
+  return 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
@@ -325,9 +412,7 @@ void* Buffer::raw_ptr() {
     return nullptr;
   }
   auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
-  if (cbuf.device != -1) {
-    copy_to_managed(cbuf);
-  }
+  cu::allocator().move_to_unified_memory(cbuf);
   return cbuf.data;
 }
 

mlx/backend/cuda/allocator.h

@@ -54,6 +54,10 @@ 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();
@@ -64,7 +68,8 @@ class CudaAllocator : public allocator::Allocator {
   void clear_cache();
 
  private:
-  void cuda_free(CudaBuffer* buf);
+  void free_cuda_buffer(CudaBuffer* buf);
+  void free_async(CudaBuffer& buf, cudaStream_t stream = nullptr);
 
   CudaAllocator();
   friend CudaAllocator& allocator();
@@ -77,7 +82,7 @@ class CudaAllocator : public allocator::Allocator {
   BufferCache<CudaBuffer> buffer_cache_;
   size_t active_memory_{0};
   size_t peak_memory_{0};
-  std::vector<cudaStream_t> free_streams_;
+  std::vector<CudaStream> free_streams_;
   std::vector<cudaMemPool_t> mem_pools_;
   SmallSizePool scalar_pool_;
 };

mlx/backend/cuda/arange.cu

@@ -56,7 +56,6 @@ 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,7 +172,6 @@ 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,8 +16,14 @@ 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__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) 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) {
@@ -36,7 +42,11 @@ __global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) 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) {
@@ -57,7 +67,11 @@ __global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) 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) {
@@ -78,7 +92,11 @@ __global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) 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) {
@@ -291,7 +309,6 @@ 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),
@@ -309,7 +326,6 @@ 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),
@@ -333,12 +349,16 @@ 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);
+                out.data_size(),
+                out.shape(),
+                out.strides(),
+                large(),
+                N_READS,
+                cu::BINARY_MAX_BLOCK_DIM);
             encoder.add_kernel_node(
                 kernel,
                 num_blocks,
                 block_dims,
-                0,
                 gpu_ptr<InType>(a),
                 gpu_ptr<InType>(b),
                 gpu_ptr<OutType>(out),
@@ -346,11 +366,12 @@ 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,7 +314,6 @@ 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),
@@ -333,7 +332,6 @@ 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),
@@ -367,7 +365,6 @@ 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),
@@ -376,11 +373,12 @@ 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,14 +36,16 @@ 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(
@@ -250,20 +252,30 @@ 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 (auto wpt : std::array<int, 2>{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 (int wpt : {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));
       }
     }
 
@@ -339,7 +351,8 @@ 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(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      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 LRUBytesKeyCache<
+  static thread_local 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(encoder.device().cudnn_handle(), dtype, compute_dtype);
+  DnnGraph graph(get_cudnn_handle(encoder.device()), 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_FE_ERROR(graph.build());
+  CHECK_CUDNN_ERROR(graph.build());
   return graph;
 }
 
@@ -252,6 +252,10 @@ 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) {
@@ -269,20 +273,19 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
 
   // Search cache.
   BytesKey<ConvCacheKey> cache_key;
-  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)};
+  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);
   if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
     auto& [backend_type, graph] = it->second;
     if (graph) {
@@ -290,7 +293,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_FE_ERROR(graph->encode_capturing(
+      CHECK_CUDNN_ERROR(graph->encode_capturing(
           encoder,
           {
               {'x', gpu_ptr<void>(in)},
@@ -372,7 +375,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_FE_ERROR(graph->encode_capturing(
+    CHECK_CUDNN_ERROR(graph->encode_capturing(
         encoder,
         {
             {'x', gpu_ptr<void>(in)},

mlx/backend/cuda/conv/gemm_conv.cu

@@ -117,7 +117,6 @@ 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,7 +120,6 @@ 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,