Add free threaded build

Use uv in macOS CI
[CUDA] Fix half type matmul in cutlass kernels (#3469 )
2026-05-06 17:27:07 -07:00 · 2026-05-06 15:41:43 +09:00 · 2026-05-06 08:35:53 +09:00 · 2026-05-05 09:53:36 +09:00 · 2026-05-05 09:22:23 +09:00 · 2026-05-05 09:06:23 +09:00
649 changed files with 59125 additions and 12002 deletions
@@ -1,724 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
-
-parameters:
-  nightly_build:
-    type: boolean
-    default: false
-  test_release:
-    type: boolean
-    default: false
-
-jobs:
-  build_documentation:
-    parameters:
-      upload-docs:
-        type: boolean
-        default: false
-    macos:
-      xcode: "16.2.0"
-    resource_class: m2pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install
-          command: |
-            brew install python@3.9
-            brew install doxygen
-            python3.9 -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install -r docs/requirements.txt
-            pip install . -v
-      - when:
-          condition:
-            not: << parameters.upload-docs >>
-          steps:
-            - run:
-               name: Build documentation
-               command: |
-                 source env/bin/activate
-                 cd docs && doxygen && make html O=-W
-      - when:
-          condition: << parameters.upload-docs >>
-          steps:
-            - add_ssh_keys:
-                fingerprints:
-                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
-            - run:
-               name: Upload documentation
-               command: |
-                 source env/bin/activate
-                 git config user.email "mlx@group.apple.com"
-                 git config user.name "CircleCI Docs"
-                 git checkout gh-pages
-                 git rebase main
-                 cd docs
-                 git rm -rf build/html
-                 doxygen && make html O=-W
-                 git add -f build/html
-                 git commit -m "rebase"
-                 git push -f origin gh-pages
-
-  linux_build_and_test:
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Run style checks
-          command: |
-            pip install pre-commit
-            pre-commit run --all
-            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
-      - run:
-          name: Install dependencies
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            uv pip install -e ".[dev]" -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            python -m unittest discover python/tests -v
-            mpirun --bind-to none -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-            if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build
-            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-            make -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests
-
-  mac_build_and_test:
-    parameters:
-      xcode_version:
-        type: string
-        default: "16.2.0"
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    resource_class: m2pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
-              brew install openmpi uv
-      - run:
-          name: Install Python package
-          command: |
-            uv venv --python 3.9
-            uv pip install \
-              nanobind==2.4.0 \
-              cmake \
-              numpy \
-              torch \
-              tensorflow \
-              unittest-xml-reporting
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
-            mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-            if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-      - run:
-          name: Build example extension
-          command: |
-            source .venv/bin/activate
-            cd examples/extensions
-            uv pip install -r requirements.txt
-            uv run --no-project setup.py build_ext --inplace
-            uv run --no-project python test.py
-      - store_test_results:
-          path: test-results
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
-      - run:
-          name: Run CPP tests
-          command: |
-            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
-      - run:
-          name: Build small binary
-          command: |
-            source .venv/bin/activate
-            cd build/
-            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
-              -DBUILD_SHARED_LIBS=ON \
-              -DMLX_BUILD_CPU=OFF \
-              -DMLX_BUILD_SAFETENSORS=OFF \
-              -DMLX_BUILD_GGUF=OFF \
-              -DMLX_METAL_JIT=ON
-            make -j `sysctl -n hw.ncpu`
-      - run:
-          name: Run Python tests with JIT
-          command: |
-            CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-              uv pip install -e .
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
-              METAL_DEBUG_ERROR_MODE=0 \
-              uv run --no-project python -m xmlrunner discover \
-                -v python/tests \
-                -o test-results/gpu_jit
-
-  cuda_build_and_test:
-    parameters:
-      image_date:
-        type: string
-        default: "2023.11.1"
-    machine:
-      image: "linux-cuda-12:<< parameters.image_date >>"
-      resource_class: gpu.nvidia.small.gen2
-    steps:
-      - checkout
-      - restore_cache:
-          keys:
-            - cuda-<< parameters.image_date >>-{{ arch }}-
-      - run:
-          name: Install dependencies
-          command: |
-            sudo apt-get update
-            sudo apt-get install libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
-            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
-            rm -rf ccache-4.11.3-linux-x86_64
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
-            LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
-      - run:
-          name: CCache report
-          command: |
-            ccache --show-stats
-            ccache --zero-stats
-            ccache --max-size 400MB
-            ccache --cleanup
-      - save_cache:
-          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
-          paths:
-            - /home/circleci/.cache/ccache
-
-  build_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.9"
-      xcode_version:
-        type: string
-        default: "16.2.0"
-      build_env:
-        type: string
-        default: ""
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    resource_class: m2pro.medium
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            brew install python@<< parameters.python_version >>
-            brew install openmpi
-            python<< parameters.python_version >> -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install nanobind==2.4.0
-            pip install --upgrade setuptools
-            pip install numpy
-            pip install twine
-            pip install build
-      - run:
-          name: Install Python package
-          command: |
-            source env/bin/activate
-            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
-              pip install . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            source env/bin/activate
-            pip install typing_extensions
-            python setup.py generate_stubs
-      - run:
-          name: Build Python package
-          command: |
-            source env/bin/activate
-            python setup.py clean --all
-            << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
-      - when:
-          condition:
-            equal: ["3.9", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  source env/bin/activate
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 python -m build -w
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  source env/bin/activate
-                  twine upload dist/*
-      - store_artifacts:
-          path: dist/
-
-  build_linux_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.9"
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            PYTHON=python<< parameters.python_version >>
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            TZ=Etc/UTC sudo apt-get -y install tzdata
-            sudo add-apt-repository -y ppa:deadsnakes/ppa
-            sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            $PYTHON -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            << parameters.build_env >> pip install ".[dev]" -v
-            pip install typing_extensions
-            python setup.py generate_stubs
-            python setup.py clean --all
-            MLX_BUILD_STAGE=1 << parameters.build_env >> python -m build -w
-            bash python/scripts/repair_linux.sh
-      - when:
-          condition:
-            equal: ["3.9", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  source env/bin/activate
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 \
-                    python -m build -w
-                  auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload packages
-                command: |
-                  source env/bin/activate
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-  build_cuda_release:
-    parameters:
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
-            sudo dpkg -i cuda-keyring_1.1-1_all.deb
-            sudo apt-get update
-            sudo apt-get install cuda-toolkit-12-9 libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install zip
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
-            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-            << parameters.build_env >> MLX_BUILD_STAGE=2 \
-              CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              python -m build -w
-            bash python/scripts/repair_cuda.sh
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-workflows:
-  build_and_test:
-    when:
-      and:
-        - matches:
-            pattern: "^(?!pull/)[-\\w]+$"
-            value: << pipeline.git.branch >>
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - mac_build_and_test:
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "14.0"]
-      - linux_build_and_test
-      - cuda_build_and_test:
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-      - build_documentation 
-
-  build_pypi_release:
-    when:
-      and:
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["PYPI_RELEASE=1"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "PYPI_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "PYPI_RELEASE=1"
-      - build_documentation:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          upload-docs: true
-      - build_linux_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              build_env: ["PYPI_RELEASE=1"]
-      - build_cuda_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              build_env: ["PYPI_RELEASE=1"]
-
-  prb:
-    when:
-      matches:
-        pattern: "^pull/\\d+(/head)?$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - hold:
-          type: approval
-      - apple/authenticate:
-          context: pr-approval
-      - mac_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "14.0"]
-      - linux_build_and_test:
-          requires: [ hold ]
-      - cuda_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-  nightly_build:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.nightly_build >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-      - build_cuda_release
-
-  build_dev_release:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["DEV_RELEASE=1"]
-              xcode_version: ["16.2.0", "15.0.0"]
-            exclude:
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "13.5"
-                xcode_version: "16.2.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "14.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.9"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.10"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.11"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.12"
-                build_env: "DEV_RELEASE=1"
-              - macosx_deployment_target: "15.0"
-                xcode_version: "15.0.0"
-                python_version: "3.13"
-                build_env: "DEV_RELEASE=1"
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              build_env: ["DEV_RELEASE=1"]
-      - build_cuda_release:
-          matrix:
-            parameters:
-              build_env: ["DEV_RELEASE=1"]
@@ -0,0 +1,31 @@
+name: 'Build CUDA wheel'
+description: 'Build CUDA wheel'
+
+inputs:
+  arch:
+    description: 'Platform architecture tag'
+    required: true
+    type: choice
+    options:
+      - x86_64
+      - aarch64
+
+runs:
+  using: "composite"
+  steps:
+    - name: Build package
+      shell: bash
+      env:
+        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
+      run: |
+        pip install auditwheel "build<=1.4.2" patchelf setuptools
+        python setup.py clean --all
+        MLX_BUILD_STAGE=2 python -m build -w
+
+        auditwheel repair dist/mlx_cuda*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libcublas* \
+          --exclude libcuda* \
+          --exclude libcudnn* \
+          --exclude libnccl* \
+          --exclude libnvrtc*
@@ -0,0 +1,38 @@
+name: 'Build Documentation'
+description: 'Build documentation'
+
+runs:
+  using: "composite"
+  steps:
+    - name: Setup machine
+      uses: ./.github/actions/setup-linux
+
+    - name: Install dependencies
+      shell: bash
+      run: |
+        sudo apt-get install -y doxygen
+        source .venv/bin/activate
+        pip install -r docs/requirements.txt
+        pip install . -v
+  
+    - name: Build documentation
+      shell: bash
+      run: |
+        source .venv/bin/activate
+        cd docs
+        doxygen
+        make html O=-W
+    
+    - name: Create artifact tar
+      shell: bash
+      run: tar -cf artifact.tar -C docs --dereference build/html index.html
+
+    # Do it manually because upload-pages-artifact requires gtar
+    - name: Upload artifact
+      id: upload-artifact
+      uses: actions/upload-artifact@v5
+      with:
+        name: github-pages
+        path: artifact.tar
+        retention-days: 1
+        if-no-files-found: error
@@ -0,0 +1,42 @@
+name: 'Build Linux wheel'
+description: 'Build Linux wheel'
+
+inputs:
+  build-backend:
+    description: 'Build the backend mlx-cpu package'
+    type: boolean
+    required: false
+    default: false
+  arch:
+    description: 'Platform architecture tag'
+    required: true
+    type: choice
+    options:
+      - x86_64
+      - aarch64
+
+runs:
+  using: "composite"
+  steps:
+    - name: Build MLX
+      shell: bash
+      run: pip install -e . -v
+
+    - name: Build Python package
+      shell: bash
+      run: |
+        pip install auditwheel patchelf "build<=1.4.2"
+        python setup.py clean --all
+        MLX_BUILD_STAGE=1 python -m build -w
+        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
+      run: |
+        python setup.py clean --all
+        MLX_BUILD_STAGE=2 python -m build -w
+        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_${{ inputs.arch }}
@@ -0,0 +1,38 @@
+name: 'Build and Test on Linux'
+
+inputs:
+  toolkit:
+    description: 'The toolkit to build with'
+    required: false
+    default: 'cpu'
+
+runs:
+  using: "composite"
+  steps:
+
+    - name: Install Python package
+      id: python_build
+      shell: sh
+      env:
+        DEBUG: 1
+        CMAKE_ARGS: >-
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }}
+      run: |
+        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
+          # There is no GPU in arm64 runner, use a common arch.
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
+          # 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: Build CPP only
+      shell: bash
+      run: |
+        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }}
+        cmake --build build -j $(nproc)
@@ -0,0 +1,36 @@
+name: 'Build macOS release'
+description: 'Build MLX releases macOS'
+
+inputs:
+  macos-target:
+    description: 'macOS build target'
+    required: false
+    default: '15.0'
+  build-backend:
+    description: 'Build the backend mlx-metal package'
+    type: boolean
+    required: false
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Build Python package
+      shell: bash -l {0}
+      env:
+        DEVELOPER_DIR: /Applications/Xcode-latest.app
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
+      run: |
+        uv pip install build
+        python setup.py clean --all
+        MLX_BUILD_STAGE=1 python -m build -w
+
+    - name: Build backend package
+      if: ${{ inputs.build-backend }}
+      shell: bash -l {0}
+      env:
+        DEVELOPER_DIR: /Applications/Xcode-latest.app
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
+      run: |
+        python setup.py clean --all
+        MLX_BUILD_STAGE=2 python -m build -w
@@ -0,0 +1,96 @@
+name: 'Build and Test on macOS'
+description: 'Build and test MLX on macOS'
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install Python package
+      env:
+        DEBUG: 1
+        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+      shell: bash
+      run: |
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
+
+    - name: Install tests dependencies
+      shell: bash
+      run: |
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
+
+    - name: Run Python tests
+      shell: bash
+      env:
+        LOW_MEMORY: 1
+      run: |
+        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
+      run: |
+        echo "::group::Build example extension"
+        cd examples/extensions
+        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
+      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
+      env:
+        DEVICE: gpu
+        METAL_DEVICE_WRAPPER_TYPE: 1
+        METAL_DEBUG_ERROR_MODE: 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 \
+          -DBUILD_SHARED_LIBS=ON \
+          -DMLX_BUILD_CPU=OFF \
+          -DMLX_BUILD_SAFETENSORS=OFF \
+          -DMLX_BUILD_GGUF=OFF \
+          -DMLX_METAL_JIT=ON
+        make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
+
+    - name: Run Python tests with JIT
+      shell: bash
+      env:
+        LOW_MEMORY: 1
+        DEVICE: gpu
+        METAL_DEVICE_WRAPPER_TYPE: 1
+        METAL_DEBUG_ERROR_MODE: 0
+      run: |
+        echo "::group::Run Python tests with JIT"
+        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
+          uv pip install -e . -v
+        python -m unittest discover -v python/tests
+        echo "::endgroup::"
@@ -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%
@@ -0,0 +1,102 @@
+name: 'Setup Linux Environment'
+description: 'Install dependencies for Linux builds'
+
+inputs:
+  toolkit:
+    description: 'Which toolkit to install'
+    required: false
+    default: 'cpu'
+  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'
+  ccache-key:
+    required: false
+    default: 'ccache'
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install common dependencies
+      shell: bash
+      run: |
+        echo "::group::Install common dependencies"
+        sudo apt-get update
+        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: ${{ 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
+
+    - uses: actions/setup-python@v6
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        python -m venv .venv
+        source .venv/bin/activate
+        pip install setuptools cmake typing_extensions
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"
+
+    - name: Set swap space
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
+      uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
+      with:
+        swap-size-gb: 16
+
+    - name: Install CUDA toolkit
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
+      shell: bash
+      env:
+        # Note: the CI machine does not meet CUDA 13's driver requirement.
+        # Compatibility matrix:
+        # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
+        PACKAGES: |
+          {
+            "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 --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: |
+        echo "::group::Installed NVIDIA and CUDA packages"
+        dpkg -l | egrep "cuda|nvidia" -i
+        echo "::endgroup::"
+        echo "::group::NVIDIA-SMI Status"
+        nvidia-smi || true
+        echo "::endgroup::"
@@ -0,0 +1,32 @@
+name: 'Setup macOS Environment'
+description: 'Install dependencies for macOS builds'
+
+inputs:
+  python-version:
+    description: 'Python version to use'
+    required: false
+    default: '3.10'
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install Homebrew packages
+      shell: sh
+      run: /opt/homebrew/bin/brew install openmpi
+
+    - name: Verify MetalToolchain installed
+      shell: bash
+      run: xcodebuild -showComponent MetalToolchain
+
+    - 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::"
@@ -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
@@ -0,0 +1,83 @@
+name: 'Run Linux tests'
+
+inputs:
+  has-gpu:
+    description: 'Run GPU tests'
+    required: false
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    # FIXME: The distributed tests fail with free-threading Python.
+    - name: Check free-threading Python
+      id: is-free-threading
+      shell: bash
+      run: |
+        echo "::group::Check free-threading Python"
+        if python -VV 2>&1 | grep "free-threading"; then
+          echo "result=true" >> $GITHUB_OUTPUT
+        else
+          echo "result=false" >> $GITHUB_OUTPUT
+        fi
+        echo "::endgroup::"
+
+    - name: Run MPI tests
+      if: ${{ steps.is-free-threading.outputs.result == 'false' }}
+      shell: bash
+      run: |
+        echo "::group::MPI tests"
+        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
+        echo "::endgroup::"
+
+    - name: Run distributed tests
+      if: ${{ steps.is-free-threading.outputs.result == 'false' && inputs.has-gpu == 'false' }}
+      shell: bash
+      run: |
+        echo "::group::Distributed tests"
+        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
+        if grep -Fq '[WARN]' stderr.log ; then
+          grep -F '[WARN]' stderr.log
+          echo "Distributed ring test failed";
+          exit 1;
+        fi
+        echo "::endgroup::"
+
+    - name: Run Python tests - CPU
+      if: ${{ inputs.has-gpu == 'false' }}
+      shell: bash
+      env:
+        DEVICE: cpu
+      run: |
+        echo "::group::Python tests - CPU"
+        python -m unittest discover python/tests -v
+        echo "::endgroup::"
+
+    - name: Run Python tests - GPU
+      if: ${{ inputs.has-gpu == 'true' }}
+      shell: bash
+      env:
+        DEVICE: gpu
+      run: |
+        echo "::group::Python tests - GPU"
+        python -m tests discover python/tests -v
+        echo "::endgroup::"
+
+    - name: Run CPP tests - CPU
+      shell: bash
+      env:
+        DEVICE: cpu
+      run: |
+        echo "::group::CPP tests - CPU"
+        ./build/tests/tests
+        echo "::endgroup::"
+
+    - name: Run CPP tests - GPU
+      if: ${{ inputs.has-gpu == 'true' }}
+      shell: bash
+      env:
+        DEVICE: gpu
+      run: |
+        echo "::group::CPP tests - GPU"
+        ./build/tests/tests -sfe="*linalg_tests.cpp"
+        echo "::endgroup::"
@@ -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::"
@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
@@ -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
@@ -0,0 +1,27 @@
+#!/bin/bash
+set -ex
+
+# [Setup] Install dependencies inside the container.
+dnf update -y
+dnf install -y \
+  blas-devel \
+  lapack-devel \
+  openblas-devel \
+  make \
+  cmake \
+  clang \
+  git
+dnf clean all
+
+# [C++] CI Build Sanity Check: Verifies code compilation, not for release.
+export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+export DEBUG=1
+export CMAKE_C_COMPILER=/usr/bin/clang
+export CMAKE_CXX_COMPILER=/usr/bin/clang++
+
+mkdir -p build
+pushd build
+cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
+make -j $(nproc)
+./tests/tests
+popd
@@ -0,0 +1,152 @@
+name: Build and Test
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      # For testing CI without starting a pull request:
+      - test/*
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+jobs:
+  check_lint:
+    name: Check Lint
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v6
+      - uses: pre-commit/action@v3.0.1
+
+  linux_build_and_test:
+    name: Linux (cpu, ${{ matrix.arch }})
+    needs: check_lint
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: ['x86_64', 'aarch64']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
+    steps:
+      - uses: actions/checkout@v6
+      - 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 }})
+    if: github.repository == 'ml-explore/mlx'
+    needs: check_lint
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: ['x86_64', 'aarch64']
+        toolkit: ['cuda-12.6', 'cuda-12.9']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'gpu-t4-4-core' || 'ubuntu-22.04-arm' }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-linux
+        with:
+          toolkit: ${{ matrix.toolkit }}
+      - uses: ./.github/actions/build-linux
+        with:
+          toolkit: ${{ matrix.toolkit }}
+      - uses: ./.github/actions/test-linux
+        if: matrix.arch == 'x86_64'
+        with:
+          has-gpu: true
+
+  mac_build_and_test:
+    name: macOS (${{ matrix.macos-target }})
+    if: github.repository == 'ml-explore/mlx'
+    strategy:
+      matrix:
+        macos-target: ["14.0", "15.0", "26.0"]
+    runs-on: [self-hosted, macos]
+    env:
+      MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }}
+    needs: check_lint
+    steps:
+      - uses: actions/checkout@v6
+      - 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'
+    runs-on: ubuntu-22.04
+    needs: check_lint
+    steps:
+      - 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
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - host: ubuntu-22.04
+            arch: x86_64
+          - host: ubuntu-22.04-arm
+            arch: aarch64
+
+    runs-on: ${{ matrix.host }}
+    container:
+      image: fedora:42
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+
+      - name: CPP Build Test - No Release
+        run: |
+          bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh
@@ -0,0 +1,28 @@
+name: Documentation
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/build-docs
+      
+  deploy:
+    needs: build
+    permissions:
+      pages: write
+      id-token: write
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v5
@@ -0,0 +1,108 @@
+name: Nightly Build
+
+on:
+  schedule:
+    - cron: 33 6 * * 1-5
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build_linux_release:
+    strategy:
+      fail-fast: false
+      matrix:
+        python_version: ["3.10", "3.14"]
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-linux
+      - uses: ./.github/actions/build-linux-release
+        with:
+          build-backend: ${{ matrix.python-version == '3.10' }}
+          arch: "x86_64"
+      - name: Upload mlx artifacts
+        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@v7
+        with:
+          name: mlx-cpu
+          path: wheelhouse/mlx_cpu-*.whl
+          retention-days: 7
+      - run: df -h
+
+  build_linux_with_tests:
+    strategy:
+      fail-fast: false
+      matrix:
+        python_version: ["3.11", "3.12", "3.13", "3.14", "3.14t"]
+        runner:
+          - ubuntu-22.04
+          - ubuntu-22.04-arm
+    runs-on: ${{ matrix.runner }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-linux
+        with:
+          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'
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.13", "3.14t"]
+    runs-on: [self-hosted, macos]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-macos
+        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:
+          macos-target: 15.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 14.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+
+  build_cuda_release:
+    if: github.repository == 'ml-explore/mlx'
+    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: ${{ matrix.toolkit }}
+          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@v7
+        with:
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          path: wheelhouse/mlx_cuda_*.whl
+          retention-days: 7
@@ -1,20 +0,0 @@
-on:
-  pull_request:
-    branches:
-      - main
-
-jobs:
-  check_lint:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v4
-        with:
-          python-version: 3.8
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install pre-commit black isort clang-format
-      - name: Run lint
-        run: |
-          pre-commit run --all-files
@@ -0,0 +1,251 @@
+name: PyPI Release
+
+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: 'Development release (DEV_RELEASE=1)'
+        required: false
+        type: boolean
+
+permissions:
+  contents: read
+
+jobs:
+  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
+      id-token: write
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v5
+
+  build_linux_release:
+    if: github.repository == 'ml-explore/mlx'
+    strategy:
+      matrix:
+        python_version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
+        arch: ['x86_64', 'aarch64']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
+    env:
+      PYPI_RELEASE: 1
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-linux
+        with:
+          python-version: ${{ matrix.python_version }}
+          use-ccache: false
+      - uses: ./.github/actions/build-linux-release
+        with:
+          build-backend: ${{ matrix.python_version == '3.10' }}
+          arch: ${{ matrix.arch }}
+      - name: Upload MLX artifacts
+        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@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:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
+    runs-on: [self-hosted, macos]
+    env:
+      PYPI_RELEASE: 1
+      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 Python package
+        run: uv pip install -e . -v
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 14.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 15 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 15.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 26 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 26.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Upload MLX artifacts
+        uses: actions/upload-artifact@v7
+        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@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'
+    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' }}
+    env:
+      PYPI_RELEASE: 1
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-linux
+        with:
+          toolkit: ${{ matrix.toolkit }}
+          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@v7
+        with:
+          overwrite: true
+          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: [build_linux_release, build_mac_release]
+    permissions:
+      id-token: write
+    environment:
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      url: https://pypi.org/p/mlx
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          pattern: linux-wheels-*
+          merge-multiple: true
+          path: dist
+      - uses: actions/download-artifact@v8
+        with:
+          pattern: mac-wheels-*
+          merge-multiple: true
+          path: dist
+      - name: Display structure of downloaded files
+        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: [build_cuda_release]
+    permissions:
+      id-token: write
+    environment:
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      url: https://pypi.org/p/mlx-cuda
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          pattern: mlx-cuda-*
+          merge-multiple: true
+          path: dist
+      - name: Display structure of downloaded files
+        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-cpu:
+    name: Upload CPU release to PyPI
+    runs-on: ubuntu-latest
+    needs: [build_linux_release]
+    permissions:
+      id-token: write
+    environment:
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      url: https://pypi.org/p/mlx-cpu
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          pattern: mlx-cpu-*
+          merge-multiple: true
+          path: dist
+      - name: Display structure of downloaded files
+        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-metal:
+    name: Upload Metal release to PyPI
+    runs-on: ubuntu-latest
+    needs: [build_mac_release]
+    permissions:
+      id-token: write
+    environment:
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
+      url: https://pypi.org/p/mlx-metal
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: mlx-metal
+          path: dist
+      - name: Display structure of downloaded files
+        run: du -ah dist
+      - name: Publish package distributions to PyPI
+        if: ${{ !inputs.dry_run }}
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://upload.pypi.org/legacy/
@@ -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
@@ -1,16 +1,22 @@
 repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+    -   id: check-yaml
+    # -   id: end-of-file-fixer
+    # -   id: trailing-whitespace
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v19.1.7
+    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:
@@ -19,12 +19,17 @@ MLX was developed with contributions from the following individuals:
 - Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
 - Paul Paczuski: Improved stability of BCE loss calculation
 - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops.
- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer.
+- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.

 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
  <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
 </a>

+# Organizations
+
+MLX has received contributions from the following companies:
+- NVIDIA Corporation & Affiliates
+
 # Third-Party Software

 MLX leverages several third-party software, listed here together with
@@ -22,10 +22,11 @@ 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)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

 # ----------------------------- Configuration -----------------------------
 option(MLX_BUILD_TESTS "Build tests for mlx" ON)
@@ -39,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(
@@ -73,12 +77,70 @@ endif()
 if(MLX_USE_CCACHE)
  find_program(CCACHE_PROGRAM ccache)
  if(CCACHE_PROGRAM)
+    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
    set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
    set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
  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)
@@ -87,22 +149,30 @@ cmake_policy(SET CMP0135 NEW)

 add_library(mlx)

-if(MLX_BUILD_METAL)
-  set(METAL_LIB "-framework Metal")
-  set(FOUNDATION_LIB "-framework Foundation")
-  set(QUARTZ_LIB "-framework QuartzCore")
-endif()
+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 AND NOT METAL_LIB)
-  message(STATUS "Metal not found. Unable to build GPU")
-  set(MLX_BUILD_METAL OFF)
-  set(MLX_METAL_DEBUG OFF)
-elseif(MLX_BUILD_METAL)
-  message(STATUS "Building METAL sources")
+if(MLX_BUILD_METAL)
+  find_library(METAL_LIB Metal)
+  find_library(FOUNDATION_LIB Foundation)
+  find_library(QUARTZ_LIB QuartzCore)
+  if(METAL_LIB)
+    message(STATUS "Metal found ${METAL_LIB}")
+  else()
+    message(
+      FATAL_ERROR
+        "Metal not found. Set MLX_BUILD_METAL=OFF to build without GPU")
+  endif()

  if(MLX_METAL_DEBUG)
    add_compile_definitions(MLX_METAL_DEBUG)
@@ -111,7 +181,8 @@ elseif(MLX_BUILD_METAL)
  # Throw an error if xcrun not found
  execute_process(
    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
-    OUTPUT_VARIABLE MACOS_SDK_VERSION COMMAND_ERROR_IS_FATAL ANY)
+    OUTPUT_VARIABLE MACOS_SDK_VERSION
+    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)

  if(${MACOS_SDK_VERSION} LESS 14.0)
    message(
@@ -121,9 +192,12 @@ elseif(MLX_BUILD_METAL)
  message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}")

  set(METAL_CPP_URL
-      https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip)

  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0)
+      message(FATAL_ERROR "MLX requires macOS >= 14.0")
+    endif()
    set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
  endif()
  execute_process(
@@ -132,7 +206,6 @@ elseif(MLX_BUILD_METAL)
      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
  FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
-
  FetchContent_MakeAvailable(metal_cpp)
  target_include_directories(
    mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
@@ -140,18 +213,27 @@ elseif(MLX_BUILD_METAL)
  target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
 endif()

+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  # With newer clang/gcc versions following libs are implicitly linked, but when
+  # building on old distributions they need to be explicitly listed.
+  target_link_libraries(mlx PRIVATE dl pthread)
+endif()
+
 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)
@@ -167,7 +249,7 @@ if(MLX_BUILD_CPU)
    message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
    set(MLX_BUILD_ACCELERATE ON)
  else()
-    message(STATUS "Accelerate or arm neon not found, using default backend.")
+    message(STATUS "Accelerate not found, using default backend.")
    set(MLX_BUILD_ACCELERATE OFF)
  endif()

@@ -175,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
@@ -234,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()
@@ -258,14 +351,16 @@ target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>)
 if(MLX_BUILD_PYTHON_BINDINGS)
  message(STATUS "Building Python bindings.")
  find_package(
-    Python 3.8
+    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()

@@ -285,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
@@ -2,7 +2,7 @@

 [**Quickstart**](#quickstart) | [**Installation**](#installation) |
 [**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
-[**Examples**](#examples) 
+[**Examples**](#examples)

 [![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)

@@ -11,37 +11,37 @@ brought to you by Apple machine learning research.

 Some key features of MLX include:

- - **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
   also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
   [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
   the Python API. MLX has higher-level packages like `mlx.nn` and
   `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
   more complex models.

- - **Composable function transformations**: MLX supports composable function
-   transformations for automatic differentiation, automatic vectorization,
-   and computation graph optimization.
+- **Composable function transformations**: MLX supports composable function
+  transformations for automatic differentiation, automatic vectorization,
+  and computation graph optimization.

- - **Lazy computation**: Computations in MLX are lazy. Arrays are only
-   materialized when needed.
+- **Lazy computation**: Computations in MLX are lazy. Arrays are only
+  materialized when needed.

- - **Dynamic graph construction**: Computation graphs in MLX are constructed
-   dynamically. Changing the shapes of function arguments does not trigger
-   slow compilations, and debugging is simple and intuitive.
+- **Dynamic graph construction**: Computation graphs in MLX are constructed
+  dynamically. Changing the shapes of function arguments does not trigger
+  slow compilations, and debugging is simple and intuitive.

- - **Multi-device**: Operations can run on any of the supported devices
-   (currently the CPU and the GPU).
+- **Multi-device**: Operations can run on any of the supported devices
+  (currently the CPU and the GPU).

- - **Unified memory**: A notable difference from MLX and other frameworks
-   is the *unified memory model*. Arrays in MLX live in shared memory.
-   Operations on MLX arrays can be performed on any of the supported
-   device types without transferring data.
+- **Unified memory**: A notable difference from MLX and other frameworks
+  is the *unified memory model*. Arrays in MLX live in shared memory.
+  Operations on MLX arrays can be performed on any of the supported
+  device types without transferring data.

 MLX is designed by machine learning researchers for machine learning
 researchers. The framework is intended to be user-friendly, but still efficient
 to train and deploy models. The design of the framework itself is also
 conceptually simple. We intend to make it easy for researchers to extend and
-improve MLX with the goal of quickly exploring new ideas. 
+improve MLX with the goal of quickly exploring new ideas.

 The design of MLX is inspired by frameworks like
 [NumPy](https://numpy.org/doc/stable/index.html),
@@ -91,7 +91,7 @@ Checkout the
 [documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
 for more information on building the C++ and Python APIs from source.

-## Contributing 
+## Contributing

 Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
 on contributing to MLX. See the
@@ -110,7 +110,7 @@ Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
 MLX useful in your research and wish to cite it, please use the following
 BibTex entry:

-```
+```text
@software{mlx2023,
  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
@@ -75,7 +75,7 @@ void time_irregular_binary_ops_3D() {

 void time_irregular_binary_ops_4D() {
  auto device = mx::default_device();
-  std::vector<int> shape = {8, 8, 512, 512};
+  mx::Shape shape = {8, 8, 512, 512};
  auto a = mx::random::uniform(shape);
  auto b = mx::random::uniform(shape);

@@ -115,7 +115,7 @@ void time_irregular_binary_ops_4D() {

 void time_irregular_reshape() {
  auto device = mx::default_device();
-  std::vector<int> shape;
+  mx::Shape shape;
  auto reshape_fn = [&shape, device](const mx::array& a) {
    return mx::reshape(a, shape, device);
  };
@@ -170,7 +170,7 @@ void time_irregular_astype_1D() {
 void time_irregular_astype_2D() {
  auto device = mx::default_device();
  int size = 2048;
-  std::vector<int> shape = {size, size};
+  mx::Shape shape = {size, size};

  auto a = mx::random::uniform(shape);
  TIMEM("2D regular", mx::astype, a, mx::int32, device);
@@ -142,9 +142,7 @@ def bench_shape(B, M, N, K, np_dtype, transpose="nn"):
    t_b = (0, 1, 2) if transpose[1] == "n" else (0, 2, 1)

    c_mlx = a_mx.transpose(t_a) @ b_mx.transpose(t_b)
-    c_npy = a_np.transpose(t_a).astype(np.float32) @ b_np.transpose(t_b).astype(
-        np.float32
-    )
+    c_npy = a_np.transpose(t_a).astype(np_dtype) @ b_np.transpose(t_b).astype(np_dtype)

    atol = 1e-5 if np_dtype == np.float32 else 1e-4

@@ -163,7 +161,7 @@ def get_gflop_count(B, M, N, K):
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run gemm benchmarks")

-    dtypes = ("float32", "float16")
+    dtypes = ("float32", "float16", "complex64")
    transposes = ("nn", "nt", "tn")
    shapes = (
        (16, 234, 768, 3072),
@@ -187,7 +185,7 @@ if __name__ == "__main__":
                diff = gflops_mx / gflops_pt - 1.0

                print(
-                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100. * diff:+5.2f}%"
+                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%"
                )
                if gflops_pt >= 2.0 * gflops_mx:
                    print("ATTENTION ^^^^^^^")
@@ -1,6 +1,5 @@
 # Copyright © 2023 Apple Inc.

-import argparse
 import os
 import subprocess
 import time
@@ -196,7 +195,7 @@ def bench_with_out_len(ax, out_vec_len, in_vector_lens, dtype, transpose):


 for transpose in (False, True):
-    for dtype in ("float32", "float16"):
+    for dtype in ("float32", "float16", "complex64"):
        fig, axs = plt.subplots(
            len(in_vec_sizes), 2, figsize=(8.5, 11), layout="constrained"
        )
@@ -215,7 +214,7 @@ for transpose in (False, True):
        fig.suptitle(f"{device_name}: {dtype} {op_name}")
        fig.savefig(
            os.path.join(
-                results_dir, f'{device_name.replace(" ", "_")}_{dtype}_{op_name}.pdf'
+                results_dir, f"{device_name.replace(' ', '_')}_{dtype}_{op_name}.pdf"
            )
        )
        plt.close(fig)
@@ -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()
@@ -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


@@ -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


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

@@ -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()
@@ -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}"
+            )
@@ -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
@@ -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()
@@ -0,0 +1,54 @@
+# FindNCCL.cmake This module finds the NVIDIA NCCL library and its include
+# directories.
+
+set(NCCL_ROOT_DIR
+    $ENV{NCCL_ROOT_DIR}
+    CACHE PATH "Folder contains NVIDIA NCCL")
+
+find_path(
+  NCCL_INCLUDE_DIRS
+  NAMES nccl.h
+  HINTS ${NCCL_INCLUDE_DIR} ${NCCL_ROOT_DIR} ${NCCL_ROOT_DIR}/include
+        ${CUDA_TOOLKIT_ROOT_DIR}/include)
+
+if($ENV{USE_STATIC_NCCL})
+  message(
+    STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library")
+  set(NCCL_LIBNAME "libnccl_static.a")
+else()
+  set(NCCL_LIBNAME "nccl")
+endif()
+
+find_library(
+  NCCL_LIBRARIES
+  NAMES ${NCCL_LIBNAME}
+  HINTS ${NCCL_LIB_DIR}
+        ${NCCL_ROOT_DIR}
+        ${NCCL_ROOT_DIR}/lib
+        ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu
+        ${NCCL_ROOT_DIR}/lib64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS
+                                  NCCL_LIBRARIES)
+
+if(NCCL_FOUND)
+  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h")
+  message(
+    STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}")
+  file(
+    STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED
+    REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$"
+    LIMIT_COUNT 1)
+  if(NCCL_MAJOR_VERSION_DEFINED)
+    string(REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" ""
+                         NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED})
+    message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}")
+  endif()
+  message(
+    STATUS
+      "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
+  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
+endif()
@@ -0,0 +1,3 @@
+# This file does nothing but to suppress the cmake warning: "By not providing
+# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
+# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.
@@ -26,6 +26,7 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
+PREDEFINED             = MLX_API=

 ################################################################################
 # Compound extraction control.                                                 #
@@ -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
+```
@@ -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>
@@ -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>
@@ -1,4 +1,5 @@
 sphinx
 breathe
 sphinx-book-theme
+sphinx-copybutton
 mlx
@@ -18,6 +18,7 @@ release = version
 # -- General configuration ---------------------------------------------------

 extensions = [
+    "sphinx_copybutton",
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.intersphinx",
@@ -127,7 +127,8 @@ relying on a copy from ``ensure_row_contiguous``:
      name="myexp_strided",
      input_names=["inp"],
      output_names=["out"],
-      source=source
+      source=source,
+      ensure_row_contiguous=False,
  )

  def exp_elementwise(a: mx.array):
@@ -138,7 +139,6 @@ relying on a copy from ``ensure_row_contiguous``:
          threadgroup=(256, 1, 1),
          output_shapes=[a.shape],
          output_dtypes=[a.dtype],
-          ensure_row_contiguous=False,
      )
      return outputs[0]

@@ -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)
@@ -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
@@ -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)


@@ -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())
@@ -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
@@ -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
@@ -70,11 +72,13 @@ are the CPU and GPU.
   python/fft
   python/linalg
   python/metal
+   python/cuda
   python/memory_management
   python/nn
   python/optimizers
   python/distributed
   python/tree_utils
+   python/printoptions

 .. toctree::
   :caption: C++ API Reference
@@ -88,5 +92,6 @@ are the CPU and GPU.

   dev/extensions
   dev/metal_debugger
+   dev/metal_logging
   dev/custom_metal_kernels
   dev/mlx_in_cpp
@@ -15,13 +15,12 @@ silicon computer is

 To install from PyPI your system must meet the following requirements:

- Using an M series chip (Apple silicon)
- Using a native Python >= 3.9
- macOS >= 13.5
+- Using `Apple silicon <https://support.apple.com/en-us/116943>`_
+- Using a native Python >= 3.10
+- macOS >= 14.0

 .. note::
-    MLX is only available on devices running macOS >= 13.5
-    It is highly recommended to use macOS 14 (Sonoma)
+    MLX is only available on devices running macOS >= 14.0 and higher.

 CUDA
 ^^^^
@@ -30,17 +29,20 @@ MLX has a CUDA backend which you can install with:

 .. code-block:: shell

-    pip install mlx[cuda]
+    pip install mlx[cuda12]
+

 To install the CUDA package from PyPi your system must meet the following
 requirements:

- Nvidia architecture >= SM 7.0 (Volta)
+- Nvidia architecture >= SM 7.5
 - Nvidia driver >= 550.54.14
 - CUDA toolkit >= 12.0
 - Linux distribution with glibc >= 2.35
- Python >= 3.9
+- Python >= 3.10

+For CUDA 13 use ``pip install mlx[cuda13]``. The CUDA 13 package requires
+an Nvidia driver >= 580 or an appropriate CUDA compatibility package.

 CPU-only (Linux)
 ^^^^^^^^^^^^^^^^
@@ -55,7 +57,7 @@ To install the CPU-only package from PyPi your system must meet the following
 requirements:

 - Linux distribution with glibc >= 2.35
- Python >= 3.9
+- Python >= 3.10


 Troubleshooting
@@ -81,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

@@ -126,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
 ^^^^^^^

@@ -271,7 +267,7 @@ and the CUDA toolkit. For example on Ubuntu, run the following:
   dpkg -i cuda-keyring_1.1-1_all.deb
   apt-get update -y
   apt-get -y install cuda-toolkit-12-9
-   apt-get install libblas-dev liblapack-dev liblapacke-dev -y
+   apt-get install libblas-dev liblapack-dev liblapacke-dev libcudnn9-dev-cuda-12 -y


 When building either the Python or C++ APIs make sure to pass the cmake flag
@@ -0,0 +1,9 @@
+CUDA
+=====
+
+.. currentmodule:: mlx.core.cuda
+
+.. autosummary::
+  :toctree: _autosummary
+
+  is_available
@@ -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 
@@ -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
@@ -13,3 +13,4 @@ Fast
  rope
  scaled_dot_product_attention
  metal_kernel
+  cuda_kernel
@@ -20,5 +20,7 @@ FFT
  irfft2
  rfftn
  irfftn
+  fftfreq
+  rfftfreq
  fftshift
  ifftshift
@@ -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
@@ -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
@@ -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
@@ -27,6 +27,7 @@ simple functions.
   mish
   prelu
   relu
+   relu2
   relu6
   selu
   sigmoid
@@ -10,6 +10,7 @@ Layers
   :template: nn-module-template.rst

   ALiBi
+   AllToShardedLinear
   AvgPool1d
   AvgPool2d
   AvgPool3d
@@ -46,15 +47,19 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
+   QuantizedAllToShardedLinear
   QuantizedEmbedding
   QuantizedLinear
+   QuantizedShardedToAllLinear
   RMSNorm
   ReLU
+   ReLU2
   ReLU6
   RNN
   RoPE
   SELU
   Sequential
+   ShardedToAllLinear
   Sigmoid
   SiLU
   SinusoidalPositionalEncoding
@@ -112,6 +112,7 @@ Operations
   max
   maximum
   mean
+   median
   meshgrid
   min
   minimum
@@ -51,14 +51,14 @@ the saved state. Here's a simple example:
   optimizer.update(model, grads)

   # Save the state
-   state = tree_flatten(optimizer.state)
-   mx.save_safetensors("optimizer.safetensors", dict(state))
+   state = tree_flatten(optimizer.state, destination={})
+   mx.save_safetensors("optimizer.safetensors", state)

   # Later on, for example when loading from a checkpoint,
   # recreate the optimizer and load the state
   optimizer = optim.Adam(learning_rate=1e-2)

-   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
+   state = tree_unflatten(mx.load("optimizer.safetensors"))
   optimizer.state = state

 Note, not every optimizer configuation parameter is saved in the state. For
@@ -0,0 +1,12 @@
+Print Options
+===============
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PrintOptions
+  set_printoptions
+  printoptions
+  get_printoptions
@@ -11,6 +11,7 @@ Transforms
   eval
   async_eval
   compile
+   checkpoint
   custom_function
   disable_compile
   enable_compile
@@ -130,8 +130,8 @@ Now make an array, and benchmark both functions:
 .. code-block:: python

  x = mx.random.uniform(shape=(32, 1000, 4096))
-  timeit(nn.gelu, x)
-  timeit(mx.compile(nn.gelu), x)
+  timeit(gelu, x)
+  timeit(mx.compile(gelu), x)

 On an M1 Max the times are 15.5 and 3.1 milliseconds. The compiled ``gelu`` is
 five times faster.
@@ -225,7 +225,7 @@ In some cases returning updated state can be pretty inconvenient. Hence,
  def fun(x, y):
      z = x + y
      state.append(z)
-      return mx.exp(z), state
+      return mx.exp(z)

  fun(mx.array(1.0), mx.array(2.0))
  # Prints [array(3, dtype=float32)]
@@ -257,7 +257,26 @@ constants. For example:

 In order to have the change of state reflected in the outputs of ``fun`` you
 again have two options. The first option is to simply pass ``state`` as input
-to the function. In some cases this can be pretty inconvenient. Hence,
+to the function.
+
+.. code-block:: python
+
+  state = [mx.array(1.0)]
+
+  @mx.compile
+  def fun(x, state):
+      return x + state[0]
+
+  # Prints array(2, dtype=float32)
+  print(fun(mx.array(1.0), state))
+
+  # Update state
+  state[0] = mx.array(5.0)
+
+  # Prints array(6, dtype=float32)
+  print(fun(mx.array(1.0), state))
+
+In some cases this can be pretty inconvenient. Hence,
 :func:`compile` also has a parameter to capture implicit inputs:

 .. code-block:: python
@@ -7,21 +7,29 @@ Distributed Communication

 MLX supports distributed communication operations that allow the computational cost
 of training or inference to be shared across many physical machines. At the
-moment we support two different communication backends:
+moment we support several different communication backends introduced below.
+
+.. list-table::
+   :widths: 20 80
+   :header-rows: 1
+
+   * - Backend
+     - Description
+   * - :ref:`MPI <mpi_section>`
+     - A full featured and mature distributed communications library.
+   * - :ref:`RING <ring_section>`
+     - Ring all reduce and all gather over TCP sockets. Always available and
+       usually faster than MPI.
+   * - :ref:`JACCL <jaccl_section>`
+     - Low latency communication with RDMA over thunderbolt. Necessary for
+       things like tensor parallelism.
+   * - :ref:`NCCL <nccl_section>`
+     - The backend of choice for CUDA environments.

-* `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
-  full-featured and mature distributed communications library
-* A **ring** backend of our own that uses native TCP sockets and should be
-  faster for thunderbolt connections.

 The list of all currently supported operations and their documentation can be
 seen in the :ref:`API docs<distributed>`.

-.. note::
-   Some operations may not be supported or not as fast as they should be.
-   We are adding more and tuning the ones we have as we are figuring out the
-   best way to do distributed computing on Macs using MLX.
-
 Getting Started
 ---------------

@@ -84,9 +92,8 @@ Selecting Backend
 ^^^^^^^^^^^^^^^^^

 You can select the backend you want to use when calling :func:`init` by passing
-one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to
-initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they
-both fail then a singleton group is created.
+one of ``{'any', 'ring', 'jaccl', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all
+available backends. If they all fail then a singleton group is created.

 .. note::
   After a distributed backend is successfully initialized :func:`init` will
@@ -110,162 +117,13 @@ 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
+----------------------------

-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.
+- :ref:`Data Parallelism <data_parallelism>`
+- :ref:`Tensor Parallelism <tensor_parallelism>`

-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 = mlx.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())
-
-
-Getting Started with MPI
------------------------
-
-MLX already comes with the ability to "talk" to MPI if it is installed on the
-machine. Launching distributed MLX programs that use MPI can be done with
-``mpirun`` as expected. However, in the following examples we will be using
-``mlx.launch --backend mpi`` which takes care of some nuisances such as setting
-absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` shared
-library.
-
-The simplest possible usage is the following which, assuming the minimal
-example in the beginning of this page, should result in:
-
-.. code:: shell
-
-    $ mlx.launch --backend mpi -n 2 test.py
-    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-
-The above launches two processes on the same (local) machine and we can see
-both standard output streams. The processes send the array of 1s to each other
-and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would
-print 4 etc.
-
-Installing MPI
-^^^^^^^^^^^^^^
-
-MPI can be installed with Homebrew, using the Anaconda package manager or
-compiled from source. Most of our testing is done using ``openmpi`` installed
-with the Anaconda package manager as follows:
-
-.. code:: shell
-
-    $ conda install conda-forge::openmpi
-
-Installing with Homebrew may require specifying the location of ``libmpi.dyld``
-so that MLX can find it and load it at runtime. This can simply be achieved by
-passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
-done automatically by ``mlx.launch``.
-
-.. code:: shell
-
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
-    $ # or simply
-    $ mlx.launch -n 2 test.py
-
-Setting up Remote Hosts
-^^^^^^^^^^^^^^^^^^^^^^^
-
-MPI can automatically connect to remote hosts and set up the communication over
-the network if the remote hosts can be accessed via ssh. A good checklist to
-debug connectivity issues is the following:
-
-* ``ssh hostname`` works from all machines to all machines without asking for
-  password or host confirmation
-* ``mpirun`` is accessible on all machines.
-* Ensure that the ``hostname`` used by MPI is the one that you have configured
-  in the ``.ssh/config`` files on all machines.
-
-Tuning MPI All Reduce
-^^^^^^^^^^^^^^^^^^^^^
-
-.. note::
-
-    For faster all reduce consider using the ring backend either with Thunderbolt
-    connections or over Ethernet.
-
-Configure MPI to use N tcp connections between each host to improve bandwidth
-by passing ``--mca btl_tcp_links N``.
-
-Force MPI to use the most performant network interface by setting ``--mca
-btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want
-to use.
+.. _ring_section:

 Getting Started with Ring
 -------------------------
@@ -275,7 +133,7 @@ available. It uses TCP sockets so the nodes need to be reachable via a network.
 As the name suggests the nodes are connected in a ring which means that rank 1
 can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3
 and so on and so forth. As a result :func:`send` and :func:`recv` with
-arbitrary sender and receiver is not supported in the ring backend.
+arbitrary sender and receiver are not supported in the ring backend.

 Defining a Ring
 ^^^^^^^^^^^^^^^
@@ -316,22 +174,13 @@ utility as follows:

 .. code:: shell

-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4
+   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --backend ring

 By default the script will attempt to discover the thunderbolt ring and provide
 you with the commands to configure each node as well as the ``hostfile.json``
 to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes
 then ``--auto-setup`` can be used to configure them automatically.

-To validate your connection without configuring anything
-``mlx.distributed_config`` can also plot the ring using DOT format.
-
-.. code:: shell
-
-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot
-   dot -Tpng ring.dot >ring.png
-   open ring.png
-
 If you want to go through the process manually, the steps are as follows:

 * Disable the thunderbolt bridge interface
@@ -342,3 +191,382 @@ If you want to go through the process manually, the steps are as follows:
  and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and
  ``192.168.0.2`` respectively to the two nodes. For more details you can see
  the commands prepared by the utility script.
+
+.. _jaccl_section:
+
+Getting Started with JACCL
+--------------------------
+
+Starting from macOS 26.2, RDMA over thunderbolt is available and
+enables low-latency communication between Macs with thunderbolt 5. MLX provides
+the JACCL backend that uses this functionality to achieve communication latency
+an order of magnitude lower than the ring backend.
+
+.. note::
+
+   The name JACCL (pronounced Jackal) stands for *Jack and Angelos' Collective
+   Communication Library* and it is an obvious pun to Nvidia's NCCL but also
+   tribute to *Jack Beasley* who led the development of RDMA over Thunderbolt
+   at Apple.
+
+Enabling RDMA
+^^^^^^^^^^^^^
+
+Until the feature matures, enabling RDMA over thunderbolt is slightly more
+involved and **cannot** be done remotely even with sudo. In fact, it has to be
+done in macOS recovery:
+
+1. `Start your computer in recovery <https://support.apple.com/en-us/102518>`_.
+2. Open the Terminal by going to Utilities -> Terminal.
+3. Run ``rdma_ctl enable``.
+4. Reboot.
+
+To verify that you have successfully enabled Thunderbolt RDMA you can run
+``ibv_devices`` which should produce something like the following for an M3 Ultra.
+
+.. code-block:: bash
+
+    ~ % ibv_devices
+    device          	   node GUID
+    ------          	----------------
+    rdma_en2        	8096a9d9edbaac05
+    rdma_en3        	8196a9d9edbaac05
+    rdma_en5        	8396a9d9edbaac05
+    rdma_en4        	8296a9d9edbaac05
+    rdma_en6        	8496a9d9edbaac05
+    rdma_en7        	8596a9d9edbaac05
+
+Defining a Mesh
+^^^^^^^^^^^^^^^
+
+The JACCL backend supports only fully connected topologies. Namely, there needs
+to be a thunderbolt cable connecting all pairs of Macs directly. For example, in
+the following topology visualizations, the left one is valid because there is a
+connection from any node to any other node, while for the one on the right M3
+Ultra 1 is not connected to M3 Ultra 2.
+
+.. raw:: html
+
+   <div style="display: flex; text-align: center; align-items: end; font-size: 80%;">
+     <div>
+       <img src="../_static/distributed/m3-ultra-mesh.png" alt="M3 Ultra thunderbolt mesh" style="width: 55%">
+       <p>Fully connected mesh of four M3 Ultra.</p>
+     </div>
+     <div>
+       <img src="../_static/distributed/m3-ultra-mesh-broken.png" alt="M3 Ultra broken thunderbolt mesh" style="width: 55%">
+       <p>Not a valid mesh (M3 Ultra 1 is not connected to M3 Ultra 2).</p>
+     </div>
+   </div>
+
+Similar to the ring backend, the easiest way to use JACCL with MLX is to write
+a JSON hostfile that will be used by ``mlx.launch``. The hostfile needs to contain
+
+- Hostnames to use for launching scripts via ssh
+- An IP for rank 0 that is reachable by all nodes
+- A list of rdma devices that connect each node to each other node
+
+The following JSON defines the valid 4-node mesh from the image above.
+
+.. code-block:: json
+
+    [
+        {
+            "ssh": "m3-ultra-1",
+            "ips": ["123.123.123.1"],
+            "rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"]
+        },
+        {
+            "ssh": "m3-ultra-2",
+            "ips": [],
+            "rdma": ["rdma_en5", null, "rdma_en3", "rdma_en4"]
+        },
+        {
+            "ssh": "m3-ultra-3",
+            "ips": [],
+            "rdma": ["rdma_en4", "rdma_en3", null, "rdma_en5"]
+        },
+        {
+            "ssh": "m3-ultra-4",
+            "ips": [],
+            "rdma": ["rdma_en3", "rdma_en4", "rdma_en5", null]
+        }
+    ]
+
+Even though TCP/IP is not used when communicating with Thunderbolt RDMA,
+disabling the thunderbolt bridge is still required as well as setting up
+isolated local networks for each thunderbolt connection.
+
+All of the above can be done instead via ``mlx.distributed_config``. This helper
+script will
+
+- ssh into each node
+- extract the thunderbolt connectivity
+- check for a valid mesh
+- provide the commands to configure each node (or run them if sudo is available)
+- generate the hostfile to be used with ``mlx.launch``
+
+Putting It All Together
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+For example launching a distributed MLX script that uses JACCL is fairly simple
+if the nodes are reachable via ssh and have password-less sudo.
+
+First, connect all the thunderbolt cables. Then we can verify the connections
+by using the ``mlx.distributed_config`` script to visualize them.
+
+.. code-block::
+
+   mlx.distributed_config --verbose \
+        --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
+        --over thunderbolt --dot | dot -Tpng | open -f -a Preview
+
+After making sure that everything looks right we can auto-configure the nodes
+and save the hostfile to ``m3-ultra-jaccl.json`` by running:
+
+.. code-block::
+
+   mlx.distributed_config --verbose \
+        --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
+        --over thunderbolt --backend jaccl \
+        --auto-setup --output m3-ultra-jaccl.json
+
+And now we are ready to run a distributed MLX script such as distributed inference
+of a gigantic model using MLX LM.
+
+.. code-block::
+
+   mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json \
+        --env MLX_METAL_FAST_SYNCH=1 -- \  # <--- important
+        /path/to/remote/python -m mlx_lm chat --model mlx-community/DeepSeek-R1-0528-4bit
+
+.. note::
+
+   Defining the environment variable ``MLX_METAL_FAST_SYNCH=1`` enables a
+   different, faster way of synchronizing between the GPU and the CPU. It is
+   not specific to the JACCL backend and can be used in all cases where the CPU
+   and GPU need to collaborate for some computation and is pretty critical for
+   low-latency communication since the communication is done by the CPU.
+
+.. _nccl_section:
+
+Getting Started with NCCL
+-------------------------
+
+MLX on CUDA environments ships with the ability to talk to `NCCL
+<https://developer.nvidia.com/nccl>`_ which is a high-performance collective
+communication library that supports both multi-gpu and multi-node setups.
+
+For CUDA environments, NCCL is the default backend for ``mlx.launch`` and all
+it takes to run a distributed job is
+
+.. code-block::
+
+   mlx.launch -n 8 test.py
+
+   # perfect for interactive scripts
+   mlx.launch -n 8 python -m mlx_lm chat --model my-model
+
+You can also use ``mlx.launch`` to ssh to a remote node and launch a script
+with the same ease
+
+.. code-block::
+
+   mlx.launch --hosts my-cuda-node -n 8 test.py
+
+In many cases you may not want to use ``mlx.launch`` with the NCCL backend
+because the cluster scheduler will be the one launching the processes. You can
+:ref:`see which environment variables need to be defined <no_mlx_launch>` in
+order for the MLX NCCL backend to be initialized correctly.
+
+.. _mpi_section:
+
+Getting Started with MPI
+------------------------
+
+MLX already comes with the ability to "talk" to `MPI
+<https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ if it is installed
+on the machine. Launching distributed MLX programs that use MPI can be done
+with ``mpirun`` as expected. However, in the following examples we will be
+using ``mlx.launch --backend mpi`` which takes care of some nuisances such as
+setting absolute paths for the ``mpirun`` executable and the ``libmpi.dyld``
+shared library.
+
+The simplest possible usage is the following which, assuming the minimal
+example in the beginning of this page, should result in:
+
+.. code:: shell
+
+    $ mlx.launch --backend mpi -n 2 test.py
+    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+
+The above launches two processes on the same (local) machine and we can see
+both standard output streams. The processes send the array of 1s to each other
+and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would
+print 4 etc.
+
+Installing MPI
+^^^^^^^^^^^^^^
+
+MPI can be installed with Homebrew, pip, using the Anaconda package manager, or
+compiled from source. Most of our testing is done using ``openmpi`` installed
+with the Anaconda package manager as follows:
+
+.. code:: shell
+
+    $ conda install conda-forge::openmpi
+
+Installing with Homebrew or pip requires specifying the location of ``libmpi.dyld``
+so that MLX can find it and load it at runtime. This can simply be achieved by
+passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
+done automatically by ``mlx.launch``. Some environments use a non-standard
+library filename that can be specified using the ``MPI_LIBNAME`` environment
+variable. This is automatically taken care of by ``mlx.launch`` as well.
+
+.. code:: shell
+
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
+    $ # or simply
+    $ mlx.launch -n 2 test.py
+
+Setting up Remote Hosts
+^^^^^^^^^^^^^^^^^^^^^^^
+
+MPI can automatically connect to remote hosts and set up the communication over
+the network if the remote hosts can be accessed via ssh. A good checklist to
+debug connectivity issues is the following:
+
+* ``ssh hostname`` works from all machines to all machines without asking for
+  password or host confirmation
+* ``mpirun`` is accessible on all machines.
+* Ensure that the ``hostname`` used by MPI is the one that you have configured
+  in the ``.ssh/config`` files on all machines.
+
+Tuning MPI All Reduce
+^^^^^^^^^^^^^^^^^^^^^
+
+.. note::
+
+    For faster all reduce consider using the ring backend either with Thunderbolt
+    connections or over Ethernet.
+
+Configure MPI to use N tcp connections between each host to improve bandwidth
+by passing ``--mca btl_tcp_links N``.
+
+Force MPI to use the most performant network interface by setting ``--mca
+btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want
+to use.
+
+.. _no_mlx_launch:
+
+Distributed Without ``mlx.launch``
+----------------------------------
+
+None of the implementations of the distributed backends require launching with
+``mlx.launch``. The script simply connects to each host. Starts a process per
+rank and sets up the necessary environment variables before delegating to your
+MLX script. See the :doc:`dedicated documentation page <launching_distributed>`
+for more details.
+
+For many use-cases this will be the easiest way to perform distributed
+computations in MLX. However, there may be reasons that you cannot or should
+not use ``mlx.launch``. A common such case is the use of a scheduler that
+starts all the processes for you on machines undetermined at the time of
+scheduling the job.
+
+Below we list the environment variables required to use each backend.
+
+Ring
+^^^^^^
+
+**MLX_RANK** should contain a single 0-based integer that defines the rank of
+the process.
+
+**MLX_HOSTFILE** should contain the path to a json file that contains IPs and
+ports for each rank to listen to, something like the following:
+
+.. code-block:: json
+
+   [
+     ["123.123.1.1:5000", "123.123.1.2:5000"],
+     ["123.123.2.1:5000", "123.123.2.2:5000"],
+     ["123.123.3.1:5000", "123.123.3.2:5000"],
+     ["123.123.4.1:5000", "123.123.4.2:5000"]
+   ]
+
+**MLX_RING_VERBOSE** is optional and if set to 1 it enables some more logging
+from the distributed backend.
+
+JACCL
+^^^^^
+
+**MLX_RANK** should contain a single 0-based integer that defines the rank of
+the process.
+
+**MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen
+to all the other ranks connect to in order to establish the RDMA connections.
+
+**MLX_IBV_DEVICES** should contain the path to a json file that contains the
+ibverbs device names that connect each node to each other node, something like
+the following:
+
+.. code-block:: json
+
+   [
+      [null, "rdma_en5", "rdma_en4", "rdma_en3"],
+      ["rdma_en5", null, "rdma_en3", "rdma_en4"],
+      ["rdma_en4", "rdma_en3", null, "rdma_en5"],
+      ["rdma_en3", "rdma_en4", "rdma_en5", null]
+   ]
+
+
+NCCL
+^^^^^
+
+**MLX_RANK** should contain a single 0-based integer that defines the rank of
+the process.
+
+**MLX_WORLD_SIZE** should contain the total number of processes that will be
+launched.
+
+**NCCL_HOST_IP** and **NCCL_PORT** should contain the IP and port that all
+hosts can connect to to establish the NCCL communication.
+
+**CUDA_VISIBLE_DEVICES** should contain the local index of the gpu that
+corresponds to this process.
+
+Of course any `other environment variable
+<https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html>`_ that is
+used by NCCL can be set.
+
+.. _tips_and_tricks:
+
+Tips and Tricks
+----------------
+
+This is a small collection of tips to help you utilize better the distributed
+communication capabilities of MLX.
+
+- *Test locally first.*
+
+  You can use the pattern ``mlx.launch -n2 -- my_script.py`` to run a small
+  scale test on a single node first.
+
+- *Batch your communication.*
+
+  As described in the :ref:`training example <training_example>`, performing a
+  lot of small communications can hurt performance. Copy the approach of
+  :func:`mlx.nn.average_gradients` to gather many small communications in a
+  single large one.
+
+- *Visualize the connectivity.*
+
+  Use ``mlx.distributed_config --hosts h1,h2,h3 --over thunderbolt --dot`` to
+  visualize the connnections and make sure that the cables are connected
+  correctly. See the :ref:`JACCL section <jaccl_section>` for examples.
+
+- *Use the debugger.*
+
+  ``mlx.launch`` is meant for interactive use. It broadcasts stdin to all
+  processes and gathers stdout from all processes. This makes using ``pdb`` a
+  breeze.
@@ -7,17 +7,17 @@ Exporting Functions

 MLX has an API to export and import functions to and from a file. This lets you
 run computations written in one MLX front-end (e.g. Python) in another MLX
-front-end (e.g. C++). 
+front-end (e.g. C++).

 This guide walks through the basics of the MLX export API with some examples.
 To see the full list of functions check-out the :ref:`API documentation
 <export>`.

-Basics of Exporting 
+Basics of Exporting
 -------------------

 Let's start with a simple example:
- 
+
 .. code-block:: python

  def fun(x, y):
@@ -67,7 +67,7 @@ specified as variable positional arguments or as a tuple of arrays:

  x = mx.array(1.0)
  y = mx.array(1.0)
-   
+
  # Both arguments to fun are positional
  mx.export_function("add.mlxfn", fun, x, y)

@@ -133,7 +133,7 @@ parameters are also saved to the ``model.mlxfn`` file.
   For enclosed arrays inside an exported function, be extra careful to ensure
   they are evaluated. The computation graph that gets exported will include
   the computation that produces enclosed inputs.
-  
+
   If the above example was missing ``mx.eval(model.parameters()``, the
   exported function would include the random initialization of the
   :obj:`mlx.nn.Module` parameters.
@@ -150,11 +150,39 @@ parameters, pass them as inputs to the ``call`` wrapper:
     # Set the model's parameters to the input parameters
     model.update(tree_unflatten(list(params.items())))
     return model(x)
- 
-   params = dict(tree_flatten(model.parameters()))
+
+   params = tree_flatten(model.parameters(), destination={})
   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
 -----------------

@@ -164,13 +192,13 @@ to export a function which can be used for inputs with variable shapes:

 .. code-block:: python

-  mx.export_function("fun.mlxfn", mx.abs, mx.array(0.0), shapeless=True)
+  mx.export_function("fun.mlxfn", mx.abs, mx.array([0.0]), shapeless=True)
  imported_abs = mx.import_function("fun.mlxfn")

  # Ok
-  out, = imported_abs(mx.array(-1.0))
-  
-  # Also ok 
+  out, = imported_abs(mx.array([-1.0]))
+
+  # Also ok
  out, = imported_abs(mx.array([-1.0, -2.0]))

 With ``shapeless=False`` (which is the default), the second call to
@@ -197,7 +225,7 @@ a single file by creating an exporting context manager with :func:`exporter`:
  def fun(x, y=None):
      constant = mx.array(3.0)
      if y is not None:
-        x += y 
+        x += y
      return x + constant

  with mx.exporter("fun.mlxfn", fun) as exporter:
@@ -215,7 +243,7 @@ a single file by creating an exporting context manager with :func:`exporter`:
  print(out)

 In the above example the function constant data, (i.e. ``constant``), is only
-saved once. 
+saved once.

 Transformations with Imported Functions
 ---------------------------------------
@@ -238,7 +266,7 @@ on imported functions just like regular Python functions:
  # Prints: array(1, dtype=float32)
  print(dfdx(x))

-  # Compile the imported function 
+  # Compile the imported function
  mx.compile(imported_fun)
  # Prints: array(0, dtype=float32)
  print(compiled_fun(x)[0])
@@ -275,7 +303,7 @@ Import and run the function in C++ with only a few lines of code:
  // Prints: array(2, dtype=float32)
  std::cout << outputs[0] << std::endl;

-Imported functions can be transformed in C++ just like in Python. Use 
+Imported functions can be transformed in C++ just like in Python. Use
 ``std::vector<mx::array>`` for positional arguments and ``std::map<std::string,
 mx::array>`` for keyword arguments when calling imported functions in C++.

@@ -70,7 +70,8 @@ Differences from NumPy

  * Indexing does not perform bounds checking. Indexing out of bounds is
    undefined behavior.
-  * Boolean mask based indexing is not yet supported.
+  * Boolean mask based indexing is supported for assignment only (see
+    :ref:`boolean-mask-assignment`).

 The reason for the lack of bounds checking is that exceptions cannot propagate
 from the GPU. Performing bounds checking for array indices before launching the
@@ -107,8 +108,20 @@ same array:
  >>> a
  array([1, 2, 0], dtype=int32)

+Note that unlike NumPy, slicing an array creates a copy, not a view. So
+mutating it does not mutate the original array:

-Note, unlike NumPy, updates to the same location are nondeterministic:
+.. code-block:: shell
+
+  >>> a = mx.array([1, 2, 3])
+  >>> b = a[:]
+  >>> b[2] = 0
+  >>> b
+  array([1, 2, 0], dtype=int32)
+  >>> a
+  array([1, 2, 3], dtype=int32)
+
+Also unlike NumPy, updates to the same location are nondeterministic:

 .. code-block:: shell

@@ -131,3 +144,51 @@ expected. For example:

 In the above ``dfdx`` will have the correct gradient, namely zeros at ``idx``
 and ones elsewhere.
+
+.. _boolean-mask-assignment:
+
+Boolean Mask Assignment
+-----------------------
+
+MLX supports boolean indices using NumPy syntax. A mask must already be
+a :class:`bool_` MLX :class:`array` or a NumPy ``ndarray`` with ``dtype=bool``.
+Other index types are routed through the standard scatter code.
+
+.. code-block:: shell
+
+   >>> a = mx.array([1.0, 2.0, 3.0])
+   >>> mask = mx.array([True, False, True])
+   >>> updates = mx.array([5.0, 6.0])
+   >>> a[mask] = updates
+   >>> a
+   array([5.0, 2.0, 6.0], dtype=float32)
+
+Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar
+assignments, ``updates`` must provide at least as many elements as there are
+``True`` entries in ``mask``.
+
+.. code-block:: shell
+
+   >>> a = mx.zeros((2, 3))
+   >>> mask = mx.array([[True, False, True],
+                        [False, False, True]])
+   >>> a[mask] = 1.0
+   >>> a
+   array([[1.0, 0.0, 1.0],
+          [0.0, 0.0, 1.0]], dtype=float32)
+
+Boolean masks follow NumPy semantics:
+
+- The mask shape must match the shape of the axes it indexes exactly. The only
+  exception is a scalar boolean mask, which broadcasts to the full array.
+- Any axes not covered by the mask are taken in full.
+
+.. code-block:: shell
+
+   >>> a = mx.arange(1000).reshape(10, 10, 10)
+   >>> a[mx.random.normal((10, 10)) > 0.0] = 0  # valid: mask covers axes 0 and 1
+
+The mask of shape ``(10, 10)`` applies to the first two axes, so ``a[mask]``
+selects the 1-D slices ``a[i, j, :]`` where ``mask[i, j]`` is ``True``.
+Shapes such as ``(1, 10, 10)`` or ``(10, 10, 1)`` do not match the indexed
+axes and therefore raise errors.
@@ -7,13 +7,106 @@ Launching Distributed Programs

 .. currentmodule:: mlx.core.distributed

-Installing the MLX python package provides a helper script ``mlx.launch`` that
-can be used to run python scripts distributed on several nodes. It allows
-launching using either the MPI backend or the ring backend. See the
-:doc:`distributed docs <distributed>` for the different backends.
+The MLX python package provides two utilities to help you configure
+your Macs for distributed computation and also launch distributed programs on
+multiple nodes or with many processes in a single node. These utilities are aptly named

-Usage
-----
+- ``mlx.launch``
+- ``mlx.distributed_config``
+
+See the :doc:`distributed docs <distributed>` for an introduction and
+getting-started guides to the various backends.
+
+``mlx.distributed_config`` 
+---------------------------
+
+Unless you are launching distributed jobs locally for development or multi-gpu
+CUDA environments, then you have several Macs that you need to configure for
+distributed communication with MLX.
+
+``mlx.distributed_config`` aims to automate the process of configuring the
+network interfaces (especially for communication over thunderbolt) and also
+creating the hostfile to be used with ``mlx.launch``.
+
+We will analyse 3 cases of using ``mlx.distributed_config``
+
+1. RDMA over thunderbolt using JACCL
+2. TCP/IP over thunderbolt using the ring backend
+3. TCP/IP over ethernet using the ring backend
+
+JACCL
+^^^^^^^
+
+After following :ref:`the steps to enable RDMA <jaccl_section>` you can run the
+following command to configure the nodes and create the hostfile.
+
+.. code-block::
+
+   mlx.distributed_config --verbose --backend jaccl \
+        --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 --over thunderbolt \
+        --auto-setup --output m3-ultra-jaccl.json
+
+Let's walk through the steps that the script takes to configure the nodes.
+
+1. ssh to all nodes to verify that they are reachable
+2. Extract the thunderbolt connectivity. Namely run commands on each node to
+   calculate which node is connected to which other node.
+3. Verify that we have a valid fully connected mesh
+4. Check that RDMA is enabled
+5. Extract the ethernet IP from interface en0
+6. Disable the thunderbolt bridge and set up peer to peer networks for each
+   thunderbolt cable
+7. Write the hostfile
+
+Knowing the above steps allows you to manually configure the nodes but also
+debug any configuration issue. For instance changing the Ethernet IP to a
+different interface directly in the config is possible (as long as it is
+reachable from all nodes).
+
+The ``--auto-setup`` argument requires password-less sudo on each node. If it
+isn't available then the configuration script will print commands to be run on
+each node.
+
+Ring over thunderbolt
+^^^^^^^^^^^^^^^^^^^^^
+
+Setting up a ring backend over thunderbolt only requires changing the
+``--backend`` from ``jaccl`` to ``ring``.
+
+The steps are very similar with the main difference being that instead of
+verifying that the nodes are fully connected, the script attempts to identify a
+ring topology (or multiple rings).
+
+Ring over Ethernet
+^^^^^^^^^^^^^^^^^^
+
+Configuring the ring backend over ethernet doesn't require setting up network
+interface and as such it simply extracts the ``en0`` IP from each node and
+writes the hostfile.
+
+Debugging cable connections
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``mlx.distributed_config`` can help you debug the connectivity of your nodes
+over thunderbolt by exporting a graph of the connections.
+
+Running
+
+.. code-block::
+
+   mlx.distributed_config --verbose \
+        --hosts host1,host2,host3,host4 \
+        --over thunderbolt --dot
+
+will export a `GraphViz <https://graphviz.org>`_ representation of the
+connections between the nodes which makes it very easy to figure out which
+cable is not connected correctly.
+
+See :ref:`the JACCL section <jaccl_section>` for an example.
+
+
+``mlx.launch``
+--------------

 The minimal usage example of ``mlx.launch`` is simply

@@ -33,6 +126,10 @@ the rest if one of them fails unexpectedly or if ``mlx.launch`` is terminated.
 It also takes care of forwarding the output of each remote process to stdout
 and stderr respectively.

+Importantly, it also broadcasts stdin to each process which enables interactive
+programs to work in distributed mode as well as debugging using the interactive
+debugger.
+
 Providing Hosts
 ^^^^^^^^^^^^^^^^

@@ -63,10 +160,62 @@ host and on the same path. A good checklist to debug errors is the following:
  ``mlx.launch --print-python`` to see what that path is.
 * the script you want to run is available on all hosts at the same path

+If you are launching from a node with a completely different setup than the
+nodes that the program will run on, you can specify ``--no-verify-script`` so
+that ``mlx.launch`` does not attempt to verify that the executable and script
+exist locally before launching the distributed job.
+
+.. _ring_specifics:
+
+Ring Specifics
+^^^^^^^^^^^^^^
+
+The :ref:`ring <ring_section>` backend, which is also the default
+backend, can be explicitly selected with the argument ``--backend ring``. The
+ring backend has some specific requirements and arguments that are different to
+other backends:
+
+* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to
+  ssh to a hostname that does not correspond to the IP we want to bind to we
+  have to provide a hostfile.
+* ``--starting-port`` defines the port to bind to on the remote hosts.
+  Specifically rank 0 for the first IP will use this port and each subsequent
+  IP or rank will add 1 to this port.
+* ``--connections-per-ip`` allows us to increase the number of connections
+  between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for
+  ``mpirun``.
+
+.. _jaccl_specifics:
+
+JACCL Specifics
+^^^^^^^^^^^^^^^^
+
+The :ref:`JACCL <jaccl_section>` backend can be selected with the argument
+``--backend jaccl``. A hostfile is necessary to launch with this backend
+because it needs to contain the RDMA devices connecting each node to each other
+node.
+
+NCCL Specifics
+^^^^^^^^^^^^^^
+
+The :ref:`NCCL <nccl_section>` backend is the default backend for CUDA
+environments. When launching from a Mac to a Linux machine with CUDA then the
+backend should be selected using ``--backend nccl``.
+
+The ``--repeat-hosts, -n`` argument should be used to launch multi-node and
+multi-gpu jobs. For instance
+
+.. code-block::
+
+   mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 --no-verify-script -- ./my-job.sh
+
+will attempt to launch 16 processes, 8 on each node that will all run
+``my-job.sh``.
+
 .. _mpi_specifics:

 MPI Specifics
-------------
+^^^^^^^^^^^^^

 One can use MPI by passing ``--backend mpi`` to ``mlx.launch``. In that case,
 ``mlx.launch`` is a thin wrapper over ``mpirun``. Moreover,
@@ -83,23 +232,3 @@ to choose a specific interface for the byte-transfer-layer of MPI we can call
 .. code:: shell

    mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py
-
-
-.. _ring_specifics:
-
-Ring Specifics
--------------
-
-The ring backend, which is also the default backend, can be explicitly selected
-with the argument ``--backend ring``. The ring backend has some specific
-requirements and arguments that are different to MPI:
-
-* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to
-  ssh to a hostname that does not correspond to the IP we want to bind to we
-  have to provide a hostfile.
-* ``--starting-port`` defines the port to bind to on the remote hosts.
-  Specifically rank 0 for the first IP will use this port and each subsequent
-  IP or rank will add 1 to this port.
-* ``--connections-per-ip`` allows us to increase the number of connections
-  between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for
-  ``mpirun``.
@@ -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
 ---
@@ -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
@@ -3,6 +3,6 @@ requires = [
  "setuptools>=42",
  "cmake>=3.25",
  "mlx>=0.18.0",
-  "nanobind==2.4.0",
+  "nanobind==2.12.0",
 ]
 build-backend = "setuptools.build_meta"
@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
-mlx>=0.21.0
-nanobind==2.4.0
+mlx>=0.31.2
+nanobind==2.12.0
@@ -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
@@ -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
@@ -0,0 +1,117 @@
+from itertools import product
+
+import mlx.core as mx
+
+
+# In mxfp8 mode, the results do not match exactly:
+# fewer than 1% of output elements differ.
+# This does not appear to be a systematic error.
+# The error can exceed 1 ULP for very small values,
+# and is always below 1 ULP for larger values.
+# For nvfp4, the results match exactly.
+# therefore I suspect that the discrepancy comes from
+# the mxfp8 matmul implementation in cuBLASLt..
+def ulp_bf16_at(x):
+    ax = mx.abs(x)
+    min_normal = mx.array(2.0**-126)
+    ax = mx.where(ax < min_normal, min_normal, ax)
+    e = mx.floor(mx.log2(ax))
+    return mx.power(2.0, e - 7.0)
+
+
+def test_qqmm():
+    key = mx.random.key(0)
+    k1, k2 = mx.random.split(key)
+    dtypes = [mx.bfloat16, mx.float32, mx.float16]
+
+    tests = (
+        (16, "nvfp4", 4),
+        (32, "mxfp8", 8),
+    )
+    shapes = (
+        [64, 65, 33, 128, 256, 1024, 1024 * 8],  # M
+        [64, 128, 256, 1024, 1024 * 8],  # N
+        [64, 128, 256, 1024, 1024 * 8],  # K
+    )
+    for group_size, mode, bits in tests:
+        for M, N, K in product(*shapes):
+            for dtype in dtypes:
+                x = mx.random.normal(shape=(M, K), key=k1, dtype=dtype)
+                w = mx.random.normal(shape=(N, K), key=k2, dtype=dtype)
+                w_q, scales_w = mx.quantize(w, group_size, bits, mode=mode)
+                w_dq = mx.dequantize(
+                    w_q,
+                    scales_w,
+                    group_size=group_size,
+                    bits=bits,
+                    mode=mode,
+                    dtype=dtype,
+                )
+                y_q = mx.qqmm(
+                    x,
+                    w_q,
+                    scales_w,
+                    group_size=group_size,
+                    bits=bits,
+                    mode=mode,
+                )
+                x_q, scales_x = mx.quantize(
+                    x, group_size=group_size, bits=bits, mode=mode
+                )
+                x_dq = mx.dequantize(
+                    x_q,
+                    scales_x,
+                    group_size=group_size,
+                    bits=bits,
+                    mode=mode,
+                    dtype=dtype,
+                )
+                y_hat = mx.matmul(x_dq, mx.transpose(w_dq))
+                ulp = ulp_bf16_at(y_hat)
+                error = (y_q - y_hat).abs()
+                if not (mx.logical_or(error < 1e-3, error <= ulp).all()):
+                    raise AssertionError(
+                        f"qqmm test failed for shape {(M, N, K)}, "
+                        f"group_size={group_size}, bits={bits}, "
+                        f"mode={mode}, dtype={dtype}"
+                    )
+
+
+def test_qqmm_vjp():
+    key = mx.random.key(0)
+    k1, k2 = mx.random.split(key)
+    M = 64
+    N = 1024
+    K = 512
+    tests = (
+        (16, "nvfp4", 4),
+        (32, "mxfp8", 8),
+    )
+    x = mx.random.normal(shape=(M, K), key=k1)
+    c = mx.ones(shape=(M, N))
+
+    for group_size, mode, bits in tests:
+        w = mx.random.normal(shape=(N, K), key=k2)
+
+        def fn(x):
+            return mx.qqmm(x, w, group_size=group_size, bits=bits, mode=mode)
+
+        _, vjp_out = mx.vjp(fn, primals=(x,), cotangents=(c,))
+        w_tq, scales_wt = mx.quantize(
+            mx.transpose(w), group_size=group_size, bits=bits, mode=mode
+        )
+        expected_out = mx.qqmm(
+            c, w_tq, scales_wt, group_size=group_size, bits=bits, mode=mode
+        )
+        ulp = ulp_bf16_at(expected_out)
+        error = (vjp_out[0] - expected_out).abs()
+        if not (mx.logical_or(error < 1e-3, error <= ulp).all()):
+            raise AssertionError(
+                f"qqmm vjp test failed for shape {(M, N, K)}, "
+                f"group_size={group_size}, bits={bits}, mode={mode}"
+            )
+
+
+if __name__ == "__main__":
+    test_qqmm()
+    test_qqmm_vjp()
@@ -1,7 +1,6 @@
 target_sources(
  mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
@@ -15,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
@@ -23,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)
@@ -1,24 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-#include <cstdlib>
-#include <sstream>
-
-#include "mlx/allocator.h"
-
-namespace mlx::core::allocator {
-
-Buffer malloc(size_t size) {
-  auto buffer = allocator().malloc(size);
-  if (size && !buffer.ptr()) {
-    std::ostringstream msg;
-    msg << "[malloc] Unable to allocate " << size << " bytes.";
-    throw std::runtime_error(msg.str());
-  }
-  return buffer;
-}
-
-void free(Buffer buffer) {
-  allocator().free(buffer);
-}
-
-} // namespace mlx::core::allocator
@@ -4,17 +4,19 @@

 #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_;

 public:
-  Buffer(void* ptr) : ptr_(ptr) {};
+  explicit Buffer(void* ptr) : ptr_(ptr) {};

  // Get the raw data pointer from the buffer
  void* raw_ptr();
@@ -28,16 +30,16 @@ class Buffer {
  };
 };

-Buffer malloc(size_t size);
-
-void free(Buffer buffer);
-
-class Allocator {
+class MLX_API Allocator {
  /** Abstract base class for a memory allocator. */
 public:
  virtual Buffer malloc(size_t size) = 0;
  virtual void free(Buffer buffer) = 0;
  virtual size_t size(Buffer buffer) const = 0;
+  virtual Buffer make_buffer(void* ptr, size_t size) {
+    return Buffer{nullptr};
+  };
+  virtual void release(Buffer buffer) {}

  Allocator() = default;
  Allocator(const Allocator& other) = delete;
@@ -47,6 +49,27 @@ class Allocator {
  virtual ~Allocator() = default;
 };

-Allocator& allocator();
+MLX_API Allocator& allocator();
+
+inline Buffer malloc(size_t size) {
+  return allocator().malloc(size);
+}
+
+inline void free(Buffer buffer) {
+  allocator().free(buffer);
+}
+
+// Make a Buffer from a raw pointer of the given size without a copy.  If a
+// no-copy conversion is not possible then the returned buffer.ptr() will be
+// nullptr. Any buffer created with this function must be released with
+// release(buffer)
+inline Buffer make_buffer(void* ptr, size_t size) {
+  return allocator().make_buffer(ptr, size);
+};
+
+// Release a buffer from the allocator made with make_buffer
+inline void release(Buffer buffer) {
+  allocator().release(buffer);
+}

 } // namespace mlx::core::allocator
@@ -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)
@@ -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) {
@@ -64,24 +65,48 @@ array array::unsafe_weak_copy(const array& other) {
      other.strides(),
      other.flags(),
      [](auto) {});
-  cpy.array_desc_->data_ptr = other.array_desc_->data_ptr;
+  cpy.array_desc_->offset = other.array_desc_->offset;
  return cpy;
 }

 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());
 }

+array::array(
+    void* data,
+    Shape shape,
+    Dtype dtype,
+    const std::function<void(void*)>& deleter)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  auto buffer = allocator::make_buffer(data, nbytes());
+  if (buffer.ptr() == nullptr) {
+    set_data(allocator::malloc(nbytes()));
+    auto ptr = static_cast<char*>(data);
+    std::copy(ptr, ptr + nbytes(), this->data<char>());
+    deleter(data);
+  } else {
+    auto wrapped_deleter = [deleter](allocator::Buffer buffer) {
+      auto ptr = buffer.raw_ptr();
+      allocator::release(buffer);
+      return deleter(ptr);
+    };
+    set_data(buffer, std::move(wrapped_deleter));
+  }
+}
+
 /* Build an array from a shared buffer */
 array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter)
    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
@@ -109,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;
  }
@@ -141,7 +167,7 @@ bool array::is_tracer() const {

 void array::set_data(allocator::Buffer buffer, Deleter d) {
  array_desc_->data = std::make_shared<Data>(buffer, d);
-  array_desc_->data_ptr = buffer.raw_ptr();
+  array_desc_->offset = 0;
  array_desc_->data_size = size();
  array_desc_->flags.contiguous = true;
  array_desc_->flags.row_contiguous = true;
@@ -156,7 +182,7 @@ void array::set_data(
    Flags flags,
    Deleter d) {
  array_desc_->data = std::make_shared<Data>(buffer, d);
-  array_desc_->data_ptr = buffer.raw_ptr();
+  array_desc_->offset = 0;
  array_desc_->data_size = data_size;
  array_desc_->strides = std::move(strides);
  array_desc_->flags = flags;
@@ -167,14 +193,13 @@ void array::copy_shared_buffer(
    const Strides& strides,
    Flags flags,
    size_t data_size,
-    size_t offset /* = 0 */) {
+    int64_t offset /* = 0 */) {
  array_desc_->data = other.array_desc_->data;
  array_desc_->strides = strides;
  array_desc_->flags = flags;
  array_desc_->data_size = data_size;
-  auto char_offset = sizeof(char) * itemsize() * offset;
-  array_desc_->data_ptr = static_cast<void*>(
-      static_cast<char*>(other.array_desc_->data_ptr) + char_offset);
+  array_desc_->offset =
+      sizeof(char) * itemsize() * offset + other.array_desc_->offset;
 }

 void array::copy_shared_buffer(const array& other) {
@@ -241,8 +266,8 @@ array::ArrayDesc::ArrayDesc(
    std::vector<array> inputs)
    : shape(std::move(shape)),
      dtype(dtype),
-      status(Status::unscheduled),
      primitive(std::move(primitive)),
+      status(Status::unscheduled),
      inputs(std::move(inputs)) {
  init();
 }
@@ -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 */

@@ -57,6 +58,16 @@ class array {
      Shape shape,
      Dtype dtype = TypeToDtype<T>());

+  /* Build an array from a raw pointer. The constructor will attempt to use the
+   * input data without a copy. The deleter will be called when the array no
+   * longer needs the underlying memory - after the array is destroyed in the
+   * no-copy case and after the copy otherwise. */
+  explicit array(
+      void* data,
+      Shape shape,
+      Dtype dtype,
+      const std::function<void(void*)>& deleter);
+
  /* Build an array from a buffer */
  explicit array(
      allocator::Buffer data,
@@ -111,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. */
@@ -125,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. */
@@ -143,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;
@@ -294,6 +305,11 @@ class array {
    return array_desc_->siblings;
  }

+  /** The array's position in the sibling list. */
+  int sibling_position() const {
+    return array_desc_->position;
+  }
+
  void set_siblings(std::vector<array> siblings, uint16_t position) {
    array_desc_->siblings = std::move(siblings);
    array_desc_->position = position;
@@ -349,15 +365,23 @@ class array {
    return array_desc_->data;
  }

-  // Return a raw pointer to the arrays data
+  // Return a raw pointer to the arrays data. This function may do a copy if
+  // the underlying buffer is not accessible on the CPU. When accessing the
+  // data for GPU kernels, be sure to use the correct method / function for the
+  // given backend to access the GPU pointer.
  template <typename T>
  T* data() {
-    return static_cast<T*>(array_desc_->data_ptr);
+    return reinterpret_cast<T*>(
+        (static_cast<char*>(buffer().raw_ptr()) + array_desc_->offset));
  }

  template <typename T>
  const T* data() const {
-    return static_cast<T*>(array_desc_->data_ptr);
+    return const_cast<array&>(*this).data<T>();
+  }
+
+  int64_t offset() const {
+    return array_desc_->offset;
  }

  enum Status {
@@ -426,7 +450,7 @@ class array {
      const Strides& strides,
      Flags flags,
      size_t data_size,
-      size_t offset = 0);
+      int64_t offset = 0);

  void copy_shared_buffer(const array& other);

@@ -441,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;
@@ -461,14 +485,14 @@ class array {
    // can share the underlying data buffer.
    std::shared_ptr<Data> data;

-    // Properly offset data pointer
-    void* data_ptr{nullptr};
+    // Offset from beginning of data pointer
+    int64_t offset{0};

    // The size in elements of the data buffer the array accesses
-    size_t data_size;
+    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
@@ -518,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());
 }

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

 void broadcast(const array& in, array& out) {
  if (out.size() == 0) {
-    out.set_data(nullptr);
+    out.set_data(allocator::malloc(0));
    return;
  }
  Strides strides(out.ndim(), 0);
--- a/Show More
+++ b/Show More